Bug 1716963 - Guard access to marionette socket with a lock, r=webdriver-reviewers,whimboo

This should ensure that we can't end up with multiple threads
interleaving reads or writes on the socket.

Differential Revision: https://phabricator.services.mozilla.com/D118148

testing/marionette/client/marionette_driver/transport.py

@@ -8,22 +8,23 @@ import json
 import socket
 import sys
 import time
+from threading import RLock
 
 import six
 
 
 class SocketTimeout(object):
-    def __init__(self, socket, timeout):
-        self.sock = socket
+    def __init__(self, socket_ctx, timeout):
+        self.socket_ctx = socket_ctx
         self.timeout = timeout
         self.old_timeout = None
 
     def __enter__(self):
-        self.old_timeout = self.sock.gettimeout()
-        self.sock.settimeout(self.timeout)
+        self.old_timeout = self.socket_ctx.socket_timeout
+        self.socket_ctx.socket_timeout = self.timeout
 
     def __exit__(self, *args, **kwargs):
-        self.sock.settimeout(self.old_timeout)
+        self.socket_ctx.socket_timeout = self.old_timeout
 
 
 class Message(object):
@@ -90,6 +91,35 @@ class Response(Message):
         return Response(data[1], data[2], data[3])
 
 
+class SocketContext(object):
+    """Object that guards access to a socket via a lock.
+
+    The socket must be accessed using this object as a context manager;
+    access to the socket outside of a context will bypass the lock."""
+
+    def __init__(self, host, port, timeout):
+        self.lock = RLock()
+
+        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self._sock.settimeout(timeout)
+        self._sock.connect((host, port))
+
+    @property
+    def socket_timeout(self):
+        return self._sock.gettimeout()
+
+    @socket_timeout.setter
+    def socket_timeout(self, value):
+        self._sock.settimeout(value)
+
+    def __enter__(self):
+        self.lock.acquire()
+        return self._sock
+
+    def __exit__(self, *args, **kwargs):
+        self.lock.release()
+
+
 class TcpTransport(object):
     """Socket client that communciates with Marionette via TCP.
 
@@ -111,11 +141,11 @@ class TcpTransport(object):
         will be used.  Setting it to `1` or `None` disables timeouts on
         socket operations altogether.
         """
-        self._sock = None
+        self._socket_context = None
 
         self.host = host
         self.port = port
-        self.socket_timeout = socket_timeout
+        self._socket_timeout = socket_timeout
 
         self.protocol = self.min_protocol_level
         self.application_type = None
@@ -130,8 +160,8 @@ class TcpTransport(object):
     def socket_timeout(self, value):
         self._socket_timeout = value
 
-        if self._sock:
-            self._sock.settimeout(value)
+        if self._socket_context is not None:
+            self._socket_context.socket_timeout = value
 
     def _unmarshal(self, packet):
         msg = None
@@ -168,6 +198,7 @@ class TcpTransport(object):
         # is 4 bytes: "2:{}". In practice the marionette format has some required fields so the
         # message is longer, but 4 bytes allows reading messages with bodies up to 999 bytes in
         # length in two reads, which is the common case.
+        with self._socket_context as sock:
             recv_bytes = 4
 
             length_prefix = b""
@@ -188,7 +219,7 @@ class TcpTransport(object):
                     )
 
                 try:
-                chunk = self._sock.recv(recv_bytes)
+                    chunk = sock.recv(recv_bytes)
                 except OSError:
                     continue
 
@@ -229,10 +260,11 @@ class TcpTransport(object):
                         body_part = parts[1]
 
                     # If we didn't find a : yet we keep reading 4 bytes at a time until we do.
-                # We could increase this here to 7 bytes (since we can't have more than 10 length
-                # bytes and a seperator byte), or just increase it to int(length_prefix) + 1 since
-                # that's the minimum total number of remaining bytes (if the : is in the next
-                # byte), but it's probably not worth optimising for large messages.
+                    # We could increase this here to 7 bytes (since we can't have more than 10
+                    # length bytes and a seperator byte), or just increase it to
+                    # int(length_prefix) + 1 since that's the minimum total number of remaining
+                    # bytes (if the : is in the next byte), but it's probably not worth optimising
+                    # for large messages.
 
                 if body_part is not None:
                     body_received += len(body_part)
@@ -259,18 +291,17 @@ class TcpTransport(object):
         Returns a tuple of the protocol level and the application type.
         """
         try:
-            self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            self._sock.settimeout(self.socket_timeout)
-
-            self._sock.connect((self.host, self.port))
+            self._socket_context = SocketContext(
+                self.host, self.port, self._socket_timeout
+            )
         except Exception:
             # Unset so that the next attempt to send will cause
             # another connection attempt.
-            self._sock = None
+            self._socket_context = None
             raise
 
         try:
-            with SocketTimeout(self._sock, 60.0):
+            with SocketTimeout(self._socket_context, 60.0):
                 # first packet is always a JSON Object
                 # which we can use to tell which protocol level we are at
                 raw = self.receive(unmarshal=False)
@@ -301,7 +332,7 @@ class TcpTransport(object):
         """Send message to the remote server.  Allowed input is a
         ``Message`` instance or a JSON serialisable object.
         """
-        if not self._sock:
+        if not self._socket_context:
             self.connect()
 
         if isinstance(obj, Message):
@@ -313,9 +344,10 @@ class TcpTransport(object):
         data = six.ensure_binary(data)
         payload = six.ensure_binary(str(len(data))) + b":" + data
 
+        with self._socket_context as sock:
             totalsent = 0
             while totalsent < len(payload):
-            sent = self._sock.send(payload[totalsent:])
+                sent = sock.send(payload[totalsent:])
                 if sent == 0:
                     raise IOError(
                         "Socket error after sending {0} of {1} bytes".format(
@@ -355,9 +387,10 @@ class TcpTransport(object):
 
         See: https://docs.python.org/2/howto/sockets.html#disconnecting
         """
-        if self._sock:
+        if self._socket_context:
+            with self._socket_context as sock:
                 try:
-                self._sock.shutdown(socket.SHUT_RDWR)
+                    sock.shutdown(socket.SHUT_RDWR)
                 except IOError as exc:
                     # If the socket is already closed, don't care about:
                     #   Errno  57: Socket not connected
@@ -365,10 +398,10 @@ class TcpTransport(object):
                     if exc.errno not in (57, 107):
                         raise
 
-            if self._sock:
+                if sock:
                     # Guard against unclean shutdown.
-                self._sock.close()
-                self._sock = None
+                    sock.close()
+                    self._socket_context = None
 
     def __del__(self):
         self.close()

testing/marionette/harness/marionette_harness/tests/unit/test_marionette.py

@@ -69,7 +69,8 @@ class TestMarionette(MarionetteTestCase):
 
         self.assertEqual(current_socket_timeout, self.marionette.client.socket_timeout)
         self.assertEqual(
-            current_socket_timeout, self.marionette.client._sock.gettimeout()
+            current_socket_timeout,
+            self.marionette.client._socket_context._sock.gettimeout(),
         )
 
     def test_application_update_disabled(self):