tunnel/internal/server/bridge.go

package server

import (
	"context"
	"encoding/json"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coder/websocket"
)

// Hub tracks the attached bridges: which tunnel they belong to and which
// routes each claimed. A route belongs to one bridge at a time.
type Hub struct {
	store *Store

	mu      sync.Mutex
	bridges map[*Bridge]struct{}
	byRoute map[string]map[string]*Bridge // tunnel ID -> route -> bridge
}

func NewHub(store *Store) *Hub {
	return &Hub{
		store:   store,
		bridges: map[*Bridge]struct{}{},
		byRoute: map[string]map[string]*Bridge{},
	}
}

// online reports whether the tunnel has an attached bridge.
func (h *Hub) online(tunnelID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	for _, b := range h.byRoute[tunnelID] {
		_ = b
		return true
	}
	return false
}

// bridgeFor returns the bridge that claimed route on the tunnel, or nil.
func (h *Hub) bridgeFor(tunnelID, route string) *Bridge {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.byRoute[tunnelID][route]
}

// closeTunnel closes every bridge of the tunnel (it was deleted).
func (h *Hub) closeTunnel(tunnelID string) {
	h.mu.Lock()
	var bridges []*Bridge
	for b := range h.bridges {
		if b.tunnelID == tunnelID {
			bridges = append(bridges, b)
		}
	}
	h.mu.Unlock()
	for _, b := range bridges {
		b.close(websocket.StatusGoingAway, "tunnel deleted")
	}
}

// drainAll announces drain to every bridge and closes it, so clients
// reconnect elsewhere (used on shutdown).
func (h *Hub) DrainAll() {
	h.mu.Lock()
	bridges := make([]*Bridge, 0, len(h.bridges))
	for b := range h.bridges {
		bridges = append(bridges, b)
	}
	h.mu.Unlock()
	for _, b := range bridges {
		drain, _ := json.Marshal(drainMessage{Type: "drain", Reason: "shutdown"})
		b.trySend(websocket.MessageText, drain)
		b.close(websocket.StatusServiceRestart, "shutdown")
	}
}

// BridgeInfo is a snapshot of one attached bridge for the admin API.
type BridgeInfo struct {
	Session       string
	TunnelID      string
	ClientVersion string
	Routes        []string
	Visitors      int
	MaxConns      int
	ConnectedAt   time.Time
}

// BridgeInfos snapshots every attached bridge. Field reads are safe without
// the bridge lock because they are immutable after construction; only the
// visitor count needs the lock.
func (h *Hub) BridgeInfos() []BridgeInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	out := make([]BridgeInfo, 0, len(h.bridges))
	for b := range h.bridges {
		b.mu.Lock()
		visitors := len(b.conns)
		b.mu.Unlock()
		out = append(out, BridgeInfo{
			Session:       b.session,
			TunnelID:      b.tunnelID,
			ClientVersion: b.clientVersion,
			Routes:        b.routes,
			Visitors:      visitors,
			MaxConns:      b.maxConns,
			ConnectedAt:   b.startedAt,
		})
	}
	return out
}

func (h *Hub) register(b *Bridge) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	claimed := h.byRoute[b.tunnelID]
	if claimed == nil {
		claimed = map[string]*Bridge{}
		h.byRoute[b.tunnelID] = claimed
	}
	for _, route := range b.routes {
		if _, taken := claimed[route]; taken {
			return false
		}
	}
	for _, route := range b.routes {
		claimed[route] = b
	}
	h.bridges[b] = struct{}{}
	return true
}

func (h *Hub) unregister(b *Bridge) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.bridges, b)
	if claimed := h.byRoute[b.tunnelID]; claimed != nil {
		for route, owner := range claimed {
			if owner == b {
				delete(claimed, route)
			}
		}
		if len(claimed) == 0 {
			delete(h.byRoute, b.tunnelID)
		}
	}
}

// Bridge is one attached client session: the WebSocket, its routes, and its
// visitor connections.
type Bridge struct {
	hub      *Hub
	tunnelID string
	session  string
	routes   []string
	maxConns int
	ws       *websocket.Conn

	clientVersion string
	startedAt     time.Time

	outbox    chan outbound
	done      chan struct{}
	closeOnce sync.Once

	mu       sync.Mutex
	conns    map[uint32]*visitorConn
	nextConn uint32

	lastActivity atomic.Int64
}

type outbound struct {
	typ  websocket.MessageType
	data []byte
}

// attach upgrades the HTTP connection and runs the bridge session. It blocks
// until the session ends.
func (h *Hub) attach(w http.ResponseWriter, r *http.Request, tunnel storedTunnel) {
	ws, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		Subprotocols:    []string{BridgeSubprotocol},
		CompressionMode: websocket.CompressionDisabled,
		OriginPatterns:  nil, // CLI and SDK send no Origin header
	})
	if err != nil {
		return // Accept already replied
	}
	ws.SetReadLimit(clientHelloLimit) // 64 KiB before attach

	ctx := context.Background()
	fail := func(code string) {
		msg, _ := json.Marshal(attachErrorMessage{Type: "attach_error", Code: code})
		writeWS(ctx, ws, websocket.MessageText, msg)
		ws.Close(websocket.StatusPolicyViolation, code)
	}

	// The attach must arrive within 10 seconds.
	readCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	typ, raw, err := ws.Read(readCtx)
	cancel()
	if err != nil {
		ws.Close(websocket.StatusPolicyViolation, "attach timeout")
		return
	}
	var attach attachMessage
	if typ != websocket.MessageText || json.Unmarshal(raw, &attach) != nil || attach.Type != "attach" {
		fail("bad_attach")
		return
	}
	if !tokenMatches(attach.Token, tunnel.TokenHash) {
		fail("bad_token")
		return
	}
	if !certReady(tunnel.Certificate) {
		fail("cert_not_ready")
		return
	}
	if len(attach.Routes) > maxRoutesPerAttach {
		fail("too_many_routes")
		return
	}
	routes := dedupe(attach.Routes)
	for _, route := range routes {
		if !validRoute(route) {
			fail("invalid_route")
			return
		}
	}

	b := &Bridge{
		hub:           h,
		tunnelID:      tunnel.ID,
		session:       sessionID(),
		routes:        routes,
		maxConns:      maxConnsPerBridge,
		ws:            ws,
		clientVersion: attach.Client.Version,
		startedAt:     time.Now(),
		outbox:        make(chan outbound, 256),
		done:          make(chan struct{}),
		conns:         map[uint32]*visitorConn{},
	}
	if attach.Client.MaxConns > 0 && attach.Client.MaxConns < b.maxConns {
		b.maxConns = attach.Client.MaxConns
	}
	b.touch()
	if !h.register(b) {
		fail("route_conflict")
		return
	}

	attached, _ := json.Marshal(attachedMessage{
		Type:          "attached",
		Session:       b.session,
		Routes:        routes,
		HeartbeatMs:   heartbeatMs,
		IdleTimeoutMs: idleTimeoutMs,
	})
	if err := writeWS(ctx, ws, websocket.MessageText, attached); err != nil {
		h.unregister(b)
		ws.Close(websocket.StatusInternalError, "write failed")
		return
	}
	ws.SetReadLimit(4 + maxPayload) // a data frame with the largest payload

	go b.writeLoop()
	go b.sweeper()
	b.readLoop() // blocks until the WebSocket ends
	b.close(websocket.StatusGoingAway, "session ended")
}

// openVisitor hands a new public connection to the bridge: an open control
// message, the ClientHello as the first data, and a pump for what follows.
func (b *Bridge) openVisitor(conn net.Conn, hello []byte, peer, sni, alpn string) {
	b.mu.Lock()
	if len(b.conns) >= b.maxConns {
		b.mu.Unlock()
		conn.Close()
		return
	}
	id := b.nextConn
	b.nextConn++
	v := &visitorConn{id: id, conn: conn}
	b.conns[id] = v
	b.mu.Unlock()

	open, _ := json.Marshal(openMessage{
		Type: "open", Conn: id, Peer: peer, SNI: sni, ALPN: alpn,
	})
	if !b.send(websocket.MessageText, open) {
		b.removeConn(id, true)
		return
	}
	for len(hello) > 0 {
		n := min(len(hello), maxPayload)
		if !b.send(websocket.MessageBinary, buildFrame(id, hello[:n])) {
			b.removeConn(id, true)
			return
		}
		hello = hello[n:]
	}
	go b.pump(v)
}

// pump copies visitor bytes to the bridge and half-closes on EOF.
func (b *Bridge) pump(v *visitorConn) {
	buf := make([]byte, maxPayload)
	for {
		n, err := v.conn.Read(buf)
		v.touch()
		if n > 0 {
			if !b.send(websocket.MessageBinary, buildFrame(v.id, buf[:n])) {
				return
			}
		}
		if err != nil {
			b.mu.Lock()
			v.visitorEnded = true
			bothEnded := v.clientEnded
			stillKnown := b.conns[v.id] == v
			b.mu.Unlock()
			if !stillKnown {
				return // reset or removed already
			}
			if err == io.EOF {
				end, _ := json.Marshal(endMessage{Type: "end", Conn: v.id})
				b.send(websocket.MessageText, end)
			} else {
				reset, _ := json.Marshal(resetMessage{Type: "reset", Conn: v.id, Code: ResetIOError})
				b.send(websocket.MessageText, reset)
				b.removeConn(v.id, true)
				return
			}
			if bothEnded {
				b.removeConn(v.id, false)
			}
			return
		}
	}
}

// readLoop applies client control messages and data frames until the
// WebSocket fails.
func (b *Bridge) readLoop() {
	ctx := context.Background()
	for {
		typ, raw, err := b.ws.Read(ctx)
		if err != nil {
			log.Printf("bridge %s: read: %v", b.session, err)
			return
		}
		b.touch()
		if typ == websocket.MessageText {
			b.handleControl(raw)
			continue
		}
		connID, payload, ok := parseFrame(raw)
		if !ok {
			continue
		}
		b.mu.Lock()
		v := b.conns[connID]
		b.mu.Unlock()
		if v == nil {
			continue
		}
		v.touch()
		if _, err := v.conn.Write(payload); err != nil {
			b.removeConn(connID, true)
		}
	}
}

func (b *Bridge) handleControl(raw []byte) {
	msg, err := parseClientControl(raw)
	if err != nil {
		return
	}
	switch msg.Type {
	case "ping":
		pong, _ := json.Marshal(pingPongMessage{Type: "pong", TimeSent: msg.TimeSent})
		b.send(websocket.MessageText, pong)
	case "end":
		if msg.Conn != nil {
			b.endFromClient(*msg.Conn)
		}
	case "reset":
		if msg.Conn != nil {
			b.removeConn(*msg.Conn, true)
		}
	}
	// Unknown types and fields are ignored.
}

func (b *Bridge) endFromClient(connID uint32) {
	b.mu.Lock()
	v := b.conns[connID]
	b.mu.Unlock()
	if v == nil {
		return
	}
	b.mu.Lock()
	v.clientEnded = true
	bothEnded := v.visitorEnded
	b.mu.Unlock()
	if tcp, ok := v.conn.(*net.TCPConn); ok {
		tcp.CloseWrite()
	}
	if bothEnded {
		b.removeConn(connID, false)
	}
}

// writeLoop is the only writer on the WebSocket, keeping frames ordered.
func (b *Bridge) writeLoop() {
	ctx := context.Background()
	for {
		select {
		case <-b.done:
			return
		case msg := <-b.outbox:
			if err := writeWS(ctx, b.ws, msg.typ, msg.data); err != nil {
				b.close(websocket.StatusInternalError, "write failed")
				return
			}
		}
	}
}

// sweeper retires a silent bridge (idle_timeout + heartbeat) and resets
// visitor connections with no traffic for 15 minutes.
func (b *Bridge) sweeper() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-b.done:
			return
		case now := <-ticker.C:
			if now.UnixMilli()-b.lastActivity.Load() > (idleTimeoutMs + heartbeatMs) {
				b.close(websocket.StatusGoingAway, "idle timeout")
				return
			}
			b.mu.Lock()
			var stale []uint32
			for id, v := range b.conns {
				if now.UnixMilli()-v.lastActivity.Load() > 15*60*1000 {
					stale = append(stale, id)
				}
			}
			b.mu.Unlock()
			for _, id := range stale {
				reset, _ := json.Marshal(resetMessage{Type: "reset", Conn: id, Code: ResetIdleTimeout})
				b.send(websocket.MessageText, reset)
				b.removeConn(id, true)
			}
		}
	}
}

// send queues a frame, applying backpressure when the client is slow.
// It returns false when the bridge is gone.
func (b *Bridge) send(typ websocket.MessageType, data []byte) bool {
	select {
	case <-b.done:
		return false
	case b.outbox <- outbound{typ, data}:
		return true
	}
}

// trySend queues a frame only if it fits immediately.
func (b *Bridge) trySend(typ websocket.MessageType, data []byte) {
	select {
	case b.outbox <- outbound{typ, data}:
	default:
	}
}

// removeConn closes the visitor socket (abortively when rst) and forgets it.
func (b *Bridge) removeConn(id uint32, rst bool) {
	b.mu.Lock()
	v := b.conns[id]
	delete(b.conns, id)
	b.mu.Unlock()
	if v == nil {
		return
	}
	if rst {
		if tcp, ok := v.conn.(*net.TCPConn); ok {
			tcp.SetLinger(0)
		}
	}
	v.conn.Close()
}

func (b *Bridge) touch() {
	b.lastActivity.Store(time.Now().UnixMilli())
}

// close ends the session exactly once: conns reset, routes released.
func (b *Bridge) close(status websocket.StatusCode, reason string) {
	b.closeOnce.Do(func() {
		close(b.done)
		b.hub.unregister(b)
		b.mu.Lock()
		ids := make([]uint32, 0, len(b.conns))
		for id := range b.conns {
			ids = append(ids, id)
		}
		b.mu.Unlock()
		for _, id := range ids {
			b.removeConn(id, true)
		}
		b.ws.Close(status, reason)
	})
}

// visitorConn is one public TCP connection carried over the bridge.
type visitorConn struct {
	id   uint32
	conn net.Conn

	visitorEnded bool // pump saw EOF: no more visitor data
	clientEnded  bool // client sent end: target will send no more

	lastActivity atomic.Int64
}

func (v *visitorConn) touch() {
	v.lastActivity.Store(time.Now().UnixMilli())
}

func writeWS(ctx context.Context, ws *websocket.Conn, typ websocket.MessageType, data []byte) error {
	writeCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	return ws.Write(writeCtx, typ, data)
}

func dedupe(routes []string) []string {
	seen := map[string]struct{}{}
	out := make([]string, 0, len(routes))
	for _, route := range routes {
		if _, ok := seen[route]; !ok {
			seen[route] = struct{}{}
			out = append(out, route)
		}
	}
	return out
}

// isWebSocketUpgrade reports the shape of the upgrade request, for the 426
// responses the contract pins.
func upgradeRequested(r *http.Request) bool {
	return strings.EqualFold(r.Header.Get("Upgrade"), "websocket")
}

func subprotocolOffered(r *http.Request) bool {
	for _, offered := range strings.Split(r.Header.Get("Sec-Websocket-Protocol"), ",") {
		if strings.TrimSpace(offered) == BridgeSubprotocol {
			return true
		}
	}
	return false
}