Use auth

pico.js

@@ -70,10 +70,12 @@ const Main = {
 m.mount(document.body, Main)
 
 const connect = (username) => {
+    const wsUrl = location.href
+        .replace('http', 'ws')
+        .replace(/:\/\//, `://${username}@`)
+
     State.websocket = new WebSocket(wsUrl)
-    State.websocket.onopen = (e) => {
-        State.websocket.send(JSON.stringify({action: 'login', username}))
-    }
+
     State.websocket.onmessage = (e) => {
         const message = JSON.parse(e.data)
         if(message.kind === 'post') {
@@ -94,6 +96,7 @@ const connect = (username) => {
         }
         m.redraw()
     }
+
     State.websocket.onclose = (e) => {
         if(!e.wasClean) {
             setTimeout(connect, 100, username)
@@ -101,7 +104,6 @@ const connect = (username) => {
         m.redraw()
     }
 }
-const wsUrl = location.toString().replace('http', 'ws')
 if(localStorage.username) {
     connect(localStorage.username)
 }

pico.py

@@ -7,9 +7,54 @@ import functools
 from pathlib import Path
 
 import websockets
+from websockets.headers import (
+    build_www_authenticate_basic,
+    parse_authorization_basic,
+)
 
 
-rooms = collections.defaultdict(dict)
+rooms = collections.defaultdict(set)
+
+
+class PicoProtocol(websockets.WebSocketServerProtocol):
+
+    def http_unauthorized(self, message, realm='pico'):
+        return (
+            http.HTTPStatus.UNAUTHORIZED,
+            [("WWW-Authenticate", build_www_authenticate_basic(realm))],
+            message.encode(),
+        )
+
+    def serve_file(self, path):
+        document = Path(__file__, '..', Path(path).name).resolve()
+        if not document.is_file():
+            document = Path(__file__, '..', 'pico.html').resolve()
+            content_type = 'text/html; charset=utf-8'
+        elif path.endswith('.js'):
+            content_type = 'application/javascript; charset=utf-8'
+        elif path.endswith('.css'):
+            content_type = 'text/css; charset=utf-8'
+        else:
+            content_type = 'text/plain; charset=utf-8'
+        return (
+            http.HTTPStatus.OK,
+            [('Content-Type', content_type)],
+            document.read_bytes(),
+        )
+
+    async def process_request(self, path, request_headers):
+        if request_headers.get('Upgrade') != 'websocket':
+            return self.serve_file(path)
+        try:
+            authorization = request_headers['Authorization']
+            self.username, password = parse_authorization_basic(authorization)
+        except Exception as error:
+            return self.http_unauthorized('No username found')
+        return await super().process_request(path, request_headers)
+
+
+def get_usernames(sockets):
+    return sorted(ws.username for ws in sockets)
 
 
 async def send_json_many(targets, **data):
@@ -33,69 +78,51 @@ async def recv_json(websocket):
         return {}
 
 
-async def core(websocket, path, server_name):
+async def core(ws, path, server_name):
     sockets = rooms[path]
+
     while True:
-        data = await recv_json(websocket)
+        if ws in sockets:
+            data = await recv_json(ws)
+        else:
+            data = {'action': 'login'}
+
         ts = datetime.datetime.now().isoformat()
-        reply = functools.partial(send_json, websocket=websocket, ts=ts)
+        reply = functools.partial(send_json, websocket=ws, ts=ts)
         error = functools.partial(reply, kind='error')
-        broadcast = functools.partial(send_json_many, targets=sockets, ts=ts)
+        broadcast = functools.partial(send_json_many, targets=set(sockets), ts=ts)
 
         if 'action' not in data:
             await error(info='Message without action is invalid')
 
-        elif websocket not in sockets and data['action'] not in {'login', 'logout'}:
-            await error(info='Not logged in')
-
         elif data['action'] == 'login':
-            username = data['username']
-            if not username:
-                await error(info='Invalid username')
-            elif username in sockets.values():
+            if not ws.username:
+                await error(info='Username not allowed')
+                break
+            if ws.username in get_usernames(sockets):
                 await error(info='Username taken')
-            else:
-                sockets[websocket] = username
-                await reply(kind='update', username=username)
-                await broadcast(kind='update', users=list(sockets.values()), info=f'{username} joined')
+                break
+
+            sockets.add(ws)
+            online = get_usernames(sockets)
+            await reply(kind='update', users=online, info=f'Welcome to {path}', username=ws.username)
+            await broadcast(kind='update', users=online, info=f'{ws.username} joined')
 
         elif data['action'] == 'post':
             text = data['text']
             if not text:
                 continue
-            await broadcast(kind='post', source=sockets[websocket], text=text)
+            await broadcast(kind='post', source=ws.username, text=text)
 
         elif data['action'] == 'logout':
-            if websocket in sockets:
-                username = sockets.pop(websocket)
-            await broadcast(kind='update', users=list(sockets.values()), info=f'{username} left')
+            sockets.discard(ws)
+            await broadcast(kind='update', users=get_usernames(sockets), info=f'{ws.username} left')
             break
 
 
-async def serve_html(path, request_headers):
-    if request_headers.get('Upgrade') != 'websocket':
-
-        document = Path(__file__, '..', Path(path).name).resolve()
-        if not document.is_file():
-            document = Path(__file__, '..', 'pico.html').resolve()
-            content_type = 'text/html; charset=utf-8'
-        elif path.endswith('.js'):
-            content_type = 'application/javascript; charset=utf-8'
-        elif path.endswith('.css'):
-            content_type = 'text/css; charset=utf-8'
-        else:
-            content_type = 'text/plain; charset=utf-8'
-
-        return (
-            http.HTTPStatus.OK,
-            [('Content-Type', content_type)],
-            document.read_bytes(),
-        )
-
-
 async def start_server(host, port, server_name):
     bound_core = functools.partial(core, server_name=server_name)
-    return await websockets.serve(bound_core, host, port, process_request=serve_html)
+    return await websockets.serve(bound_core, host, port, create_protocol=PicoProtocol)
 
 
 if __name__ == '__main__':

test.py

@@ -38,8 +38,8 @@ async def _make_client(path, timeout, script):
             elif kind == 'send':
                 await send_json(websocket, **message)
 
-        while True:
-            if state['users'][0] == state['username']:
+        while 'users' in state:
+            if min(state['users']) == state['username']:
                 break
             message = await recv_json(websocket)
             message.pop('ts')
@@ -59,52 +59,48 @@ async def _test(*clients):
             start_server('localhost', 8642, 'TestServer'),
             *clients,
         )
-        server, *_ = await asyncio.wait_for(gather, timeout=2)
+        server, *results = await asyncio.wait_for(gather, timeout=2)
         server.close()
     except asyncio.TimeoutError:
         return
 
+    return results
+
 
 def test_happy_path():
-    client = _make_client('ws://localhost:8642/', 0.1, Script()
-        + {'action': 'login', 'username': 'TestUser'}
-        - {'kind': 'update', 'username': 'TestUser'}
-        - {'kind': 'update', 'users': ['TestUser'], 'info': 'TestUser joined'}
+    client = _make_client('ws://TestUser@localhost:8642/A', 0.1, Script()
+        - {'kind': 'update', 'users': ['TestUser'], 'info': 'Welcome to /A', 'username': 'TestUser'}
         + {'action': 'post', 'text': 'Hello World!'}
         - {'kind': 'post', 'source': 'TestUser', 'text': 'Hello World!'}
     )
     return _test(client)
 
 
-def test_post_before_login():
-    client = _make_client('ws://localhost:8642/', 0.1, Script()
-        + {}
-        - {'kind': 'error', 'info': 'Message without action is invalid'}
-        + {'action': 'post', 'text': ''}
-        - {'kind': 'error', 'info': 'Not logged in'}
-        + {'action': 'login', 'username': ''}
-        - {'kind': 'error', 'info': 'Invalid username'}
-        + {'action': 'login', 'username': 'Joe'}
-        - {'kind': 'update', 'username': 'Joe'}
-        - {'kind': 'update', 'users': ['Joe'], 'info': 'Joe joined'}
+def test_name_taken():
+    client1 = _make_client('ws://A@localhost:8642/', 0.10, Script()
+        - {'kind': 'update', 'users': ['A'], 'info': 'Welcome to /', 'username': 'A'}
+        - {'kind': 'update', 'users': ['A', 'B'], 'info': 'B joined'}
     )
-    return _test(client)
+    client2 = _make_client('ws://A@localhost:8642/', 0.11, Script()
+        - {'kind': 'error', 'info': 'Username taken'}
+    )
+    client3 = _make_client('ws://B@localhost:8642/', 0.12, Script()
+        - {'kind': 'update', 'users': ['A', 'B'], 'info': 'Welcome to /', 'username': 'B'}
+        - {'kind': 'update', 'users': ['B'], 'info': 'A left'}
+    )
+    return _test(client1, client2, client3)
 
 
 def test_interact():
-    client1 = _make_client('ws://localhost:8642/', 0.10, Script()
-        + {'action': 'login', 'username': 'Alice'}
-        - {'kind': 'update', 'username': 'Alice'}
-        - {'kind': 'update', 'users': ['Alice'], 'info': 'Alice joined'}
+    client1 = _make_client('ws://Alice@localhost:8642/', 0.10, Script()
+        - {'kind': 'update', 'users': ['Alice'], 'info': 'Welcome to /', 'username': 'Alice'}
         - {'kind': 'update', 'users': ['Alice', 'Bob'], 'info': 'Bob joined'}
         + {'action': 'post', 'text': 'Hey Bob!'}
         - {'kind': 'post', 'source': 'Alice', 'text': 'Hey Bob!'}
         - {'kind': 'post', 'source': 'Bob', 'text': 'Howdy!'}
     )
-    client2 = _make_client('ws://localhost:8642/', 0.11, Script()
-        + {'action': 'login', 'username': 'Bob'}
-        - {'kind': 'update', 'username': 'Bob'}
-        - {'kind': 'update', 'users': ['Alice', 'Bob'], 'info': 'Bob joined'}
+    client2 = _make_client('ws://Bob@localhost:8642/', 0.11, Script()
+        - {'kind': 'update', 'users': ['Alice', 'Bob'], 'info': 'Welcome to /', 'username': 'Bob'}
         - {'kind': 'post', 'source': 'Alice', 'text': 'Hey Bob!'}
         + {'action': 'post', 'text': 'Howdy!'}
         - {'kind': 'post', 'source': 'Bob', 'text': 'Howdy!'}
@@ -115,46 +111,36 @@ def test_interact():
 
 
 def test_party():
-    client1 = _make_client('ws://localhost:8642/', 0.10, Script()
-        + {'action': 'login', 'username': 'Norman'}
-        - {'kind': 'update', 'username': 'Norman'}
-        - {'kind': 'update', 'users': ['Norman'], 'info': 'Norman joined'}
+    client1 = _make_client('ws://Norman@localhost:8642/', 0.10, Script()
+        - {'kind': 'update', 'users': ['Norman'], 'info': 'Welcome to /', 'username': 'Norman'}
         - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Ray joined'}
-        - {'kind': 'update', 'users': ['Norman', 'Ray', 'Emma'], 'info': 'Emma joined'}
+        - {'kind': 'update', 'users': ['Emma', 'Norman', 'Ray'], 'info': 'Emma joined'}
         + {'action': 'post', 'text': 'なに？'}
         - {'kind': 'post', 'source': 'Norman', 'text': 'なに？'}
+        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Emma left'}
     )
-    client2 = _make_client('ws://localhost:8642/', 0.11, Script()
-        + {'action': 'login', 'username': 'Ray'}
-        - {'kind': 'update', 'username': 'Ray'}
-        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Ray joined'}
-        - {'kind': 'update', 'users': ['Norman', 'Ray', 'Emma'], 'info': 'Emma joined'}
+    client2 = _make_client('ws://Ray@localhost:8642/', 0.11, Script()
+        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Welcome to /', 'username': 'Ray'}
+        - {'kind': 'update', 'users': ['Emma', 'Norman', 'Ray'], 'info': 'Emma joined'}
         - {'kind': 'post', 'source': 'Norman', 'text': 'なに？'}
-        - {'kind': 'update', 'users': ['Ray', 'Emma'], 'info': 'Norman left'}
+        - {'kind': 'update', 'users': ['Norman', 'Ray'], 'info': 'Emma left'}
+        - {'kind': 'update', 'users': ['Ray'], 'info': 'Norman left'}
     )
-    client3 = _make_client('ws://localhost:8642/', 0.12, Script()
-        + {'action': 'login', 'username': 'Emma'}
-        - {'kind': 'update', 'username': 'Emma'}
-        - {'kind': 'update', 'users': ['Norman', 'Ray', 'Emma'], 'info': 'Emma joined'}
+    client3 = _make_client('ws://Emma@localhost:8642/', 0.12, Script()
+        - {'kind': 'update', 'users': ['Emma', 'Norman', 'Ray'], 'info': 'Welcome to /', 'username': 'Emma'}
         - {'kind': 'post', 'source': 'Norman', 'text': 'なに？'}
-        - {'kind': 'update', 'users': ['Ray', 'Emma'], 'info': 'Norman left'}
-        - {'kind': 'update', 'users': ['Emma'], 'info': 'Ray left'}
     )
     return _test(client1, client2, client3)
 
 
 def test_rooms():
-    client1 = _make_client('ws://localhost:8642/A', 0.10, Script()
-        + {'action': 'login', 'username': 'Dandy'}
-        - {'kind': 'update', 'username': 'Dandy'}
-        - {'kind': 'update', 'users': ['Dandy'], 'info': 'Dandy joined'}
+    client1 = _make_client('ws://Dandy@localhost:8642/A', 0.10, Script()
+        - {'kind': 'update', 'users': ['Dandy'], 'info': 'Welcome to /A', 'username': 'Dandy'}
         + {'action': 'post', 'text': 'Hi'}
         - {'kind': 'post', 'source': 'Dandy', 'text': 'Hi'}
     )
-    client2 = _make_client('ws://localhost:8642/B', 0.10, Script()
-        + {'action': 'login', 'username': 'Dandy'}
-        - {'kind': 'update', 'username': 'Dandy'}
-        - {'kind': 'update', 'users': ['Dandy'], 'info': 'Dandy joined'}
+    client2 = _make_client('ws://Dandy@localhost:8642/B', 0.10, Script()
+        - {'kind': 'update', 'users': ['Dandy'], 'info': 'Welcome to /B', 'username': 'Dandy'}
         + {'action': 'post', 'text': 'Howdy'}
         - {'kind': 'post', 'source': 'Dandy', 'text': 'Howdy'}
     )