From 23adf3a61aa4b9a0129362a184e91a2c4bbfcf70 Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Tue, 3 Nov 2020 18:38:37 -0500
Subject: [PATCH] ARROW-10475: [C++][FlightRPC] handle IPv6 hosts

---
 cpp/src/arrow/flight/client.cc      |  3 ++-
 cpp/src/arrow/flight/flight_test.cc | 17 +++++++++++++++++
 cpp/src/arrow/flight/server.cc      |  3 ++-
 cpp/src/arrow/util/uri.cc           | 13 +++++++++++++
 cpp/src/arrow/util/uri.h            |  5 +++++
 cpp/src/arrow/util/uri_test.cc      |  6 ++++++
 6 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/cpp/src/arrow/flight/client.cc b/cpp/src/arrow/flight/client.cc
index 821da214a59e..cdffa7f06567 100644
--- a/cpp/src/arrow/flight/client.cc
+++ b/cpp/src/arrow/flight/client.cc
@@ -876,7 +876,8 @@ class FlightClient::FlightClientImpl {
     std::stringstream grpc_uri;
     std::shared_ptr<grpc::ChannelCredentials> creds;
     if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp || scheme == kSchemeGrpcTls) {
-      grpc_uri << location.uri_->host() << ":" << location.uri_->port_text();
+      grpc_uri << arrow::internal::UriEncodeHost(location.uri_->host()) << ':'
+               << location.uri_->port_text();
 
       if (scheme == kSchemeGrpcTls) {
         if (options.disable_server_verification) {
diff --git a/cpp/src/arrow/flight/flight_test.cc b/cpp/src/arrow/flight/flight_test.cc
index a5fdfe8184e1..8d7e944bcddd 100644
--- a/cpp/src/arrow/flight/flight_test.cc
+++ b/cpp/src/arrow/flight/flight_test.cc
@@ -289,6 +289,23 @@ TEST(TestFlight, GetPort) {
   ASSERT_GT(server->port(), 0);
 }
 
+// CI environments don't have an IPv6 interface configured
+TEST(TestFlight, DISABLED_IpV6Port) {
+  Location location, location2;
+  std::unique_ptr<FlightServerBase> server = ExampleTestServer();
+
+  ASSERT_OK(Location::ForGrpcTcp("[::1]", 0, &location));
+  FlightServerOptions options(location);
+  ASSERT_OK(server->Init(options));
+  ASSERT_GT(server->port(), 0);
+
+  ASSERT_OK(Location::ForGrpcTcp("[::1]", server->port(), &location2));
+  std::unique_ptr<FlightClient> client;
+  ASSERT_OK(FlightClient::Connect(location2, &client));
+  std::unique_ptr<FlightListing> listing;
+  ASSERT_OK(client->ListFlights(&listing));
+}
+
 TEST(TestFlight, BuilderHook) {
   Location location;
   std::unique_ptr<FlightServerBase> server = ExampleTestServer();
diff --git a/cpp/src/arrow/flight/server.cc b/cpp/src/arrow/flight/server.cc
index 0cdd3b43f516..87c96ce49102 100644
--- a/cpp/src/arrow/flight/server.cc
+++ b/cpp/src/arrow/flight/server.cc
@@ -822,7 +822,8 @@ Status FlightServerBase::Init(const FlightServerOptions& options) {
   const std::string scheme = location.scheme();
   if (scheme == kSchemeGrpc || scheme == kSchemeGrpcTcp || scheme == kSchemeGrpcTls) {
     std::stringstream address;
-    address << location.uri_->host() << ':' << location.uri_->port_text();
+    address << arrow::internal::UriEncodeHost(location.uri_->host()) << ':'
+            << location.uri_->port_text();
 
     std::shared_ptr<grpc::ServerCredentials> creds;
     if (scheme == kSchemeGrpcTls) {
diff --git a/cpp/src/arrow/util/uri.cc b/cpp/src/arrow/util/uri.cc
index 1261607b6c14..f48a8cb6ed17 100644
--- a/cpp/src/arrow/util/uri.cc
+++ b/cpp/src/arrow/util/uri.cc
@@ -80,6 +80,19 @@ std::string UriEscape(const std::string& s) {
   return escaped;
 }
 
+std::string UriEncodeHost(const std::string& host) {
+  // Fairly naive check: if it contains a ':', it's IPv6 and needs
+  // brackets, else it's OK
+  if (host.find(":") != std::string::npos) {
+    std::string result = "[";
+    result += host;
+    result += ']';
+    return result;
+  } else {
+    return host;
+  }
+}
+
 struct Uri::Impl {
   Impl() : string_rep_(""), port_(-1) { memset(&uri_, 0, sizeof(uri_)); }
 
diff --git a/cpp/src/arrow/util/uri.h b/cpp/src/arrow/util/uri.h
index 6ef54900020a..480e35474a31 100644
--- a/cpp/src/arrow/util/uri.h
+++ b/cpp/src/arrow/util/uri.h
@@ -91,5 +91,10 @@ class ARROW_EXPORT Uri {
 ARROW_EXPORT
 std::string UriEscape(const std::string& s);
 
+/// Encode a host for use within a URI, such as "localhost",
+/// "127.0.0.1", or "[::1]".
+ARROW_EXPORT
+std::string UriEncodeHost(const std::string& host);
+
 }  // namespace internal
 }  // namespace arrow
diff --git a/cpp/src/arrow/util/uri_test.cc b/cpp/src/arrow/util/uri_test.cc
index 44804704c436..169e9c81b36b 100644
--- a/cpp/src/arrow/util/uri_test.cc
+++ b/cpp/src/arrow/util/uri_test.cc
@@ -35,6 +35,12 @@ TEST(UriEscape, Basics) {
   ASSERT_EQ(UriEscape("/El Niño/"), "%2FEl%20Ni%C3%B1o%2F");
 }
 
+TEST(UriEncodeHost, Basics) {
+  ASSERT_EQ(UriEncodeHost("::1"), "[::1]");
+  ASSERT_EQ(UriEscape("arrow.apache.org"), "arrow.apache.org");
+  ASSERT_EQ(UriEscape("192.168.1.1"), "192.168.1.1");
+}
+
 TEST(Uri, Empty) {
   Uri uri;
   ASSERT_EQ(uri.scheme(), "");