Fix streaming RPC frame validation

src/brpc/controller.cpp

@@ -1745,7 +1745,13 @@ void Controller::HandleStreamConnection(Socket *host_socket) {
             if(!ptrs[i]) continue;
             Stream* extra_stream = ptrs[i].get();
             _remote_stream_settings->set_stream_id(extra_stream_ids[i - 1]);
-            extra_stream->SetHostSocket(host_socket);
+            if (extra_stream->SetHostSocket(host_socket) != 0) {
+                SetFailed(EREQUEST, "Fail to bind response stream=%" PRIu64,
+                          extra_stream_ids[i - 1]);
+                Stream::SetFailed(_request_streams, _error_code,
+                                  "%s", _error_text.c_str());
+                return;
+            }
             extra_stream->SetConnected(_remote_stream_settings);
         }
     }

src/brpc/policy/baidu_rpc_protocol.cpp

@@ -372,7 +372,14 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
             Stream* s = stream_ptr.get();
             StreamSettings *stream_settings = meta.mutable_stream_settings();
             s->FillSettings(stream_settings);
-            s->SetHostSocket(sock);
+            if (s->SetHostSocket(sock) != 0) {
+                cntl->SetFailed(EINVAL, "Fail to bind stream=%" PRIu64
+                                " to %s", response_stream_id,
+                                sock->description().c_str());
+                Stream::SetFailed(response_stream_ids, EINVAL,
+                                  "%s", cntl->ErrorText().c_str());
+                return;
+            }
             for (size_t i = 1; i < response_stream_ids.size(); ++i) {
                 stream_settings->mutable_extra_stream_ids()->Add(response_stream_ids[i]);
             }
@@ -438,8 +445,13 @@ void SendRpcResponse(int64_t correlation_id, Controller* cntl,
             StreamUniquePtr extra_stream_ptr;
             if (Stream::Address(extra_stream_id, &extra_stream_ptr) == 0) {
                 Stream* extra_stream = extra_stream_ptr.get();
-                extra_stream->SetHostSocket(sock);
-                extra_stream->SetConnected();
+                if (extra_stream->SetHostSocket(sock) == 0) {
+                    extra_stream->SetConnected();
+                } else {
+                    Stream::SetFailed(extra_stream_id, EINVAL,
+                                      "Fail to bind stream to %s",
+                                      sock->description().c_str());
+                }
             } else {
                 LOG(WARNING) << "Stream=" << extra_stream_id
                              << " was closed before sending response";

src/brpc/policy/streaming_rpc_protocol.cpp

@@ -27,6 +27,7 @@
 #include "butil/iobuf.h"                         // butil::IOBuf
 #include "butil/raw_pack.h"                      // RawPacker RawUnpacker
 #include "brpc/log.h"
+#include "brpc/server.h"
 #include "brpc/socket.h"                        // Socket
 #include "brpc/streaming_rpc_meta.pb.h"         // StreamFrameMeta
 #include "brpc/policy/most_common_message.h"
@@ -59,7 +60,7 @@ void PackStreamMessage(butil::IOBuf* out,
 }
 
 ParseResult ParseStreamingMessage(butil::IOBuf* source,
-                            Socket* socket, bool /*read_eof*/, const void* /*arg*/) {
+                            Socket* socket, bool /*read_eof*/, const void* arg) {
     char header_buf[12];
     const size_t n = source->copy_to(header_buf, sizeof(header_buf));
     if (n >= 4) {
@@ -90,6 +91,16 @@ ParseResult ParseStreamingMessage(butil::IOBuf* source,
         source->pop_front(sizeof(header_buf) + body_size);
         return MakeParseError(PARSE_ERROR_TRY_OTHERS);
     }
+    if (arg != nullptr) {
+        // Stream frames are consumed here and never reach InputMessenger's
+        // authentication hook.
+        const Server* server = static_cast<const Server*>(arg);
+        if (server->options().auth != nullptr && !socket->IsAuthenticated()) {
+            LOG(WARNING) << "Reject streaming frame from unauthenticated "
+                         << *socket;
+            return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
+        }
+    }
     source->pop_front(sizeof(header_buf));
     butil::IOBuf meta_buf;
     source->cutn(&meta_buf, meta_size);

src/brpc/socket.cpp

@@ -2206,6 +2206,13 @@ void Socket::SetAuthentication(int error_code) {
     }
 }
 
+bool Socket::IsAuthenticated() const {
+    const uint64_t flag_error =
+        _auth_flag_error.load(butil::memory_order_acquire);
+    return (flag_error & AUTH_FLAG) &&
+        (int32_t)(flag_error & 0xFFFFFFFFul) == 0;
+}
+
 AuthContext* Socket::mutable_auth_context() {
     if (_auth_context != nullptr) {
         LOG(FATAL) << "Impossible! This function is supposed to be called "

src/brpc/socket.h

@@ -556,6 +556,10 @@ friend class TransportFactory;
     // `FightAuthentication', otherwise it's regarded as an error
     void SetAuthentication(int error_code);
 
+    // Returns true iff authentication over this socket has completed
+    // successfully, i.e. `SetAuthentication(0)' was called.
+    bool IsAuthenticated() const;
+
     // Since some protocols are not able to store correlation id in their
     // headers (such as nova-pbrpc, http), we have to store it here. Note
     // that there can only be 1 RPC call on this socket at any time, otherwise

src/brpc/stream.cpp

@@ -54,6 +54,7 @@ Stream::Stream(Forbidden f)
     , _local_consumed(0)
     , _atomic_local_consumed(0)
     , _parse_rpc_response(false)
+    , _server_accepted_stream(false)
     , _pending_buf(nullptr)
     , _start_idle_timer_us(0)
     , _idle_timer(0) {
@@ -94,6 +95,7 @@ int Stream::OnCreated(const StreamOptions& options,
     _local_consumed = 0;
     _atomic_local_consumed.store(0, butil::memory_order_relaxed);
     _parse_rpc_response = parse_rpc_response;
+    _server_accepted_stream = (remote_settings != nullptr);
     _pending_buf = nullptr;
     _start_idle_timer_us = 0;
     _idle_timer = 0;
@@ -584,11 +586,20 @@ void Stream::SetConnected(const StreamSettings* remote_settings) {
 
 int Stream::OnReceived(const StreamFrameMeta& fm, butil::IOBuf *buf, Socket* sock) {
     if (!_connected.load(butil::memory_order_acquire)) {
-        // Before connection is published, let the locked slow path initialize
-        // the host socket or confirm that another thread already did so.
-        if (SetHostSocket(sock) != 0) {
+        if (_server_accepted_stream) {
+            BAIDU_SCOPED_LOCK(_connect_mutex);
+            if (_host_socket == nullptr || _host_socket->id() != sock->id()) {
+                LOG(WARNING) << "stream=" << id()
+                             << " dropped a frame from a foreign socket";
+                return -1;
+            }
+        } else if (SetHostSocket(sock) != 0) {
             return -1;
         }
+    } else if (_host_socket == nullptr || _host_socket->id() != sock->id()) {
+        LOG(WARNING) << "stream=" << id()
+                     << " dropped a frame from a foreign socket";
+        return -1;
     }
 
     switch (fm.frame_type()) {
@@ -758,7 +769,7 @@ int Stream::SetHostSocket(Socket* host_socket) {
         return -1;
     }
     if (_host_socket != nullptr) {
-        return 0;
+        return _host_socket->id() == host_socket->id() ? 0 : -1;
     }
 
     SocketUniquePtr ptr;

src/brpc/stream_impl.h

@@ -159,6 +159,9 @@ friend class VersionedRefWithId<Stream>;
     StreamSettings _remote_settings;
 
     bool _parse_rpc_response;
+    // Server-accepted streams must be bound by their creating RPC, never by
+    // the first frame that happens to carry their id.
+    bool _server_accepted_stream;
     bthread::ExecutionQueueId<butil::IOBuf*> _consumer_queue;
     butil::IOBuf* _pending_buf;
     int64_t _start_idle_timer_us;

test/brpc_streaming_rpc_unittest.cpp

@@ -81,6 +81,128 @@ class StreamingRpcTest : public testing::Test {
     test::EchoResponse response;
 };
 
+class StreamingRpcAuthenticator : public brpc::Authenticator {
+public:
+    int GenerateCredential(std::string* auth_str) const override {
+        *auth_str = "credential";
+        return 0;
+    }
+
+    int VerifyCredential(const std::string& auth_str,
+                         const butil::EndPoint&,
+                         brpc::AuthContext*) const override {
+        return auth_str == "credential" ? 0 : brpc::ERPCAUTH;
+    }
+};
+
+class AuthenticatedStreamHandler : public brpc::StreamInputHandler {
+public:
+    int on_received_messages(brpc::StreamId,
+                             butil::IOBuf* const messages[],
+                             size_t size) override {
+        for (size_t i = 0; i < size; ++i) {
+            if (messages[i]->to_string() == "authenticated stream frame") {
+                _received.store(true, std::memory_order_release);
+            }
+        }
+        return 0;
+    }
+
+    void on_idle_timeout(brpc::StreamId) override {}
+    void on_closed(brpc::StreamId) override {}
+    void on_failed(brpc::StreamId, int, const std::string&) override {}
+
+    bool received() const {
+        return _received.load(std::memory_order_acquire);
+    }
+
+private:
+    std::atomic<bool> _received{false};
+};
+
+TEST_F(StreamingRpcTest, reject_stream_frame_from_unauthenticated_socket) {
+    StreamingRpcAuthenticator auth;
+    brpc::Server server;
+    brpc::ServerOptions server_options;
+    server_options.auth = &auth;
+    ASSERT_EQ(0, server.Start(0, &server_options));
+
+    brpc::SocketId socket_id;
+    brpc::SocketOptions socket_options;
+    ASSERT_EQ(0, brpc::Socket::Create(socket_options, &socket_id));
+    brpc::SocketUniquePtr socket;
+    ASSERT_EQ(0, brpc::Socket::Address(socket_id, &socket));
+
+    brpc::StreamFrameMeta frame_meta;
+    frame_meta.set_stream_id(brpc::INVALID_STREAM_ID);
+    frame_meta.set_frame_type(brpc::FRAME_TYPE_CLOSE);
+    butil::IOBuf frame;
+    brpc::policy::PackStreamMessage(&frame, frame_meta, nullptr);
+    const size_t frame_size = frame.size();
+
+    brpc::ParseResult result = brpc::policy::ParseStreamingMessage(
+        &frame, socket.get(), false, &server);
+    ASSERT_EQ(brpc::PARSE_ERROR_ABSOLUTELY_WRONG, result.error());
+    ASSERT_EQ(frame_size, frame.size());
+}
+
+TEST_F(StreamingRpcTest, authenticate_before_exchanging_stream_frames) {
+    StreamingRpcAuthenticator auth;
+    AuthenticatedStreamHandler handler;
+    brpc::StreamOptions server_stream_options;
+    server_stream_options.handler = &handler;
+
+    brpc::Server server;
+    MyServiceWithStream service(server_stream_options);
+    ASSERT_EQ(0, server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE));
+    brpc::ServerOptions server_options;
+    server_options.auth = &auth;
+    ASSERT_EQ(0, server.Start(0, &server_options));
+
+    // A client without credentials is rejected before the service accepts its
+    // stream, and the server-side stream handler does not observe any frame.
+    brpc::Channel unauthenticated_channel;
+    brpc::ChannelOptions unauthenticated_options;
+    unauthenticated_options.max_retry = 0;
+    ASSERT_EQ(0, unauthenticated_channel.Init(
+        server.listen_address(), &unauthenticated_options));
+    brpc::Controller unauthenticated_cntl;
+    brpc::StreamId unauthenticated_stream;
+    ASSERT_EQ(0, StreamCreate(
+        &unauthenticated_stream, unauthenticated_cntl, nullptr));
+    brpc::ScopedStream unauthenticated_stream_guard(unauthenticated_stream);
+    test::EchoResponse unauthenticated_response;
+    test::EchoService_Stub unauthenticated_stub(&unauthenticated_channel);
+    unauthenticated_stub.Echo(&unauthenticated_cntl, &request,
+                              &unauthenticated_response, nullptr);
+    ASSERT_TRUE(unauthenticated_cntl.Failed());
+    ASSERT_EQ(brpc::ERPCAUTH, unauthenticated_cntl.ErrorCode());
+    ASSERT_FALSE(handler.received());
+
+    // A normal client authenticates in the RPC service path before either
+    // endpoint exchanges stream frames.
+    brpc::ChannelOptions channel_options;
+    channel_options.auth = &auth;
+    brpc::Channel channel;
+    ASSERT_EQ(0, channel.Init(server.listen_address(), &channel_options));
+    brpc::Controller cntl;
+    brpc::StreamId request_stream;
+    ASSERT_EQ(0, StreamCreate(&request_stream, cntl, nullptr));
+    brpc::ScopedStream stream_guard(request_stream);
+    test::EchoService_Stub stub(&channel);
+    stub.Echo(&cntl, &request, &response, nullptr);
+    ASSERT_FALSE(cntl.Failed()) << cntl.ErrorText();
+
+    butil::IOBuf message;
+    message.append("authenticated stream frame");
+    ASSERT_EQ(0, brpc::StreamWrite(request_stream, message));
+    const int64_t deadline = butil::gettimeofday_us() + 3000000L;
+    while (!handler.received() && butil::gettimeofday_us() < deadline) {
+        usleep(1000);
+    }
+    ASSERT_TRUE(handler.received());
+}
+
 struct BatchStreamFeedbackRaceState {
     brpc::StreamId server_first_stream_id{brpc::INVALID_STREAM_ID};
     brpc::StreamId server_extra_stream_id{brpc::INVALID_STREAM_ID};