| 1 | """ |
|---|
| 2 | Ported to Python 3. |
|---|
| 3 | """ |
|---|
| 4 | |
|---|
| 5 | from autobahn.twisted.resource import WebSocketResource |
|---|
| 6 | from autobahn.twisted.websocket import ( |
|---|
| 7 | WebSocketServerFactory, |
|---|
| 8 | WebSocketServerProtocol, |
|---|
| 9 | ) |
|---|
| 10 | import eliot |
|---|
| 11 | |
|---|
| 12 | from twisted.web.resource import ( |
|---|
| 13 | Resource, |
|---|
| 14 | ) |
|---|
| 15 | |
|---|
| 16 | from allmydata.util import jsonbytes as json |
|---|
| 17 | |
|---|
| 18 | |
|---|
| 19 | class TokenAuthenticatedWebSocketServerProtocol(WebSocketServerProtocol): |
|---|
| 20 | """ |
|---|
| 21 | A WebSocket protocol that looks for an `Authorization:` header |
|---|
| 22 | with a `tahoe-lafs` scheme and a token matching our private config |
|---|
| 23 | for `api_auth_token`. |
|---|
| 24 | """ |
|---|
| 25 | |
|---|
| 26 | def onConnect(self, req): |
|---|
| 27 | """ |
|---|
| 28 | WebSocket callback |
|---|
| 29 | """ |
|---|
| 30 | # we don't care what WebSocket sub-protocol is |
|---|
| 31 | # negotiated, nor do we need to send headers to the |
|---|
| 32 | # client, so we ask Autobahn to just allow this |
|---|
| 33 | # connection with the defaults. We could return a |
|---|
| 34 | # (headers, protocol) pair here instead if required. |
|---|
| 35 | return None |
|---|
| 36 | |
|---|
| 37 | def _received_eliot_log(self, message): |
|---|
| 38 | """ |
|---|
| 39 | While this WebSocket connection is open, this function is |
|---|
| 40 | registered as an eliot destination |
|---|
| 41 | """ |
|---|
| 42 | # probably want a try/except around here? what do we do if |
|---|
| 43 | # transmission fails or anything else bad happens? |
|---|
| 44 | encoded = json.dumps_bytes(message, any_bytes=True) |
|---|
| 45 | self.sendMessage(encoded) |
|---|
| 46 | |
|---|
| 47 | def onOpen(self): |
|---|
| 48 | """ |
|---|
| 49 | WebSocket callback |
|---|
| 50 | """ |
|---|
| 51 | eliot.add_destination(self._received_eliot_log) |
|---|
| 52 | |
|---|
| 53 | def onClose(self, wasClean, code, reason): |
|---|
| 54 | """ |
|---|
| 55 | WebSocket callback |
|---|
| 56 | """ |
|---|
| 57 | try: |
|---|
| 58 | eliot.remove_destination(self._received_eliot_log) |
|---|
| 59 | except ValueError: |
|---|
| 60 | pass |
|---|
| 61 | |
|---|
| 62 | |
|---|
| 63 | def create_log_streaming_resource(): |
|---|
| 64 | factory = WebSocketServerFactory() |
|---|
| 65 | factory.protocol = TokenAuthenticatedWebSocketServerProtocol |
|---|
| 66 | return WebSocketResource(factory) |
|---|
| 67 | |
|---|
| 68 | |
|---|
| 69 | def create_log_resources(): |
|---|
| 70 | logs = Resource() |
|---|
| 71 | logs.putChild(b"v1", create_log_streaming_resource()) |
|---|
| 72 | return logs |
|---|