portal 3.6.0__tar.gz → 3.7.0__tar.gz

{portal-3.6.0/portal.egg-info → portal-3.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: portal
-Version: 3.6.0
+Version: 3.7.0
 Summary: Fast and reliable distributed systems in Python
 Home-page: http://github.com/danijar/portal
 Author: Danijar Hafner

{portal-3.6.0 → portal-3.7.0}/portal/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '3.6.0'
+__version__ = '3.7.0'
 
 import multiprocessing as mp
 try:

{portal-3.6.0 → portal-3.7.0}/portal/client_socket.py

@@ -31,6 +31,7 @@ class Options:
   logging: bool = True
   logging_color: str = 'yellow'
   connect_wait: float = 0.1
+  handshake: str = 'portal_handshake'
 
 
 class ClientSocket:
@@ -205,8 +206,9 @@ class ClientSocket:
       try:
         sock.settimeout(10)
         sock.connect(addr)
-        sock.settimeout(0)
         self._log('Connection established')
+        self._send_handshake(sock)
+        sock.settimeout(0)
         return sock
       except TimeoutError as e:
         error = e
@@ -250,6 +252,11 @@ class ClientSocket:
 
     return sock
 
+  def _send_handshake(self, sock):
+    buf = buffers.SendBuffer(self.options.handshake.encode('utf-8'))
+    while not buf.done():
+      buf.send(sock)
+
   def _log(self, *args):
     if not self.options.logging:
       return

{portal-3.6.0 → portal-3.7.0}/portal/server_socket.py

@@ -16,6 +16,7 @@ class Connection:
     self.sock = sock
     self.addr = addr
     self.recvbuf = None
+    self.handshake = False
     self.sendbufs = collections.deque()
 
   def fileno(self):
@@ -32,6 +33,7 @@ class Options:
   max_send_queue: int = 4096
   logging: bool = True
   logging_color: str = 'blue'
+  handshake: str = 'portal_handshake'
 
 
 class ServerSocket:
@@ -157,11 +159,20 @@ class ServerSocket:
       # - TimeoutError: [Errno 110] Connection timed out
       self._disconnect(conn, e)
       return
-    if conn.recvbuf.done():
-      if self.recvq.qsize() > self.options.max_recv_queue:
-        raise RuntimeError('Too many incoming messages enqueued')
-      self.recvq.put((conn.addr, conn.recvbuf.result()))
-      conn.recvbuf = None
+    if not conn.recvbuf.done():
+      return
+    if self.recvq.qsize() > self.options.max_recv_queue:
+      raise RuntimeError('Too many incoming messages enqueued')
+    message = conn.recvbuf.result()
+    conn.recvbuf = None
+    if not conn.handshake:
+      if message != self.options.handshake.encode('utf-8'):
+        e = ValueError(f"Handshake '{self.options.handshake}' expected")
+        self._disconnect(conn, e)
+        return
+      conn.handshake = True
+    else:
+      self.recvq.put((conn.addr, message))
 
   def _disconnect(self, conn, e):
     detail = f'{type(e).__name__}'

{portal-3.6.0 → portal-3.7.0/portal.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: portal
-Version: 3.6.0
+Version: 3.7.0
 Summary: Fast and reliable distributed systems in Python
 Home-page: http://github.com/danijar/portal
 Author: Danijar Hafner

{portal-3.6.0 → portal-3.7.0}/tests/test_socket.py

@@ -193,3 +193,14 @@ class TestSocket:
        portal.Process(server, port),
        portal.Process(client, port),
    ])
+
+  @pytest.mark.parametrize('ipv6', (False, True))
+  def test_wrong_protocol(self, ipv6):
+    port = portal.free_port()
+    server = portal.ServerSocket(port, ipv6=ipv6)
+    client = portal.ClientSocket(
+        port, ipv6=ipv6, autoconn=False, handshake='unknown_word')
+    client.connect()
+    with pytest.raises(portal.Disconnected):
+      client.recv()
+    server.close()