From 317ff3f5b40f2bd496159250bf0bf4328249b1ab Mon Sep 17 00:00:00 2001 From: Wang Xiaofeng Date: Sun, 23 Aug 2026 22:20:24 +0800 Subject: [PATCH 1/2] Validate returned stream identifiers --- src/brpc/controller.cpp | 6 +++ test/brpc_streaming_rpc_unittest.cpp | 69 ++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/src/brpc/controller.cpp b/src/brpc/controller.cpp index 629332f48a..cef24afcae 100644 --- a/src/brpc/controller.cpp +++ b/src/brpc/controller.cpp @@ -1711,6 +1711,12 @@ void Controller::HandleStreamConnection(Socket *host_socket) { if (!FailedInline()) { SetFailed(EREQUEST, "The server didn't accept the stream"); } + } else if (_remote_stream_settings->extra_stream_ids_size() != + (int)stream_num - 1) { + SetFailed(ERESPONSE, "Server returned %d extra_stream_ids, " + "expected %d", + _remote_stream_settings->extra_stream_ids_size(), + (int)stream_num - 1); } else { for (size_t i = 0; i < stream_num; ++i) { if (Stream::Address(_request_streams[i], &ptrs[i]) != 0) { diff --git a/test/brpc_streaming_rpc_unittest.cpp b/test/brpc_streaming_rpc_unittest.cpp index 2e4c046a13..8c80aec853 100644 --- a/test/brpc_streaming_rpc_unittest.cpp +++ b/test/brpc_streaming_rpc_unittest.cpp @@ -27,6 +27,7 @@ #include "brpc/controller.h" #include "brpc/channel.h" #include "brpc/callback.h" +#include "brpc/details/controller_private_accessor.h" #include "brpc/socket.h" #include "brpc/stream_impl.h" #include "brpc/policy/streaming_rpc_protocol.h" @@ -1130,6 +1131,74 @@ class MyServiceWithExtraStream : public test::EchoService { int _n; }; +class MyServiceWithMismatchedExtraStreamIds : public test::EchoService { +public: + MyServiceWithMismatchedExtraStreamIds(size_t stream_count, int adjustment) + : _stream_count(stream_count), _adjustment(adjustment) {} + + void Echo(::google::protobuf::RpcController* controller, + const ::test::EchoRequest* request, + ::test::EchoResponse* response, + ::google::protobuf::Closure* done) override { + brpc::ClosureGuard done_guard(done); + brpc::Controller* cntl = static_cast(controller); + response->set_message(request->message()); + + brpc::ControllerPrivateAccessor accessor(cntl); + brpc::StreamSettings* settings = accessor.remote_stream_settings(); + if (_adjustment < 0) { + settings->mutable_extra_stream_ids()->RemoveLast(); + } else { + settings->add_extra_stream_ids(settings->extra_stream_ids(0)); + } + + brpc::StreamIds response_streams; + ASSERT_EQ(0, brpc::StreamAccept(response_streams, *cntl, nullptr)); + ASSERT_EQ((int)_stream_count + _adjustment, + (int)response_streams.size()); + } + +private: + size_t _stream_count; + int _adjustment; +}; + +TEST_F(StreamingRpcTest, reject_mismatched_returned_stream_identifiers) { + const size_t STREAM_COUNT = 3; + + for (int adjustment : {-1, 1}) { + brpc::Server server; + MyServiceWithMismatchedExtraStreamIds service(STREAM_COUNT, adjustment); + ASSERT_EQ(0, server.AddService( + &service, brpc::SERVER_DOESNT_OWN_SERVICE)); + ASSERT_EQ(0, server.Start(0, nullptr)); + + brpc::Channel channel; + ASSERT_EQ(0, channel.Init(server.listen_address(), nullptr)); + + brpc::Controller cntl; + brpc::StreamIds request_streams; + ASSERT_EQ(0, brpc::StreamCreate(request_streams, STREAM_COUNT, cntl, + nullptr)); + ASSERT_EQ(STREAM_COUNT, request_streams.size()); + + test::EchoService_Stub stub(&channel); + stub.Echo(&cntl, &request, &response, nullptr); + ASSERT_TRUE(cntl.Failed()); + ASSERT_EQ(brpc::ERESPONSE, cntl.ErrorCode()); + ASSERT_NE(std::string::npos, + cntl.ErrorText().find("extra_stream_ids, expected 2")); + + for (brpc::StreamId stream_id : request_streams) { + brpc::StreamUniquePtr stream; + ASSERT_NE(0, brpc::Stream::Address(stream_id, &stream)); + } + + server.Stop(0); + server.Join(); + } +} + TEST_F(StreamingRpcTest, batch_create_extra_stream) { const size_t STREAM_COUNT = 3; // 1 first stream + 2 extra streams const int N = 1000; From 186ec05b0b0720017dfcc1e9bf35dc60e6a8ff71 Mon Sep 17 00:00:00 2001 From: Wang Xiaofeng Date: Mon, 24 Aug 2026 23:26:42 +0800 Subject: [PATCH 2/2] Avoid narrowing stream count validation --- src/brpc/controller.cpp | 10 ++++++---- test/brpc_streaming_rpc_unittest.cpp | 4 +++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/brpc/controller.cpp b/src/brpc/controller.cpp index cef24afcae..4fff9fd2f4 100644 --- a/src/brpc/controller.cpp +++ b/src/brpc/controller.cpp @@ -1705,18 +1705,20 @@ void Controller::HandleStreamConnection(Socket *host_socket) { return; } size_t stream_num = _request_streams.size(); + const size_t expected_extra_streams = stream_num - 1; std::vector ptrs(stream_num); if (!FailedInline()) { if (_remote_stream_settings == nullptr) { if (!FailedInline()) { SetFailed(EREQUEST, "The server didn't accept the stream"); } - } else if (_remote_stream_settings->extra_stream_ids_size() != - (int)stream_num - 1) { + } else if (static_cast( + _remote_stream_settings->extra_stream_ids_size()) != + expected_extra_streams) { SetFailed(ERESPONSE, "Server returned %d extra_stream_ids, " - "expected %d", + "expected %zu", _remote_stream_settings->extra_stream_ids_size(), - (int)stream_num - 1); + expected_extra_streams); } else { for (size_t i = 0; i < stream_num; ++i) { if (Stream::Address(_request_streams[i], &ptrs[i]) != 0) { diff --git a/test/brpc_streaming_rpc_unittest.cpp b/test/brpc_streaming_rpc_unittest.cpp index 8c80aec853..7ae820aa64 100644 --- a/test/brpc_streaming_rpc_unittest.cpp +++ b/test/brpc_streaming_rpc_unittest.cpp @@ -1186,8 +1186,10 @@ TEST_F(StreamingRpcTest, reject_mismatched_returned_stream_identifiers) { stub.Echo(&cntl, &request, &response, nullptr); ASSERT_TRUE(cntl.Failed()); ASSERT_EQ(brpc::ERESPONSE, cntl.ErrorCode()); + const std::string expected_error = + "extra_stream_ids, expected " + std::to_string(STREAM_COUNT - 1); ASSERT_NE(std::string::npos, - cntl.ErrorText().find("extra_stream_ids, expected 2")); + cntl.ErrorText().find(expected_error)); for (brpc::StreamId stream_id : request_streams) { brpc::StreamUniquePtr stream;