Arthur de Jong

Open Source / Free Software developer

summaryrefslogtreecommitdiffstats
path: root/channel.py
blob: 5dfc19ffc4dbaa5ed77d693b577110986a6690d3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
# Webchat WebRTC application
# https://arthurdejong.org/webchat/
#
# Copyright (C) 2020 Arthur de Jong
#
# Released under the GNU General Public License, either version 3, or
# (at your option) any later version.

"""Simple server that relays channel messages across websockets."""

import asyncio
import collections
import logging

import aiohttp
import aiohttp.web


routes = aiohttp.web.RouteTableDef()


websockets = collections.defaultdict(set)


async def send_to_all(channel, data):
    """Send the data to all websockets for the channel."""
    tasks = set(
        asyncio.ensure_future(ws.send_bytes(data))
        for ws in websockets[channel])
    while tasks:
        done, tasks = await asyncio.wait(tasks)


@routes.get('/channel/{channel}')
async def get_channel(request):
    """Process a new connection by returning a websocket."""
    channel = request.match_info['channel']
    ws = aiohttp.web.WebSocketResponse()
    await ws.prepare(request)
    websockets[channel].add(ws)
    try:
        # keep relaying all messages
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.BINARY:
                await send_to_all(channel, msg.data)
    finally:
        websockets[channel].remove(ws)
    return ws


if __name__ == '__main__':
    app = aiohttp.web.Application()
    app.add_routes(routes)
    logging.basicConfig(level=logging.INFO)
    aiohttp.web.run_app(app, host='localhost')