Fix segfault in grpc_server_lib when an invalid hostname-port pair is passed.

mrry · tensorflower-gardener · commit 95ae233e25f8 · 2016-04-08T17:11:28.000-07:00
Previously, if the port was undefined, an out-of-bounds access would
be made. This change adds the appropriate checks.
Change: 119424297
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc
@@ -103,10 +103,14 @@ Status GrpcServer::Init() {
         return errors::InvalidArgument("Task ", server_def_.task_index(),
                                        " was not defined in job \"",
                                        server_def_.job_name(), "\"");
-      } else if (!strings::safe_strto32(str_util::Split(iter->second, ':')[1],
-                                        &requested_port_)) {
-        return errors::Internal("Could not parse port for local server from \"",
-                                iter->second, "\"");
+      }
+      const std::vector<string> hostname_port =
+          str_util::Split(iter->second, ':');
+      if (hostname_port.size() != 2 ||
+          !strings::safe_strto32(hostname_port[1], &requested_port_)) {
+        return errors::InvalidArgument(
+            "Could not parse port for local server from \"", iter->second,
+            "\"");
       } else {
         break;
       }
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h
@@ -89,12 +89,12 @@ class GrpcServer : public ServerInterface {
 
   // Implementation of a TensorFlow master, and RPC polling thread.
   MasterEnv master_env_;
-  AsyncServiceInterface* master_service_;
+  AsyncServiceInterface* master_service_ = nullptr;
   std::unique_ptr<Thread> master_thread_ GUARDED_BY(mu_);
 
   // Implementation of a TensorFlow worker, and RPC polling thread.
   WorkerEnv worker_env_;
-  AsyncServiceInterface* worker_service_;
+  AsyncServiceInterface* worker_service_ = nullptr;
   std::unique_ptr<Thread> worker_thread_ GUARDED_BY(mu_);
 
   std::unique_ptr<::grpc::Server> server_ GUARDED_BY(mu_);
diff --git a/tensorflow/python/training/server_lib.py b/tensorflow/python/training/server_lib.py
@@ -22,6 +22,7 @@
 
 from tensorflow.core.protobuf import tensorflow_server_pb2
 from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.framework import errors
 from tensorflow.python.util import compat
 
 
@@ -128,7 +129,12 @@ def __init__(self,
     """
     server_def = _make_server_def(server_or_cluster_def,
                                   job_name, task_index, protocol)
-    self._server = pywrap_tensorflow.NewServer(server_def.SerializeToString())
+    try:
+      self._server = pywrap_tensorflow.NewServer(server_def.SerializeToString())
+    except pywrap_tensorflow.StatusNotOK as e:
+      # pylint: disable=protected-access
+      raise errors._make_specific_exception(None, None, e.error_message, e.code)
+      # pylint: enable=protected-access
     if start:
       self.start()
 
diff --git a/tensorflow/python/training/server_lib_test.py b/tensorflow/python/training/server_lib_test.py
@@ -79,6 +79,12 @@ def testLargeFeed(self):
       self.assertEqual(0.5, min_val)
       self.assertEqual(0.5, max_val)
 
+  def testInvalidHostname(self):
+    with self.assertRaisesRegexp(tf.errors.InvalidArgumentError, "port"):
+      _ = tf.train.Server({"local": ["localhost"]},
+                          job_name="local",
+                          task_index=0)
+
 
 class ServerDefTest(tf.test.TestCase):
 
