diff --git a/cpp/src/arrow/flight/flight_test.cc b/cpp/src/arrow/flight/flight_test.cc index eb931011906f..ab74e1fd6b89 100644 --- a/cpp/src/arrow/flight/flight_test.cc +++ b/cpp/src/arrow/flight/flight_test.cc @@ -453,6 +453,47 @@ class TestBasicAuthHandler : public ::testing::Test { std::unique_ptr server_; }; +// A client auth handler that reads and writes after the server +// finishes the handshake. +class ReadWriteAfterFinishClientAuthHandler : public ClientAuthHandler { + public: + Status Authenticate(ClientAuthSender* outgoing, ClientAuthReader* incoming) override { + std::string token; + // The server finishes the handshake without sending anything. So + // this finishes the stream. + RETURN_NOT_OK(incoming->Read(&token)); + // The stream is already finished. + return outgoing->Write("token"); + } + + Status GetToken(std::string* token) override { + *token = ""; + return Status::OK(); + } +}; + +class TestNoOpAuthHandler : public ::testing::Test { + public: + void SetUp() { + ASSERT_OK(MakeServer( + &server_, &client_, + [](FlightServerOptions* options) { + options->auth_handler = std::make_unique(); + return Status::OK(); + }, + [](FlightClientOptions* options) { return Status::OK(); })); + } + + void TearDown() { + ASSERT_OK(client_->Close()); + ASSERT_OK(server_->Shutdown()); + } + + protected: + std::unique_ptr client_; + std::unique_ptr server_; +}; + class TestTls : public ::testing::Test { public: void SetUp() { @@ -1325,6 +1366,12 @@ TEST_F(TestBasicAuthHandler, CheckPeerIdentity) { ASSERT_EQ(result->body->ToString(), "user"); } +TEST_F(TestNoOpAuthHandler, ReadWriteAfterFinish) { + // GH-49818: This must not call Finish() on the handshake stream twice. + ASSERT_OK(client_->Authenticate( + {}, std::make_unique())); +} + TEST_F(TestTls, DoAction) { FlightCallOptions options; options.timeout = TimeoutDuration{5.0}; diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_client.cc b/cpp/src/arrow/flight/transport/grpc/grpc_client.cc index f1a5dcd68635..5f2a2aa0a0df 100644 --- a/cpp/src/arrow/flight/transport/grpc/grpc_client.cc +++ b/cpp/src/arrow/flight/transport/grpc/grpc_client.cc @@ -214,13 +214,47 @@ class GrpcClientInterceptorAdapterFactory std::vector> middleware_; }; -class GrpcClientAuthSender : public ClientAuthSender { +/// \brief A Handshake stream that calls Finish() at most once. +/// +/// gRPC aborts the process if Finish() is called twice on the same call. +class GrpcClientHandshakeStream { public: - explicit GrpcClientAuthSender( + explicit GrpcClientHandshakeStream( std::shared_ptr< ::grpc::ClientReaderWriter> stream) - : stream_(stream) {} + : stream_(std::move(stream)) {} + + bool finished() const { return finished_; } + + bool Write(const pb::HandshakeRequest& request) { + return !finished_ && stream_->Write(request); + } + + bool Read(pb::HandshakeResponse* response) { + return !finished_ && stream_->Read(response); + } + + bool WritesDone() { return !finished_ && stream_->WritesDone(); } + + const ::grpc::Status& Finish() { + if (!finished_) { + status_ = stream_->Finish(); + finished_ = true; + } + return status_; + } + + private: + std::shared_ptr<::grpc::ClientReaderWriter> + stream_; + bool finished_ = false; + ::grpc::Status status_; +}; + +class GrpcClientAuthSender : public ClientAuthSender { + public: + explicit GrpcClientAuthSender(GrpcClientHandshakeStream* stream) : stream_(stream) {} Status Write(const std::string& token) override { pb::HandshakeRequest response; @@ -232,17 +266,12 @@ class GrpcClientAuthSender : public ClientAuthSender { } private: - std::shared_ptr<::grpc::ClientReaderWriter> - stream_; + GrpcClientHandshakeStream* stream_; }; class GrpcClientAuthReader : public ClientAuthReader { public: - explicit GrpcClientAuthReader( - std::shared_ptr< - ::grpc::ClientReaderWriter> - stream) - : stream_(stream) {} + explicit GrpcClientAuthReader(GrpcClientHandshakeStream* stream) : stream_(stream) {} Status Read(std::string* token) override { pb::HandshakeResponse request; @@ -254,8 +283,7 @@ class GrpcClientAuthReader : public ClientAuthReader { } private: - std::shared_ptr<::grpc::ClientReaderWriter> - stream_; + GrpcClientHandshakeStream* stream_; }; /// \brief The base of the ClientDataStream implementation for gRPC. @@ -1065,16 +1093,18 @@ class GrpcClientImpl : public internal::ClientTransport { private: Status AuthenticateInternal(ClientRpc& rpc) { - std::shared_ptr< - ::grpc::ClientReaderWriter> - stream = stub_->Handshake(&rpc.context); + GrpcClientHandshakeStream stream{stub_->Handshake(&rpc.context)}; if (auth_handler_) { - GrpcClientAuthSender outgoing{stream}; - GrpcClientAuthReader incoming{stream}; + GrpcClientAuthSender outgoing{&stream}; + GrpcClientAuthReader incoming{&stream}; RETURN_NOT_OK(auth_handler_->Authenticate(&outgoing, &incoming)); } + // The server may have finished the handshake while authenticating + if (stream.finished()) { + return FromGrpcStatus(stream.Finish(), &rpc.context); + } // Explicitly close our side of the connection - bool finished_writes = stream->WritesDone(); + bool finished_writes = stream.WritesDone(); if (!finished_writes) { return MakeFlightError(FlightStatusCode::Internal, "Could not finish writing before closing"); @@ -1084,9 +1114,9 @@ class GrpcClientImpl : public internal::ClientTransport { // reader finishes, so it's OK to assume the client no longer // wants to read and drain the read side. pb::HandshakeResponse response; - while (stream->Read(&response)) { + while (stream.Read(&response)) { } - RETURN_NOT_OK(FromGrpcStatus(stream->Finish(), &rpc.context)); + RETURN_NOT_OK(FromGrpcStatus(stream.Finish(), &rpc.context)); return Status::OK(); }