Replace a server's accepted short IDs at runtime and revoke sessions

server.go

@@ -9,6 +9,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"slices"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -26,12 +27,19 @@ var ErrServerClosed = errors.New("server closed")
 type Server struct {
 	config       ServerConfig
 	serverPubKey []byte // derived from config.PrivateKey
-	listenerMu   sync.Mutex
-	listener     net.Listener
-	masquerade   *Masquerade
-	ctx          context.Context
-	cancel       context.CancelCauseFunc
-	wg           sync.WaitGroup
+	// shortIDs is the set a client may authenticate with. It starts as
+	// config.ShortIDs; SetShortIDs replaces it while the server runs.
+	shortIDs   atomic.Pointer[[][8]byte]
+	sessionsMu sync.Mutex
+	// sessions maps each authenticated connection to the short ID it
+	// authenticated with, so SetShortIDs can close the ones it revokes.
+	sessions   map[net.Conn][8]byte
+	listenerMu sync.Mutex
+	listener   net.Listener
+	masquerade *Masquerade
+	ctx        context.Context
+	cancel     context.CancelCauseFunc
+	wg         sync.WaitGroup
 }
 
 // NewServer creates a new Samizdat server.
@@ -59,9 +67,12 @@ func NewServer(config ServerConfig) (*Server, error) {
 	s := &Server{
 		config:       config,
 		serverPubKey: serverPubKey,
+		sessions:     make(map[net.Conn][8]byte),
 		ctx:          ctx,
 		cancel:       cancel,
 	}
+	ids := append([][8]byte(nil), config.ShortIDs...)
+	s.shortIDs.Store(&ids)
 
 	// Set up masquerade if configured
 	if config.MasqueradeDomain != "" {
@@ -140,6 +151,59 @@ func (s *Server) Addr() net.Addr {
 	return nil
 }
 
+// SetShortIDs replaces the short IDs clients may authenticate with, while the
+// server runs. A connection that authenticated with an ID the new set leaves
+// out is closed, so removing an ID ends its sessions now rather than whenever
+// they would end on their own. Stream handlers see this as an ordinary
+// disconnect: stream errors and context cancellation.
+//
+// An empty set is rejected with an error and the previous set stays in force,
+// so callers must check the error. ids is copied.
+func (s *Server) SetShortIDs(ids [][8]byte) error {
+	if len(ids) == 0 {
+		return fmt.Errorf("at least one ShortID is required")
+	}
+	next := append([][8]byte(nil), ids...)
+	s.sessionsMu.Lock()
+	s.shortIDs.Store(&next)
+	var revoked []net.Conn
+	for conn, id := range s.sessions {
+		if !slices.Contains(next, id) {
+			revoked = append(revoked, conn)
+			delete(s.sessions, conn)
+		}
+	}
+	s.sessionsMu.Unlock()
+	for _, conn := range revoked {
+		conn.Close()
+	}
+	return nil
+}
+
+// ShortIDs returns a copy of the short IDs clients may authenticate with.
+func (s *Server) ShortIDs() [][8]byte {
+	return append([][8]byte(nil), *s.shortIDs.Load()...)
+}
+
+// trackSession records conn as authenticated with id, unless id has been
+// revoked since it was verified. The check and the record share
+// SetShortIDs' lock, so no revocation falls between them.
+func (s *Server) trackSession(conn net.Conn, id [8]byte) bool {
+	s.sessionsMu.Lock()
+	defer s.sessionsMu.Unlock()
+	if !slices.Contains(*s.shortIDs.Load(), id) {
+		return false
+	}
+	s.sessions[conn] = id
+	return true
+}
+
+func (s *Server) untrackSession(conn net.Conn) {
+	s.sessionsMu.Lock()
+	delete(s.sessions, conn)
+	s.sessionsMu.Unlock()
+}
+
 // handleConnection processes a new TCP connection:
 // 1. Read the ClientHello (buffer raw bytes)
 // 2. Attempt Samizdat auth verification
@@ -168,15 +232,15 @@ func (s *Server) handleConnection(conn net.Conn) {
 	}
 
 	// Verify Samizdat auth using PSK derived from server public key
-	_, authenticated, err := VerifySessionID(sessionID, s.serverPubKey, s.config.ShortIDs)
+	shortID, authenticated, err := VerifySessionID(sessionID, s.serverPubKey, *s.shortIDs.Load())
 	if err != nil || !authenticated {
 		// Auth failed — enter masquerade mode
 		s.doMasquerade(conn, clientHelloRecord)
 		return
 	}
 
 	// Auth passed — complete TLS handshake and enter H2 proxy mode
-	s.handleAuthenticated(conn, clientHelloRecord)
+	s.handleAuthenticated(conn, clientHelloRecord, shortID)
 }
 
 // doMasquerade forwards the connection to the real masquerade domain.
@@ -190,7 +254,14 @@ func (s *Server) doMasquerade(conn net.Conn, clientHelloRecord []byte) {
 
 // handleAuthenticated completes the TLS handshake with the authenticated
 // client and serves HTTP/2 CONNECT requests.
-func (s *Server) handleAuthenticated(conn net.Conn, clientHelloRecord []byte) {
+func (s *Server) handleAuthenticated(conn net.Conn, clientHelloRecord []byte, shortID [8]byte) {
+	// Registered before the handshake: a SetShortIDs that revokes shortID
+	// from here on closes conn, and one that already did refuses it.
+	if !s.trackSession(conn, shortID) {
+		return
+	}
+	defer s.untrackSession(conn)
+
 	// Create a connection that replays the ClientHello record first,
 	// then reads from the real connection
 	replayConn := newReplayConn(conn, clientHelloRecord)

shortids_test.go

@@ -0,0 +1,213 @@
+package samizdat
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"sync"
+	"testing"
+	"time"
+)
+
+// startEcho starts a TCP echo server and returns its address.
+func startEcho(t *testing.T) string {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("echo listener: %v", err)
+	}
+	t.Cleanup(func() { ln.Close() })
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func() {
+				defer c.Close()
+				io.Copy(c, c)
+			}()
+		}
+	}()
+	return ln.Addr().String()
+}
+
+// relayHandler dials the CONNECT destination and copies both ways.
+func relayHandler(_ context.Context, conn net.Conn, destination string) {
+	defer conn.Close()
+	target, err := net.DialTimeout("tcp", destination, 5*time.Second)
+	if err != nil {
+		return
+	}
+	defer target.Close()
+	var wg sync.WaitGroup
+	wg.Add(2)
+	// Closing both ends when either direction finishes unblocks the other copy.
+	go func() { defer wg.Done(); io.Copy(target, conn); target.Close(); conn.Close() }()
+	go func() { defer wg.Done(); io.Copy(conn, target); target.Close(); conn.Close() }()
+	wg.Wait()
+}
+
+// dialEcho opens a tunnel to echoAddr authenticating as shortID, and proves it
+// with one round trip. It returns the open tunnel.
+func dialEcho(t *testing.T, serverAddr string, pub []byte, shortID [8]byte, echoAddr string) (net.Conn, error) {
+	t.Helper()
+	client, err := NewClient(ClientConfig{
+		ServerAddr:  serverAddr,
+		ServerName:  "test.example.com",
+		PublicKey:   pub,
+		ShortID:     shortID,
+		Fingerprint: "chrome",
+	})
+	if err != nil {
+		t.Fatalf("NewClient: %v", err)
+	}
+	t.Cleanup(func() { client.Close() })
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	conn, err := client.DialContext(ctx, "tcp", echoAddr)
+	if err != nil {
+		return nil, err
+	}
+	msg := []byte("ping")
+	if _, err := conn.Write(msg); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	buf := make([]byte, len(msg))
+	conn.SetReadDeadline(time.Now().Add(5 * time.Second))
+	if _, err := io.ReadFull(conn, buf); err != nil {
+		conn.Close()
+		return nil, err
+	}
+	conn.SetReadDeadline(time.Time{})
+	return conn, nil
+}
+
+func TestSetShortIDsAddsAndRevokes(t *testing.T) {
+	priv, pub, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateKeyPair: %v", err)
+	}
+	base, _ := GenerateShortID()
+	paired, _ := GenerateShortID()
+	certPEM, keyPEM := generateSelfSignedCert(t)
+	echo := startEcho(t)
+	server, ln := startTestServer(t, ServerConfig{
+		ListenAddr: "127.0.0.1:0",
+		PrivateKey: priv,
+		ShortIDs:   [][8]byte{base},
+		CertPEM:    certPEM,
+		KeyPEM:     keyPEM,
+		Handler:    relayHandler,
+	})
+	addr := ln.Addr().String()
+
+	if c, err := dialEcho(t, addr, pub, paired, echo); err == nil {
+		c.Close()
+		t.Fatal("a short ID the server was never given authenticated")
+	}
+
+	if err := server.SetShortIDs([][8]byte{base, paired}); err != nil {
+		t.Fatalf("SetShortIDs: %v", err)
+	}
+	held, err := dialEcho(t, addr, pub, paired, echo)
+	if err != nil {
+		t.Fatalf("an added short ID was rejected: %v", err)
+	}
+	defer held.Close()
+
+	if err := server.SetShortIDs([][8]byte{base}); err != nil {
+		t.Fatalf("SetShortIDs: %v", err)
+	}
+	if c, err := dialEcho(t, addr, pub, paired, echo); err == nil {
+		c.Close()
+		t.Fatal("a removed short ID still authenticates")
+	}
+	c, err := dialEcho(t, addr, pub, base, echo)
+	if err != nil {
+		t.Fatalf("the remaining short ID was rejected: %v", err)
+	}
+	c.Close()
+
+	// The session that authenticated with the removed ID was closed, not left
+	// to run until it ends on its own.
+	held.SetReadDeadline(time.Now().Add(5 * time.Second))
+	_, err = held.Read(make([]byte, 1))
+	var ne net.Error
+	if err == nil || (errors.As(err, &ne) && ne.Timeout()) {
+		t.Fatalf("a session using a removed short ID survived: %v", err)
+	}
+}
+
+func TestSetShortIDsRejectsEmptySetAndCopies(t *testing.T) {
+	priv, _, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateKeyPair: %v", err)
+	}
+	a, _ := GenerateShortID()
+	b, _ := GenerateShortID()
+	certPEM, keyPEM := generateSelfSignedCert(t)
+	s, err := NewServer(ServerConfig{PrivateKey: priv, ShortIDs: [][8]byte{a}, CertPEM: certPEM, KeyPEM: keyPEM, Handler: relayHandler})
+	if err != nil {
+		t.Fatalf("NewServer: %v", err)
+	}
+	if err := s.SetShortIDs(nil); err == nil {
+		t.Fatal("an empty set must be refused")
+	}
+	if got := s.ShortIDs(); len(got) != 1 || got[0] != a {
+		t.Fatalf("a refused update changed the set: %v", got)
+	}
+	ids := [][8]byte{b}
+	if err := s.SetShortIDs(ids); err != nil {
+		t.Fatalf("SetShortIDs: %v", err)
+	}
+	ids[0] = a
+	if got := s.ShortIDs(); got[0] != b {
+		t.Fatal("SetShortIDs must copy its argument")
+	}
+	got := s.ShortIDs()
+	got[0] = a
+	if s.ShortIDs()[0] != b {
+		t.Fatal("ShortIDs must return a copy")
+	}
+}
+
+func TestTrackSessionRefusesIDRevokedAfterVerify(t *testing.T) {
+	priv, pub, err := GenerateKeyPair()
+	if err != nil {
+		t.Fatalf("GenerateKeyPair: %v", err)
+	}
+	base, _ := GenerateShortID()
+	paired, _ := GenerateShortID()
+	certPEM, keyPEM := generateSelfSignedCert(t)
+	s, err := NewServer(ServerConfig{PrivateKey: priv, ShortIDs: [][8]byte{base, paired}, CertPEM: certPEM, KeyPEM: keyPEM, Handler: relayHandler})
+	if err != nil {
+		t.Fatalf("NewServer: %v", err)
+	}
+
+	sessionID, err := BuildSessionID(pub, paired)
+	if err != nil {
+		t.Fatalf("BuildSessionID: %v", err)
+	}
+	id, ok, err := VerifySessionID(sessionID[:], s.serverPubKey, s.ShortIDs())
+	if err != nil || !ok || id != paired {
+		t.Fatalf("verify against the old set: id=%v ok=%v err=%v", id, ok, err)
+	}
+
+	if err := s.SetShortIDs([][8]byte{base}); err != nil {
+		t.Fatalf("SetShortIDs: %v", err)
+	}
+
+	c1, c2 := net.Pipe()
+	defer c1.Close()
+	defer c2.Close()
+	if s.trackSession(c1, id) {
+		t.Fatal("a session verified before its ID was revoked was tracked")
+	}
+	if !s.trackSession(c1, base) {
+		t.Fatal("a session using a remaining ID was refused")
+	}
+	s.untrackSession(c1)
+}