#!/usr/bin/env python3 """postgres-relay.py — a tiny WebSocket-to-TCP bridge for the Postgres playground on geekychris.github.io. NOTE: the gitwebsite app (and `gw serve`) now has this built in — its *companion* gateway on 127.0.0.1:7317 relays the same things, and the tool uses it first when the app is running. This script stays as the fallback for when the app is not running. Browsers cannot open TCP connections, so the playground speaks the PostgreSQL wire protocol over a WebSocket, and this script copies the bytes, unchanged, between each WebSocket and a new TCP connection to ONE Postgres server that you choose. Every browser session (A, B, C …) is its own WebSocket and therefore its own Postgres backend, so locks, SKIP LOCKED and LISTEN behave exactly as they do for any other client. Python standard library only. python3 postgres-relay.py # ws://127.0.0.1:6432 -> localhost:5432 python3 postgres-relay.py --target db.internal:5432 --port 6433 Safety: * It listens on 127.0.0.1 only unless you pass --host (do not expose it to a network). * It refuses WebSocket handshakes from web pages other than https://geekychris.github.io and http://localhost / http://127.0.0.1 (local previews). Add others with --allow-origin. * It forwards to the --target only: it is not an open proxy. * Postgres still authenticates every connection (SCRAM-SHA-256 or md5) — the relay never sees or stores a password in a form it can use, but the protocol between the browser and the relay is not encrypted (ws:// on localhost). Use a throwaway database and stop the relay when done. """ import argparse, base64, hashlib, socket, struct, sys, threading, time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer VERSION = '1.0' WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' ALLOWED = {'https://geekychris.github.io'} def origin_ok(o): if not o: return True # not a browser (e.g. a test client); browsers always send Origin return o in ALLOWED or o.startswith(('http://127.0.0.1:', 'http://localhost:')) or o in ('http://127.0.0.1', 'http://localhost') def read_exact(f, n): data = b'' while len(data) < n: chunk = f.read(n - len(data)) if not chunk: raise EOFError data += chunk return data class Relay(BaseHTTPRequestHandler): server_version = 'postgres-relay/' + VERSION protocol_version = 'HTTP/1.1' def log_message(self, fmt, *args): if not self.server.quiet: sys.stderr.write('%s %s\n' % (time.strftime('%H:%M:%S'), fmt % args)) def plain(self, code, text): body = text.encode() self.send_response(code) self.send_header('Content-Type', 'text/plain; charset=utf-8') self.send_header('Content-Length', str(len(body))) self.end_headers() self.wfile.write(body) def do_GET(self): if self.headers.get('Upgrade', '').lower() != 'websocket': return self.plain(200, f'postgres-relay {VERSION}: WebSocket -> {self.server.target}\n') o = self.headers.get('Origin', '') if not origin_ok(o): return self.plain(403, f'origin {o} not allowed (see --allow-origin)\n') key = self.headers.get('Sec-WebSocket-Key', '') if not key: return self.plain(400, 'missing Sec-WebSocket-Key\n') host, _, port = self.server.target.rpartition(':') try: up = socket.create_connection((host, int(port)), timeout=10) except Exception as e: return self.plain(502, f'cannot connect to {self.server.target}: {e}\n') up.settimeout(None) up.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) accept = base64.b64encode(hashlib.sha1((key + WS_GUID).encode()).digest()).decode() self.send_response(101, 'Switching Protocols') self.send_header('Upgrade', 'websocket') self.send_header('Connection', 'Upgrade') self.send_header('Sec-WebSocket-Accept', accept) proto = self.headers.get('Sec-WebSocket-Protocol', '') if 'binary' in [p.strip() for p in proto.split(',')]: self.send_header('Sec-WebSocket-Protocol', 'binary') self.end_headers() self.wfile.flush() self.pipe(up) def pipe(self, up): sock = self.connection sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) send_lock = threading.Lock() alive = [True] closed = [False] # send at most one close frame def frame(op, payload): if op == 0x8: if closed[0]: return closed[0] = True n = len(payload) head = bytes([0x80 | op]) + (bytes([n]) if n < 126 else struct.pack('!BH', 126, n) if n < 65536 else struct.pack('!BQ', 127, n)) with send_lock: sock.sendall(head + payload) def from_browser(): # client frames are masked; data frames (and continuations) go to Postgres try: while alive[0]: h = read_exact(self.rfile, 2) op, n = h[0] & 0x0F, h[1] & 0x7F if n == 126: n = struct.unpack('!H', read_exact(self.rfile, 2))[0] elif n == 127: n = struct.unpack('!Q', read_exact(self.rfile, 8))[0] mask = read_exact(self.rfile, 4) if h[1] & 0x80 else b'\0\0\0\0' raw = read_exact(self.rfile, n) data = bytes(b ^ mask[i % 4] for i, b in enumerate(raw)) if h[1] & 0x80 else raw if op == 0x8: try: frame(0x8, data[:2]) except Exception: pass break if op == 0x9: frame(0xA, data) continue if op in (0x0, 0x1, 0x2): up.sendall(data) except Exception: pass alive[0] = False try: up.shutdown(socket.SHUT_RDWR) except Exception: pass t = threading.Thread(target=from_browser, daemon=True) t.start() try: while alive[0]: chunk = up.recv(65536) if not chunk: break frame(0x2, chunk) except Exception: pass alive[0] = False try: frame(0x8, struct.pack('!H', 1000)) except Exception: pass try: up.close() except Exception: pass self.close_connection = True def main(): p = argparse.ArgumentParser(description='WebSocket bridge to one Postgres server, for the geekychris.github.io Postgres playground.') p.add_argument('--port', type=int, default=6432, help='WebSocket port (default 6432)') p.add_argument('--host', default='127.0.0.1', help='address to listen on (default 127.0.0.1 — keep it local)') p.add_argument('--target', default='localhost:5432', help='Postgres host:port (default localhost:5432)') p.add_argument('--allow-origin', action='append', default=[], help='another web origin allowed to connect') p.add_argument('--quiet', action='store_true') a = p.parse_args() ALLOWED.update(a.allow_origin) srv = ThreadingHTTPServer((a.host, a.port), Relay) srv.daemon_threads = True srv.target, srv.quiet = a.target, a.quiet print(f'postgres-relay {VERSION}: ws://{a.host}:{a.port}/ -> {a.target} (Ctrl+C to stop)', flush=True) try: srv.serve_forever() except KeyboardInterrupt: pass if __name__ == '__main__': main()