Refactor server code (use connections)

Past version was unusable with >1 client.

secret_handshake/network.py

@@ -19,7 +19,7 @@
 # SOFTWARE.
 
 
-from asyncio import open_connection, start_server, ensure_future
+import asyncio
 
 from async_generator import async_generator, yield_
 
@@ -31,10 +31,10 @@ class SHSClientException(Exception):
     pass
 
 
-class SHSSocket(object):
-    def __init__(self, loop):
-        self.loop = loop
-        self._on_connect = None
+class SHSDuplexStream(object):
+    def __init__(self):
+        self.write_stream = None
+        self.read_stream = None
 
     def write(self, data):
         self.write_stream.write(data)
@@ -42,24 +42,35 @@ class SHSSocket(object):
     async def read(self):
         return await self.read_stream.read()
 
-    def disconnect(self):
-        self.writer.close()
+    def close(self):
+        self.write_stream.close()
+        self.read_stream.close()
 
     @async_generator
     async def __aiter__(self):
         async for msg in self.read_stream:
             await yield_(msg)
 
+
+class SHSEndpoint(object):
+    def __init__(self):
+        self._on_connect = None
+        self.crypto = None
+
     def on_connect(self, cb):
         self._on_connect = cb
 
+    def disconnect(self):
+        raise NotImplementedError
 
-class SHSServer(SHSSocket):
-    def __init__(self, host, port, server_kp, application_key=None, loop=None):
-        super(SHSServer, self).__init__(loop)
+
+class SHSServer(SHSEndpoint):
+    def __init__(self, host, port, server_kp, application_key=None):
+        super(SHSServer, self).__init__()
         self.host = host
         self.port = port
         self.crypto = SHSServerCrypto(server_kp, application_key=application_key)
+        self.connections = []
 
     async def _handshake(self, reader, writer):
         data = await reader.readexactly(64)
@@ -77,23 +88,39 @@ class SHSServer(SHSSocket):
     async def handle_connection(self, reader, writer):
         self.crypto.clean()
         await self._handshake(reader, writer)
-
         keys = self.crypto.get_box_keys()
         self.crypto.clean()
 
-        self.read_stream, self.write_stream = get_stream_pair(reader, writer, **keys)
-        self.writer = writer
+        conn = SHSServerConnection.from_byte_streams(reader, writer, **keys)
+        self.connections.append(conn)
 
         if self._on_connect:
-            ensure_future(self._on_connect(), loop=self.loop)
+            asyncio.ensure_future(self._on_connect(conn))
 
-    def listen(self):
-        self.loop.run_until_complete(start_server(self.handle_connection, self.host, self.port, loop=self.loop))
+    async def listen(self):
+        await asyncio.start_server(self.handle_connection, self.host, self.port)
+
+    def disconnect(self):
+        for connection in self.connections:
+            connection.close()
 
 
-class SHSClient(SHSSocket):
-    def __init__(self, host, port, client_kp, server_pub_key, ephemeral_key=None, application_key=None, loop=None):
-        super(SHSClient, self).__init__(loop)
+class SHSServerConnection(SHSDuplexStream):
+    def __init__(self, read_stream, write_stream):
+        super(SHSServerConnection, self).__init__()
+        self.read_stream = read_stream
+        self.write_stream = write_stream
+
+    @classmethod
+    def from_byte_streams(cls, reader, writer, **keys):
+        reader, writer = get_stream_pair(reader, writer, **keys)
+        return cls(reader, writer)
+
+
+class SHSClient(SHSDuplexStream, SHSEndpoint):
+    def __init__(self, host, port, client_kp, server_pub_key, ephemeral_key=None, application_key=None):
+        SHSDuplexStream.__init__(self)
+        SHSEndpoint.__init__(self)
         self.host = host
         self.port = port
         self.crypto = SHSClientCrypto(client_kp, server_pub_key, ephemeral_key=ephemeral_key,
@@ -112,8 +139,8 @@ class SHSClient(SHSSocket):
         if not self.crypto.verify_server_accept(data):
             raise SHSClientException('Server accept is not valid')
 
-    async def connect(self):
-        reader, writer = await open_connection(self.host, self.port, loop=self.loop)
+    async def open(self):
+        reader, writer = await asyncio.open_connection(self.host, self.port)
         await self._handshake(reader, writer)
 
         keys = self.crypto.get_box_keys()
@@ -123,3 +150,6 @@ class SHSClient(SHSSocket):
         self.writer = writer
         if self._on_connect:
             await self._on_connect()
+
+    def disconnect(self):
+        self.close()