From 0bd392e076e36c80a152abe00cbcd0bc9efedd9c Mon Sep 17 00:00:00 2001
From: bndw <ben@bdw.to>
Date: Mon, 9 Mar 2026 12:48:39 -0700
Subject: feat: add axon CLI for pub/sub; fix relay hijack context bug

- cmd/axon: new CLI module with keygen, pub, and req subcommands
  - keygen: generate Ed25519 keypair, print hex seed and pubkey
  - pub: sign and publish an event; accepts --kind, --content, --tag
  - req: query/stream events as JSON lines; accepts --kind, --author,
    --tag, --since, --until, --limit, --stream
  - key loaded from --key flag or AXON_KEY env var

- relay/websocket: add Dial() for client-side WebSocket handshake
  (ws:// and wss://, RFC 6455 masking via client:true)

- relay/server: fix broken-pipe on auth by switching hijacked conn
  goroutine from r.Context() to context.Background(); r.Context() is
  cancelled by net/http immediately after Hijack is called
---
 relay/server.go              |  5 ++-
 relay/websocket/websocket.go | 75 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

(limited to 'relay')

diff --git a/relay/server.go b/relay/server.go
index 085929c..d4a1edd 100644
--- a/relay/server.go
+++ b/relay/server.go
@@ -98,7 +98,10 @@ func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
 	s.mu.Add(1)
 	go func() {
 		defer s.mu.Done()
-		ctx := r.Context()
+		// r.Context() is cancelled by the HTTP server when Hijack is called,
+		// so we use a fresh context. The connection manages its own lifecycle
+		// via the ping loop and WebSocket close frames.
+		ctx := context.Background()
 		h.serve(ctx)
 		if err := c.CloseConn(); err != nil {
 			// Ignore close errors — connection may already be gone.
diff --git a/relay/websocket/websocket.go b/relay/websocket/websocket.go
index cfc3289..2ae0dec 100644
--- a/relay/websocket/websocket.go
+++ b/relay/websocket/websocket.go
@@ -7,12 +7,14 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/sha1"
+	"crypto/tls"
 	"encoding/base64"
 	"encoding/binary"
 	"fmt"
 	"io"
 	"net"
 	"net/http"
+	"net/url"
 	"strings"
 	"sync"
 	"time"
@@ -207,6 +209,79 @@ func acceptKey(key string) string {
 	return base64.StdEncoding.EncodeToString(h.Sum(nil))
 }
 
+// Dial connects to a WebSocket server at rawURL and performs the client-side
+// RFC 6455 handshake. Supports ws:// and wss:// schemes.
+func Dial(rawURL string) (*Conn, error) {
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return nil, fmt.Errorf("websocket: parse url: %w", err)
+	}
+
+	host := u.Host
+	var netConn net.Conn
+	switch u.Scheme {
+	case "ws":
+		if !strings.Contains(host, ":") {
+			host += ":80"
+		}
+		netConn, err = net.Dial("tcp", host)
+	case "wss":
+		if !strings.Contains(host, ":") {
+			host += ":443"
+		}
+		netConn, err = tls.Dial("tcp", host, &tls.Config{ServerName: u.Hostname()})
+	default:
+		return nil, fmt.Errorf("websocket: unsupported scheme %q (use ws:// or wss://)", u.Scheme)
+	}
+	if err != nil {
+		return nil, fmt.Errorf("websocket: dial %s: %w", host, err)
+	}
+
+	// Generate a random 16-byte key and base64-encode it.
+	var keyBytes [16]byte
+	if _, err := rand.Read(keyBytes[:]); err != nil {
+		netConn.Close()
+		return nil, fmt.Errorf("websocket: generate key: %w", err)
+	}
+	key := base64.StdEncoding.EncodeToString(keyBytes[:])
+
+	path := u.RequestURI()
+	if path == "" {
+		path = "/"
+	}
+
+	req := "GET " + path + " HTTP/1.1\r\n" +
+		"Host: " + u.Host + "\r\n" +
+		"Upgrade: websocket\r\n" +
+		"Connection: Upgrade\r\n" +
+		"Sec-WebSocket-Key: " + key + "\r\n" +
+		"Sec-WebSocket-Version: 13\r\n\r\n"
+
+	if _, err := netConn.Write([]byte(req)); err != nil {
+		netConn.Close()
+		return nil, fmt.Errorf("websocket: send handshake: %w", err)
+	}
+
+	br := bufio.NewReader(netConn)
+	resp, err := http.ReadResponse(br, nil)
+	if err != nil {
+		netConn.Close()
+		return nil, fmt.Errorf("websocket: read handshake response: %w", err)
+	}
+	resp.Body.Close()
+
+	if resp.StatusCode != 101 {
+		netConn.Close()
+		return nil, fmt.Errorf("websocket: server returned status %d, want 101", resp.StatusCode)
+	}
+	if resp.Header.Get("Sec-WebSocket-Accept") != acceptKey(key) {
+		netConn.Close()
+		return nil, fmt.Errorf("websocket: bad Sec-WebSocket-Accept header")
+	}
+
+	return &Conn{rwc: netConn, br: br, client: true}, nil
+}
+
 // Accept performs the server-side WebSocket handshake, hijacking the HTTP
 // connection and returning a Conn ready for framed I/O.
 func Accept(w http.ResponseWriter, r *http.Request) (*Conn, error) {
-- 
cgit v1.2.3