1 files changed, 180 insertions, 0 deletions
diff --git a/lib/thrift/transport/TSocket.py b/lib/thrift/transport/TSocket.py
new file mode 100644
index 0000000..7b564aa
--- /dev/null
+++ b/lib/thrift/transport/TSocket.py
@@ -0,0 +1,180 @@
+#
+# 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.
+#
+
+import errno
+import os
+import socket
+import sys
+
+from TTransport import *
+
+
+class TSocketBase(TTransportBase):
+  def _resolveAddr(self):
+    if self._unix_socket is not None:
+      return [(socket.AF_UNIX, socket.SOCK_STREAM, None, None,
+               self._unix_socket)]
+    else:
+      return socket.getaddrinfo(self.host,
+                                self.port,
+                                self._socket_family,
+                                socket.SOCK_STREAM,
+                                0,
+                                socket.AI_PASSIVE | socket.AI_ADDRCONFIG)
+
+  def close(self):
+    if self.handle:
+      self.handle.close()
+      self.handle = None
+
+
+class TSocket(TSocketBase):
+  """Socket implementation of TTransport base."""
+
+  def __init__(self, host='localhost', port=9090, unix_socket=None, socket_family=socket.AF_UNSPEC):
+    """Initialize a TSocket
+
+    @param host(str)  The host to connect to.
+    @param port(int)  The (TCP) port to connect to.
+    @param unix_socket(str)  The filename of a unix socket to connect to.
+                             (host and port will be ignored.)
+    @param socket_family(int)  The socket family to use with this socket.
+    """
+    self.host = host
+    self.port = port
+    self.handle = None
+    self._unix_socket = unix_socket
+    self._timeout = None
+    self._socket_family = socket_family
+
+  def setHandle(self, h):
+    self.handle = h
+
+  def isOpen(self):
+    return self.handle is not None
+
+  def setTimeout(self, ms):
+    if ms is None:
+      self._timeout = None
+    else:
+      self._timeout = ms / 1000.0
+
+    if self.handle is not None:
+      self.handle.settimeout(self._timeout)
+
+  def open(self):
+    try:
+      res0 = self._resolveAddr()
+      for res in res0:
+        self.handle = socket.socket(res[0], res[1])
+        self.handle.settimeout(self._timeout)
+        try:
+          self.handle.connect(res[4])
+        except socket.error as e:
+          if res is not res0[-1]:
+            continue
+          else:
+            raise e
+        break
+    except socket.error as e:
+      if self._unix_socket:
+        message = 'Could not connect to socket %s' % self._unix_socket
+      else:
+        message = 'Could not connect to %s:%d' % (self.host, self.port)
+      raise TTransportException(type=TTransportException.NOT_OPEN,
+                                message=message)
+
+  def read(self, sz):
+    try:
+      buff = self.handle.recv(sz)
+    except socket.error as e:
+      if (e.args[0] == errno.ECONNRESET and
+          (sys.platform == 'darwin' or sys.platform.startswith('freebsd'))):
+        # freebsd and Mach don't follow POSIX semantic of recv
+        # and fail with ECONNRESET if peer performed shutdown.
+        # See corresponding comment and code in TSocket::read()
+        # in lib/cpp/src/transport/TSocket.cpp.
+        self.close()
+        # Trigger the check to raise the END_OF_FILE exception below.
+        buff = ''
+      else:
+        raise
+    if len(buff) == 0:
+      raise TTransportException(type=TTransportException.END_OF_FILE,
+                                message='TSocket read 0 bytes')
+    return buff
+
+  def write(self, buff):
+    if not self.handle:
+      raise TTransportException(type=TTransportException.NOT_OPEN,
+                                message='Transport not open')
+    sent = 0
+    have = len(buff)
+    while sent < have:
+      plus = self.handle.send(buff)
+      if plus == 0:
+        raise TTransportException(type=TTransportException.END_OF_FILE,
+                                  message='TSocket sent 0 bytes')
+      sent += plus
+      buff = buff[plus:]
+
+  def flush(self):
+    pass
+
+
+class TServerSocket(TSocketBase, TServerTransportBase):
+  """Socket implementation of TServerTransport base."""
+
+  def __init__(self, host=None, port=9090, unix_socket=None, socket_family=socket.AF_UNSPEC):
+    self.host = host
+    self.port = port
+    self._unix_socket = unix_socket
+    self._socket_family = socket_family
+    self.handle = None
+
+  def listen(self):
+    res0 = self._resolveAddr()
+    socket_family = self._socket_family == socket.AF_UNSPEC and socket.AF_INET6 or self._socket_family
+    for res in res0:
+      if res[0] is socket_family or res is res0[-1]:
+        break
+
+    # We need remove the old unix socket if the file exists and
+    # nobody is listening on it.
+    if self._unix_socket:
+      tmp = socket.socket(res[0], res[1])
+      try:
+        tmp.connect(res[4])
+      except socket.error as err:
+        eno, message = err.args
+        if eno == errno.ECONNREFUSED:
+          os.unlink(res[4])
+
+    self.handle = socket.socket(res[0], res[1])
+    self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    if hasattr(self.handle, 'settimeout'):
+      self.handle.settimeout(None)
+    self.handle.bind(res[4])
+    self.handle.listen(128)
+
+  def accept(self):
+    client, addr = self.handle.accept()
+    result = TSocket()
+    result.setHandle(client)
+    return result