Allow HTTP Server to bind to a specific address, as opposed to listening on all addresses.

PiperOrigin-RevId: 710784609

Author: tensorflower-gardener

tensorflow_serving/util/net_http/server/internal/evhttp_server.cc

@@ -223,14 +223,25 @@ bool EvHTTPServer::StartAcceptingRequests() {
 
   const int port = server_options_->ports().front();
 
-  // "::"  =>  in6addr_any
-  ev_uint16_t ev_port = static_cast<ev_uint16_t>(port);
-  ev_listener_ = evhttp_bind_socket_with_handle(ev_http_, "::", ev_port);
-  if (ev_listener_ == nullptr) {
-    // in case ipv6 is not supported, fallback to inaddr_any
-    ev_listener_ = evhttp_bind_socket_with_handle(ev_http_, nullptr, ev_port);
+  if (server_options_->ip_addresses().empty()) {
+    // "::"  =>  in6addr_any
+    ev_uint16_t ev_port = static_cast<ev_uint16_t>(port);
+    ev_listener_ = evhttp_bind_socket_with_handle(ev_http_, "::", ev_port);
     if (ev_listener_ == nullptr) {
-      NET_LOG(ERROR, "Couldn't bind to port %d", port);
+      // in case ipv6 is not supported, fallback to inaddr_any
+      ev_listener_ = evhttp_bind_socket_with_handle(ev_http_, nullptr, ev_port);
+      if (ev_listener_ == nullptr) {
+        NET_LOG(ERROR, "Couldn't bind to port %d", port);
+        return false;
+      }
+    }
+  } else {
+    const std::string& ip_address = server_options_->ip_addresses().front();
+    ev_listener_ =
+        evhttp_bind_socket_with_handle(ev_http_, ip_address.c_str(), port);
+    if (ev_listener_ == nullptr) {
+      NET_LOG(ERROR, "Couldn't bind address %s to port %d", ip_address.c_str(),
+              port);
       return false;
     }
   }

tensorflow_serving/util/net_http/server/internal/evhttp_server_test.cc

@@ -57,16 +57,18 @@ class EvHTTPServerTest : public ::testing::Test {
   }
 
  protected:
-  std::unique_ptr<HTTPServerInterface> server;
-
- private:
-  void InitServer() {
+  virtual std::unique_ptr<ServerOptions> GetOptions() {
     auto options = absl::make_unique<ServerOptions>();
     options->AddPort(0);
     options->SetExecutor(absl::make_unique<MyExecutor>(4));
+    return options;
+  }
 
-    server = CreateEvHTTPServer(std::move(options));
+  std::unique_ptr<HTTPServerInterface> server;
 
+ private:
+  void InitServer() {
+    server = CreateEvHTTPServer(GetOptions());
     ASSERT_TRUE(server != nullptr);
   }
 };
@@ -339,7 +341,27 @@ TEST_F(EvHTTPServerTest, ActiveRequestCountInShutdown) {
   // response.status etc are undefined as the server is terminated
 }
 
+class EvHTTPServerTestWithAddress : public EvHTTPServerTest {
+ protected:
+  std::unique_ptr<ServerOptions> GetOptions() override {
+    auto options = EvHTTPServerTest::GetOptions();
+    options->AddIPAddress("::1");
+    return options;
+  }
+};
+
+TEST_F(EvHTTPServerTestWithAddress, BasicListen) {
+  server->StartAcceptingRequests();
+
+  EXPECT_TRUE(server->is_accepting_requests());
+  EXPECT_NE(server->listen_port(), 0);
+
+  server->Terminate();
+  server->WaitForTermination();
+}
+
 }  // namespace
+
 }  // namespace net_http
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/util/net_http/server/public/httpserver_interface.h

@@ -19,14 +19,13 @@ limitations under the License.
 #define TENSORFLOW_SERVING_UTIL_NET_HTTP_SERVER_PUBLIC_HTTPSERVER_INTERFACE_H_
 
 #include <cassert>
-
 #include <functional>
 #include <memory>
+#include <string>
 #include <vector>
 
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
-
 #include "tensorflow_serving/util/net_http/server/public/server_request_interface.h"
 
 namespace tensorflow {
@@ -62,18 +61,26 @@ class ServerOptions {
     ports_.emplace_back(port);
   }
 
+  // Add an IP address to listen on.  If not specified, the server will listen
+  // on all addresses.
+  void AddIPAddress(absl::string_view ip_address) {
+    ip_addresses_.emplace_back(ip_address);
+  }
+
   // The default executor for running I/O event polling.
   // This is a mandatory option.
   void SetExecutor(std::unique_ptr<EventExecutor> executor) {
     executor_ = std::move(executor);
   }
 
+  const std::vector<std::string>& ip_addresses() const { return ip_addresses_; }
   const std::vector<int>& ports() const { return ports_; }
 
   EventExecutor* executor() const { return executor_.get(); }
 
  private:
   std::vector<int> ports_;
+  std::vector<std::string> ip_addresses_;
   std::unique_ptr<EventExecutor> executor_;
 };