ARROW-10487 [FlightRPC][C++] Header-based auth in clients

cpp/src/arrow/flight/CMakeLists.txt

@@ -118,6 +118,7 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS_BACKUP}")
 # protobuf-internal.cc
 set(ARROW_FLIGHT_SRCS
     client.cc
+    client_header_internal.cc
     internal.cc
     protocol_internal.cc
     serialization_internal.cc

cpp/src/arrow/flight/client.cc

@@ -50,6 +50,7 @@
 #include "arrow/util/uri.h"
 
 #include "arrow/flight/client_auth.h"
+#include "arrow/flight/client_header_internal.h"
 #include "arrow/flight/client_middleware.h"
 #include "arrow/flight/internal.h"
 #include "arrow/flight/middleware.h"
@@ -104,6 +105,9 @@ struct ClientRpc {
               std::chrono::system_clock::now() + options.timeout);
       context.set_deadline(deadline);
     }
+    for (auto header : options.headers) {
+      context.AddMetadata(header.first, header.second);
+    }
   }
 
   /// \brief Add an auth token via an auth handler
@@ -994,6 +998,31 @@ class FlightClient::FlightClientImpl {
     return Status::OK();
   }
 
+  Status AuthenticateBasicToken(const FlightCallOptions& options,
+                                const std::string& username, const std::string& password,
+                                std::pair<std::string, std::string>* bearer_token) {
+    // Add basic auth headers to outgoing headers.
+    ClientRpc rpc(options);
+    internal::AddBasicAuthHeaders(&rpc.context, username, password);
+
+    std::shared_ptr<grpc::ClientReaderWriter<pb::HandshakeRequest, pb::HandshakeResponse>>
+        stream = stub_->Handshake(&rpc.context);
+    GrpcClientAuthSender outgoing{stream};
+    GrpcClientAuthReader incoming{stream};
+
+    // Explicitly close our side of the connection.
+    bool finished_writes = stream->WritesDone();
+    RETURN_NOT_OK(internal::FromGrpcStatus(stream->Finish(), &rpc.context));
+    if (!finished_writes) {
+      return MakeFlightError(FlightStatusCode::Internal,
+                             "Could not finish writing before closing");
+    }
+
+    // Grab bearer token from incoming headers.
+    internal::GetBearerTokenHeader(rpc.context, bearer_token);
+    return Status::OK();
+  }
+
   Status ListFlights(const FlightCallOptions& options, const Criteria& criteria,
                      std::unique_ptr<FlightListing>* listing) {
     pb::Criteria pb_criteria;
@@ -1198,6 +1227,12 @@ Status FlightClient::Authenticate(const FlightCallOptions& options,
   return impl_->Authenticate(options, std::move(auth_handler));
 }
 
+Status FlightClient::AuthenticateBasicToken(
+    const FlightCallOptions& options, const std::string& username,
+    const std::string& password, std::pair<std::string, std::string>* bearer_token) {
+  return impl_->AuthenticateBasicToken(options, username, password, bearer_token);
+}
+
 Status FlightClient::DoAction(const FlightCallOptions& options, const Action& action,
                               std::unique_ptr<ResultStream>* results) {
   return impl_->DoAction(options, action, results);

cpp/src/arrow/flight/client.h

@@ -65,6 +65,9 @@ class ARROW_FLIGHT_EXPORT FlightCallOptions {
 
   /// \brief IPC writer options, if applicable for the call.
   ipc::IpcWriteOptions write_options;
+
+  /// \brief Headers for client to add to context.
+  std::vector<std::pair<std::string, std::string>> headers;
 };
 
 /// \brief Indicate that the client attempted to write a message
@@ -191,6 +194,16 @@ class ARROW_FLIGHT_EXPORT FlightClient {
   Status Authenticate(const FlightCallOptions& options,
                       std::unique_ptr<ClientAuthHandler> auth_handler);
 
+  /// \brief Authenticate to the server using basic HTTP style authentication.
+  /// \param[in] options Per-RPC options
+  /// \param[in] username Username to use
+  /// \param[in] password Password to use
+  /// \param[in] bearer_token Bearer token retreived if applicable
+  /// \return Status OK if the client authenticated successfully
+  Status AuthenticateBasicToken(const FlightCallOptions& options,
+                                const std::string& username, const std::string& password,
+                                std::pair<std::string, std::string>* bearer_token);
+
   /// \brief Perform the indicated action, returning an iterator to the stream
   /// of results, if any
   /// \param[in] options Per-RPC options

cpp/src/arrow/flight/client_header_internal.cc

@@ -0,0 +1,89 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Interfaces for defining middleware for Flight clients. Currently
+// experimental.
+
+#include "arrow/flight/client_header_internal.h"
+#include "arrow/flight/client.h"
+#include "arrow/flight/client_auth.h"
+#include "arrow/util/base64.h"
+#include "arrow/util/make_unique.h"
+
+#include <algorithm>
+#include <cctype>
+#include <memory>
+#include <string>
+
+const char kAuthHeader[] = "authorization";
+const char kBearerPrefix[] = "Bearer ";
+const char kBasicPrefix[] = "Basic ";
+
+namespace arrow {
+namespace flight {
+namespace internal {
+
+// Add base64 encoded credentials to the outbound headers.
+//
+// @param context Context object to add the headers to.
+// @param username Username to format and encode.
+// @param password Password to format and encode.
+void AddBasicAuthHeaders(grpc::ClientContext* context, const std::string& username,
+                         const std::string& password) {
+  const std::string credentials = username + ":" + password;
+  context->AddMetadata(
+      kAuthHeader,
+      kBasicPrefix + arrow::util::base64_encode(
+                         reinterpret_cast<const unsigned char*>(credentials.c_str()),
+                         static_cast<unsigned int>(credentials.size())));
+}
+
+// Get bearer token from inbound headers.
+//
+// @param context Incoming ClientContext that contains headers.
+// @param bearer_token[out] Bearer token pointer to set.
+void GetBearerTokenHeader(grpc::ClientContext& context,
+                          std::pair<std::string, std::string>* bearer_token) {
+  // Lambda function to compare characters without case sensitivity.
+  auto char_compare = [](const char& char1, const char& char2) {
+    return (::toupper(char1) == ::toupper(char2));
+  };
+
+  // Get the auth token if it exists, this can be in the initial or the trailing metadata.
+  auto trailing_headers = context.GetServerTrailingMetadata();
+  auto initial_headers = context.GetServerInitialMetadata();
+  auto bearer_iter = trailing_headers.find(kAuthHeader);
+  if (bearer_iter == trailing_headers.end()) {
+    bearer_iter = initial_headers.find(kAuthHeader);
+    if (bearer_iter == initial_headers.end()) {
+      return;
+    }
+  }
+
+  // Check if the value of the auth token starts with the bearer prefix and latch it.
+  std::string bearer_val(bearer_iter->second.data(), bearer_iter->second.size());
+  if (bearer_val.size() > strlen(kBearerPrefix)) {
+    if (std::equal(bearer_val.begin(), bearer_val.begin() + strlen(kBearerPrefix),
+                   kBearerPrefix, char_compare)) {
+      *bearer_token = std::make_pair(kAuthHeader, bearer_val);
+    }
+  }
+}
+
+}  // namespace internal
+}  // namespace flight
+}  // namespace arrow

cpp/src/arrow/flight/client_header_internal.h

@@ -0,0 +1,57 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Interfaces for defining middleware for Flight clients. Currently
+// experimental.
+
+#pragma once
+
+#include "arrow/flight/client_middleware.h"
+
+#ifdef GRPCPP_PP_INCLUDE
+#include <grpcpp/grpcpp.h>
+#if defined(GRPC_NAMESPACE_FOR_TLS_CREDENTIALS_OPTIONS)
+#include <grpcpp/security/tls_credentials_options.h>
+#endif
+#else
+#include <grpc++/grpc++.h>
+#endif
+
+namespace arrow {
+namespace flight {
+namespace internal {
+
+/// \brief Add basic authentication header key value pair to context.
+///
+/// \param context grpc context variable to add header to.
+/// \param username username to encode into header.
+/// \param password password to to encode into header.
+void ARROW_FLIGHT_EXPORT AddBasicAuthHeaders(grpc::ClientContext* context,
+                                             const std::string& username,
+                                             const std::string& password);
+
+/// \brief Get bearer token from incoming headers.
+///
+/// \param context context that contains headers which hold the bearer token.
+/// \param[out] bearer_token pointer to a std::pair of std::strings that the factory
+///     will populate with the bearer token that is received from the server.
+void ARROW_FLIGHT_EXPORT GetBearerTokenHeader(
+    grpc::ClientContext& context, std::pair<std::string, std::string>* bearer_token);
+
+}  // namespace internal
+}  // namespace flight
+}  // namespace arrow