From 50ba0dbf5f78df8a0d6b0c2d5b30048e3f6d3084 Mon Sep 17 00:00:00 2001 From: Wang Xiaofeng Date: Sun, 23 Aug 2026 22:20:05 +0800 Subject: [PATCH] Fix streaming RPC frame validation --- src/brpc/controller.cpp | 8 +- src/brpc/policy/baidu_rpc_protocol.cpp | 18 ++- src/brpc/policy/streaming_rpc_protocol.cpp | 13 ++- src/brpc/socket.cpp | 7 ++ src/brpc/socket.h | 4 + src/brpc/stream.cpp | 19 +++- src/brpc/stream_impl.h | 3 + test/brpc_streaming_rpc_unittest.cpp | 122 +++++++++++++++++++++ 8 files changed, 185 insertions(+), 9 deletions(-) diff --git a/src/brpc/controller.cpp b/src/brpc/controller.cpp index 6c3ecee372..629332f48a 100644 --- a/src/brpc/controller.cpp +++ b/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); } } diff --git a/src/brpc/policy/baidu_rpc_protocol.cpp b/src/brpc/policy/baidu_rpc_protocol.cpp index b74fb1805a..f39463e888 100644 --- a/src/brpc/policy/baidu_rpc_protocol.cpp +++ b/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"; diff --git a/src/brpc/policy/streaming_rpc_protocol.cpp b/src/brpc/policy/streaming_rpc_protocol.cpp index bdad1f2385..155ec18875 100644 --- a/src/brpc/policy/streaming_rpc_protocol.cpp +++ b/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(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); diff --git a/src/brpc/socket.cpp b/src/brpc/socket.cpp index 38727571ba..becb094486 100644 --- a/src/brpc/socket.cpp +++ b/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 " diff --git a/src/brpc/socket.h b/src/brpc/socket.h index 6d321f8bc3..bbde3ecda9 100644 --- a/src/brpc/socket.h +++ b/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 diff --git a/src/brpc/stream.cpp b/src/brpc/stream.cpp index 8f5ef3c50b..8db9352320 100644 --- a/src/brpc/stream.cpp +++ b/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; diff --git a/src/brpc/stream_impl.h b/src/brpc/stream_impl.h index 06e5044859..4d3dd0c5fb 100644 --- a/src/brpc/stream_impl.h +++ b/src/brpc/stream_impl.h @@ -159,6 +159,9 @@ friend class VersionedRefWithId; 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 _consumer_queue; butil::IOBuf* _pending_buf; int64_t _start_idle_timer_us; diff --git a/test/brpc_streaming_rpc_unittest.cpp b/test/brpc_streaming_rpc_unittest.cpp index 46401d6db0..2e4c046a13 100644 --- a/test/brpc_streaming_rpc_unittest.cpp +++ b/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 _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};