diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..1075d2d --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,91 @@ +# Create a version tag and GitHub Release for the proxy Go module. +# Run from Actions → Release → Run workflow (typically on develop or master). +name: Release + +on: + workflow_dispatch: + inputs: + version: + description: "Version (e.g. v0.1.0 or 0.1.0)" + required: true + type: string + prerelease: + description: "Mark GitHub Release as prerelease" + required: false + type: boolean + default: false + +permissions: + contents: write + +jobs: + release: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + # Need write access to push the tag back to the repo. + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Normalize version + id: ver + run: | + set -euo pipefail + RAW="${{ inputs.version }}" + RAW="${RAW#v}" + if [[ ! "$RAW" =~ ^[0-9]+\.[0-9]+\.[0-9]+([.-][0-9A-Za-z.-]+)?$ ]]; then + echo "Invalid version: ${{ inputs.version }} (expected semver like v0.1.0)" + exit 1 + fi + TAG="v${RAW}" + echo "tag=${TAG}" >> "$GITHUB_OUTPUT" + echo "Using tag ${TAG}" + + - name: Ensure tag does not already exist + run: | + set -euo pipefail + TAG="${{ steps.ver.outputs.tag }}" + if git rev-parse "$TAG" >/dev/null 2>&1; then + echo "Tag ${TAG} already exists" + exit 1 + fi + if git ls-remote --tags origin "refs/tags/${TAG}" | grep -q .; then + echo "Tag ${TAG} already exists on origin" + exit 1 + fi + + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version-file: go.mod + cache: true + + - name: Test + run: go test ./... -count=1 + + - name: Create and push tag + run: | + set -euo pipefail + TAG="${{ steps.ver.outputs.tag }}" + git config user.name "github-actions[bot]" + git config user.email "41898282+github-actions[bot]@users.noreply.github.com" + git tag -a "$TAG" -m "Release ${TAG}" + git push origin "$TAG" + + - name: Create GitHub Release + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -euo pipefail + TAG="${{ steps.ver.outputs.tag }}" + ARGS=(release create "$TAG" + --title "$TAG" + --generate-notes + --target "${{ github.sha }}" + ) + if [[ "${{ inputs.prerelease }}" == "true" ]]; then + ARGS+=(--prerelease) + fi + gh "${ARGS[@]}" diff --git a/README.md b/README.md index d2c49b4..c16050f 100644 --- a/README.md +++ b/README.md @@ -27,3 +27,11 @@ Name-based egress via HTTP CONNECT over the mesh (underlay remains WireGuard). S ```bash go test ./... -race ``` + +## Release + +From GitHub Actions → **Release** → **Run workflow**: enter a semver (`v0.1.0` or `0.1.0`). The workflow runs tests, pushes an annotated tag, and creates a GitHub Release (notes auto-generated). Consumers can then: + +```bash +go get github.com/gravitl/proxy/l7@v0.1.0 +``` diff --git a/docs/PROXY_WSS_UPLINK.md b/docs/PROXY_WSS_UPLINK.md new file mode 100644 index 0000000..1586485 --- /dev/null +++ b/docs/PROXY_WSS_UPLINK.md @@ -0,0 +1,68 @@ +# TCP Uplink WebSocket Transport + +Netmaker TCP uplink uses **WebSocket** (`/uplink/v1`) as the application transport in every deployment. TLS termination is configured independently via `tcp_proxy_tls_mode`. + +## TLS modes + +| Mode | Value | Behavior | +|------|-------|----------| +| Self-signed (default) | `selfsigned` | Uplink terminates TLS with a persisted self-signed certificate and serves `wss://` | +| Reverse proxy | `proxy` | Upstream reverse proxy terminates TLS; uplink serves plain `ws://` on an internal bind | + +Invalid modes fail startup/API with: `unsupported uplink TLS mode: …` + +## Standalone gateway + +```json +{ + "enabled": true, + "listen_port": 443, + "tls_mode": "selfsigned" +} +``` + +Client dials: `wss://:443/uplink/v1` + +## Behind a reverse proxy + +```json +{ + "enabled": true, + "listen_port": 51822, + "listen_addr": "127.0.0.1", + "tls_mode": "proxy", + "public_hostname": "gateway.example.com" +} +``` + +- `listen_port` / `listen_addr` are the **backend** bind (what Caddy/nginx proxies to). +- Set gateway `tcp_proxy_public_hostname` (e.g. `gateway.example.com`) so clients dial the reverse-proxy DNS name. +- Clients are published `wss://:/uplink/v1` in `proxy` mode. +- Public port defaults to **443** (server env `TCP_PROXY_PUBLIC_PORT` can override). +- The reverse proxy must expose that public port and forward `/uplink/v1` to the backend. + +Generic reverse-proxy route (works with Caddy, nginx, HAProxy, Ingress, etc.): + +```text +:443 /uplink/v1 → 127.0.0.1:51822 +``` + +Example Caddy (documentation only — not required by the implementation): + +```caddy +gateway.example.com { + reverse_proxy /uplink/v1 127.0.0.1:51822 +} +``` + +Client dials: `wss://gateway.example.com:443/uplink/v1` (falls back to gateway IP if hostname unset) + +## Client trust + +- Public / reverse-proxy certificates: normal system TLS verification +- Self-signed gateways: SHA-256 leaf fingerprint published as `tcp_proxy_cert_fingerprint` +- Temporary legacy: IP dial without fingerprint may skip verify (deprecated) + +## Origin policy + +Netclient is a machine client and typically sends no browser `Origin` header. The uplink WebSocket upgrader accepts missing Origin; authentication remains the HELLO/WG-proof security boundary. diff --git a/go.mod b/go.mod index 0141656..1d5b4fa 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,5 @@ module github.com/gravitl/proxy go 1.22 + +require github.com/gorilla/websocket v1.5.3 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..25a9fc4 --- /dev/null +++ b/go.sum @@ -0,0 +1,2 @@ +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= diff --git a/uplink/client.go b/uplink/client.go index 49e181f..c0557b8 100644 --- a/uplink/client.go +++ b/uplink/client.go @@ -6,16 +6,24 @@ import ( "encoding/json" "errors" "io" - "net" + "math/rand" + "net/url" + "strings" "sync" "sync/atomic" "time" ) -// ClientOptions configures the TCP/TLS uplink client (spec §9.1). +// ClientOptions configures the WebSocket uplink client. type ClientOptions struct { - Addr string - ServerName string + // Addr is the uplink endpoint. Prefer a full URL: + // wss://gateway.example.com/uplink/v1 + // Legacy host:port is accepted and rewritten to wss://host:port/uplink/v1. + Addr string + // ServerName is TLS SNI (optional; derived from URL host when empty). + ServerName string + // TLSConfig configures TLS verification for wss:// dials. + // Required for wss; ignored for plain ws:// (proxy-mode internal tests). TLSConfig *tls.Config HelloFactory func() (ClientHello, error) PacketHandler func([]byte) error @@ -25,28 +33,36 @@ type ClientOptions struct { WriteTimeout time.Duration ReconnectBackoff BackoffConfig MaxFrameSize int + PingInterval time.Duration + PongWait time.Duration + HandshakeTimeout time.Duration } -// Client maintains a TCP/TLS session to the relay (spec §9.1). +// Client maintains a WebSocket session to the relay/gateway. type Client struct { - opts ClientOptions - addr string - log Logger - metrics MetricsSink + opts ClientOptions + addr string // normalised dial URL + log Logger + metrics MetricsSink maxPayload uint32 backoff BackoffConfig keepAlive time.Duration writeTimeout time.Duration + pingInterval time.Duration + pongWait time.Duration tlsConfig *tls.Config mu sync.Mutex state atomic.Value // ClientState - tlsConn *tls.Conn + conn Conn sessID uint32 cancel context.CancelFunc runWG sync.WaitGroup closed atomic.Bool started atomic.Bool + + // writeMu serialises frames. Never acquire writeMu while holding mu. + writeMu sync.Mutex } // NewClient validates options and returns a Client. @@ -54,8 +70,13 @@ func NewClient(opts ClientOptions) (*Client, error) { if opts.Addr == "" { return nil, errors.New("proxy: Addr is required") } - if opts.TLSConfig == nil { - return nil, errors.New("proxy: TLSConfig is required") + dialURL, err := NormaliseUplinkURL(opts.Addr) + if err != nil { + return nil, err + } + needsTLS := strings.HasPrefix(dialURL, "wss://") + if needsTLS && opts.TLSConfig == nil { + return nil, errors.New("proxy: TLSConfig is required for wss://") } if opts.HelloFactory == nil { return nil, errors.New("proxy: HelloFactory is required") @@ -83,26 +104,74 @@ func NewClient(opts ClientOptions) (*Client, error) { if wt <= 0 { wt = 2 * time.Second } - cfg := opts.TLSConfig.Clone() - if opts.ServerName != "" { - cfg.ServerName = opts.ServerName + pi := opts.PingInterval + if pi <= 0 { + pi = defaultWSPingInterval + } + pw := opts.PongWait + if pw <= 0 { + pw = defaultWSPongWait + } + var cfg *tls.Config + if opts.TLSConfig != nil { + cfg = opts.TLSConfig.Clone() + if opts.ServerName != "" { + cfg.ServerName = opts.ServerName + } else if cfg.ServerName == "" { + if u, err := url.Parse(dialURL); err == nil { + cfg.ServerName = u.Hostname() + } + } } c := &Client{ opts: opts, - addr: opts.Addr, + addr: dialURL, log: log, metrics: metrics, maxPayload: uint32(max), backoff: opts.ReconnectBackoff.withDefaults(), keepAlive: ka, writeTimeout: wt, + pingInterval: pi, + pongWait: pw, tlsConfig: cfg, } c.state.Store(ClientState(StateDisconnected)) return c, nil } -// State returns the current client state (spec §13.1). +// NormaliseUplinkURL accepts wss://…/uplink/v1 or legacy host:port. +func NormaliseUplinkURL(addr string) (string, error) { + addr = strings.TrimSpace(addr) + if addr == "" { + return "", errors.New("proxy: empty uplink address") + } + if strings.Contains(addr, "://") { + u, err := url.Parse(addr) + if err != nil { + return "", err + } + switch u.Scheme { + case "wss", "ws", "https", "http": + default: + return "", errors.New("proxy: unsupported uplink URL scheme") + } + switch u.Scheme { + case "https": + u.Scheme = "wss" + case "http": + u.Scheme = "ws" + } + if u.Path == "" || u.Path == "/" { + u.Path = UplinkWSPath + } + return u.String(), nil + } + // Legacy host:port → wss://host:port/uplink/v1 + return "wss://" + addr + UplinkWSPath, nil +} + +// State returns the current client state. func (c *Client) State() ClientState { v := c.state.Load() if v == nil { @@ -138,8 +207,8 @@ func (c *Client) Stop(ctx context.Context) error { c.mu.Lock() cancel := c.cancel c.cancel = nil - conn := c.tlsConn - c.tlsConn = nil + conn := c.conn + c.conn = nil c.mu.Unlock() if cancel != nil { cancel() @@ -160,7 +229,7 @@ func (c *Client) Stop(ctx context.Context) error { return nil } -// SendPacket sends a DATA frame with the current session ID (spec §9.1). +// SendPacket sends a DATA frame with the current session ID. Safe for concurrent use. func (c *Client) SendPacket(ctx context.Context, pkt []byte) error { _ = ctx if c.closed.Load() { @@ -170,28 +239,43 @@ func (c *Client) SendPacket(ctx context.Context, pkt []byte) error { return ErrInvalidFrame } c.mu.Lock() - conn := c.tlsConn - sid := c.sessID - st := c.State() - if conn == nil || st != StateActive { - c.mu.Unlock() + conn := c.conn + c.mu.Unlock() + if conn == nil || c.State() != StateActive { return ErrClientClosed } - hdr := FrameHeader{ - Version: ProtocolVersion, - MsgType: MsgData, - SessionID: sid, - PayloadLen: uint32(len(pkt)), + + if err := c.writeOnConn(conn, MsgData, pkt); err != nil { + return err } - // Copy conn under lock; release before write so Stop can close the conn and - // unblock a stuck writer via SetWriteDeadline / Close. + c.metrics.IncCounter("proxy_client_frames_out_total", nil) + c.metrics.IncCounter("uplink_packets_tx_total", map[string]string{"transport": "wss"}) + return nil +} + +func (c *Client) writeOnConn(conn Conn, msgType uint8, payload []byte) error { + c.mu.Lock() + if c.conn != conn { + c.mu.Unlock() + return ErrClientClosed + } + sid := c.sessID c.mu.Unlock() + c.writeMu.Lock() _ = conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)) - err := writeFrame(conn, hdr, pkt) + err := conn.WriteFrame(FrameHeader{ + Version: ProtocolVersion, + MsgType: msgType, + SessionID: sid, + PayloadLen: uint32(len(payload)), + }, payload) _ = conn.SetWriteDeadline(time.Time{}) - if err == nil { - c.metrics.IncCounter("proxy_client_frames_out_total", nil) + c.writeMu.Unlock() + + if err != nil { + c.log.Debug("uplink write failed, closing session", "msgType", msgType, "err", err) + _ = conn.Close() } return err } @@ -211,37 +295,48 @@ func (c *Client) run(ctx context.Context) { } c.setState(StateConnecting) - d := net.Dialer{Timeout: 15 * time.Second} - raw, err := d.DialContext(ctx, "tcp", c.addr) - if err != nil { - c.log.Warn("dial failed", "addr", c.addr, "err", err) - c.setState(StateFailed) - if !c.sleepBackoff(ctx, cur) { - return - } - cur = nextBackoff(cur, b) - continue - } + c.log.Info("uplink: connecting", + "transport", "wss", + "url", c.addr, + ) - tlsConn := tls.Client(raw, c.tlsConfig) - if err := tlsConn.HandshakeContext(ctx); err != nil { - c.log.Warn("tls handshake failed", "err", err) - _ = tlsConn.Close() + wsConn, err := DialWebSocket(ctx, c.addr, c.tlsConfig, c.opts.HandshakeTimeout) + if err != nil { + c.log.Warn("uplink: websocket dial failed", + "transport", "wss", + "url", c.addr, + "err", err, + ) c.setState(StateFailed) - if !c.sleepBackoff(ctx, cur) { + c.metrics.IncCounter("uplink_connection_errors_total", map[string]string{"transport": "wss"}) + c.metrics.IncCounter("uplink_reconnects_total", map[string]string{"transport": "wss"}) + c.log.Info("uplink: reconnect scheduled", + "transport", "wss", + "url", c.addr, + "backoff", cur.String(), + ) + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) continue } + wsConn.pingInterval = c.pingInterval + wsConn.pongWait = c.pongWait c.setState(StateTLSReady) + c.log.Info("uplink: websocket connection established", + "transport", "wss", + "url", c.addr, + "remote", wsConn.RemoteAddr(), + "status", wsConn.HandshakeStatus(), + ) hello, err := c.opts.HelloFactory() if err != nil { c.log.Error("hello factory", "err", err) - _ = tlsConn.Close() + _ = wsConn.Close() c.setState(StateFailed) - if !c.sleepBackoff(ctx, cur) { + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) @@ -250,52 +345,62 @@ func (c *Client) run(ctx context.Context) { c.setState(StateAuthenticating) payload, err := json.Marshal(hello) if err != nil { - _ = tlsConn.Close() + _ = wsConn.Close() c.setState(StateFailed) continue } - if err := writeFrame(tlsConn, FrameHeader{ + _ = wsConn.SetWriteDeadline(time.Now().Add(c.writeTimeout)) + err = wsConn.WriteFrame(FrameHeader{ Version: ProtocolVersion, MsgType: MsgHello, PayloadLen: uint32(len(payload)), - }, payload); err != nil { - _ = tlsConn.Close() + }, payload) + _ = wsConn.SetWriteDeadline(time.Time{}) + if err != nil { + _ = wsConn.Close() c.setState(StateFailed) - if !c.sleepBackoff(ctx, cur) { + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) continue } - _ = tlsConn.SetReadDeadline(time.Now().Add(30 * time.Second)) - h, ack, err := readFrame(tlsConn, c.maxPayload) + _ = wsConn.SetReadDeadline(time.Now().Add(30 * time.Second)) + h, ack, err := wsConn.ReadFrame(c.maxPayload) if err != nil { c.log.Warn("read hello ack", "err", err) - _ = tlsConn.Close() + _ = wsConn.Close() c.setState(StateFailed) - if !c.sleepBackoff(ctx, cur) { + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) continue } if h.MsgType == MsgError { - c.log.Warn("server rejected hello") - _ = tlsConn.Close() + c.log.Warn("uplink: authentication failed", + "transport", "wss", + "url", c.addr, + ) + _ = wsConn.Close() c.setState(StateFailed) c.metrics.IncCounter("proxy_client_auth_failures_total", nil) - if !c.sleepBackoff(ctx, cur) { + c.metrics.IncCounter("uplink_auth_failures_total", map[string]string{"transport": "wss"}) + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) continue } if h.MsgType != MsgHelloAck { - c.log.Warn("unexpected hello response", "type", h.MsgType) - _ = tlsConn.Close() + c.log.Warn("uplink: unexpected hello response", + "transport", "wss", + "type", h.MsgType, + ) + _ = wsConn.Close() c.setState(StateFailed) - if !c.sleepBackoff(ctx, cur) { + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) @@ -303,52 +408,73 @@ func (c *Client) run(ctx context.Context) { } var hw helloAckWire if err := json.Unmarshal(ack, &hw); err != nil { - _ = tlsConn.Close() + _ = wsConn.Close() c.setState(StateFailed) continue } + sessCtx, sessCancel := context.WithCancel(ctx) + wsConn.StartPingLoop(sessCtx) + c.mu.Lock() - c.tlsConn = tlsConn + c.conn = wsConn c.sessID = hw.SessionID c.mu.Unlock() c.setState(StateActive) cur = b.Initial c.metrics.IncCounter("proxy_client_sessions_active_total", nil) + c.log.Info("uplink: authenticated", + "transport", "wss", + "url", c.addr, + "remote", wsConn.RemoteAddr(), + "session_id", hw.SessionID, + ) - pingCtx, pingCancel := context.WithCancel(ctx) + pingCtx, pingCancel := context.WithCancel(sessCtx) pingDone := make(chan struct{}) go func() { defer close(pingDone) - c.pingLoop(pingCtx, tlsConn) + c.pingLoop(pingCtx, wsConn) }() readErr := make(chan error, 1) go func() { - readErr <- c.readLoop(ctx, tlsConn) + readErr <- c.readLoop(sessCtx, wsConn) }() select { case <-ctx.Done(): pingCancel() + sessCancel() <-pingDone - _ = tlsConn.Close() + _ = wsConn.Close() c.clearConn() c.setState(StateClosing) return case err := <-readErr: pingCancel() + sessCancel() <-pingDone if err != nil && !errors.Is(err, io.EOF) { c.log.Debug("read loop ended", "err", err) } - _ = tlsConn.Close() + _ = wsConn.Close() c.clearConn() c.setState(StateFailed) + c.log.Info("uplink: websocket disconnected", + "transport", "wss", + "url", c.addr, + ) + c.metrics.IncCounter("uplink_reconnects_total", map[string]string{"transport": "wss"}) if c.closed.Load() { return } - if !c.sleepBackoff(ctx, cur) { + c.log.Info("uplink: reconnect scheduled", + "transport", "wss", + "url", c.addr, + "backoff", cur.String(), + ) + if !c.sleepBackoff(ctx, jitterDuration(cur)) { return } cur = nextBackoff(cur, b) @@ -358,7 +484,7 @@ func (c *Client) run(ctx context.Context) { func (c *Client) clearConn() { c.mu.Lock() - c.tlsConn = nil + c.conn = nil c.sessID = 0 c.mu.Unlock() } @@ -382,7 +508,16 @@ func nextBackoff(cur time.Duration, b BackoffConfig) time.Duration { return next } -func (c *Client) pingLoop(ctx context.Context, conn *tls.Conn) { +func jitterDuration(d time.Duration) time.Duration { + if d <= 0 { + return d + } + // Full jitter in [d/2, d]. + half := d / 2 + return half + time.Duration(rand.Int63n(int64(half)+1)) +} + +func (c *Client) pingLoop(ctx context.Context, conn Conn) { t := time.NewTicker(c.keepAlive) defer t.Stop() for { @@ -391,7 +526,7 @@ func (c *Client) pingLoop(ctx context.Context, conn *tls.Conn) { return case <-t.C: c.mu.Lock() - active := c.tlsConn == conn && c.State() == StateActive + active := c.conn == conn && c.State() == StateActive c.mu.Unlock() if !active { return @@ -404,31 +539,28 @@ func (c *Client) pingLoop(ctx context.Context, conn *tls.Conn) { } } -func (c *Client) writeControl(conn *tls.Conn, msgType uint8, payload []byte) error { - c.mu.Lock() - defer c.mu.Unlock() - if c.tlsConn != conn { - return ErrClientClosed - } - sid := c.sessID - return writeFrame(conn, FrameHeader{ - Version: ProtocolVersion, - MsgType: msgType, - SessionID: sid, - PayloadLen: uint32(len(payload)), - }, payload) +func (c *Client) writeControl(conn Conn, msgType uint8, payload []byte) error { + return c.writeOnConn(conn, msgType, payload) } -func (c *Client) readLoop(ctx context.Context, conn *tls.Conn) error { +func (c *Client) readLoop(ctx context.Context, conn Conn) error { + idle := c.pongWait * 2 + if idle < c.keepAlive*3 { + idle = c.keepAlive * 3 + } for { select { case <-ctx.Done(): return ctx.Err() default: } - _ = conn.SetReadDeadline(time.Now().Add(c.keepAlive * 3)) - h, payload, err := readFrame(conn, c.maxPayload) + _ = conn.SetReadDeadline(time.Now().Add(idle)) + h, payload, err := conn.ReadFrame(c.maxPayload) if err != nil { + if ne, ok := err.(interface{ Timeout() bool }); ok && ne.Timeout() { + c.log.Info("uplink pong timeout") + c.metrics.IncCounter("uplink_ping_timeout_total", map[string]string{"transport": "wss"}) + } return err } switch h.MsgType { @@ -437,6 +569,7 @@ func (c *Client) readLoop(ctx context.Context, conn *tls.Conn) error { c.log.Warn("packet handler error", "err", err) } c.metrics.IncCounter("proxy_client_frames_in_total", nil) + c.metrics.IncCounter("uplink_packets_rx_total", map[string]string{"transport": "wss"}) case MsgPing: _ = c.writeControl(conn, MsgPong, nil) case MsgPong: diff --git a/uplink/clientip_test.go b/uplink/clientip_test.go new file mode 100644 index 0000000..11a8353 --- /dev/null +++ b/uplink/clientip_test.go @@ -0,0 +1,28 @@ +package uplink + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestClientIPFromRequest(t *testing.T) { + t.Parallel() + + req := httptest.NewRequest(http.MethodGet, "/uplink/v1", nil) + req.RemoteAddr = "203.0.113.10:54321" + if got := ClientIPFromRequest(req); got != "203.0.113.10" { + t.Fatalf("RemoteAddr: got %q", got) + } + + req.Header.Set("X-Forwarded-For", "198.51.100.7, 10.0.0.1") + if got := ClientIPFromRequest(req); got != "198.51.100.7" { + t.Fatalf("X-Forwarded-For: got %q", got) + } + + req.Header.Del("X-Forwarded-For") + req.Header.Set("X-Real-IP", "192.0.2.55") + if got := ClientIPFromRequest(req); got != "192.0.2.55" { + t.Fatalf("X-Real-IP: got %q", got) + } +} diff --git a/uplink/doc.go b/uplink/doc.go index 3b116ad..95bf29f 100644 --- a/uplink/doc.go +++ b/uplink/doc.go @@ -1,9 +1,19 @@ -// Package uplink provides a TCP/TLS framed transport for carrying WireGuard packet -// payloads between a relay-attached peer and its relay/gateway (Phase 1 uplink). +// Package uplink provides a WebSocket Secure (WSS) framed transport for carrying +// WireGuard packet payloads between a relay-attached peer and its relay/gateway. // -// It owns connection setup, TLS, framing, session lifecycle, keepalive, and -// peer→session registration for reverse traffic. It does not implement routing policy, -// relay selection, or Netmaker control-plane logic—integrate those in a separate adapter. +// The application transport is always WebSocket (/uplink/v1). TLS termination is +// configurable independently: +// +// - selfsigned (default): uplink terminates TLS with a persisted certificate +// - proxy: an upstream reverse proxy terminates TLS; uplink serves plain WS +// +// Each WebSocket binary message carries one uplink protocol frame (12-byte header +// + payload). The length prefix is retained as compatibility framing. +// +// It owns connection setup, framing, session lifecycle, keepalive, and +// peer→session registration for reverse traffic. It does not implement routing +// policy, relay selection, or Netmaker control-plane logic—integrate those in a +// separate adapter. // // For application-layer (HTTP CONNECT) egress, see package l7. package uplink diff --git a/uplink/example_test.go b/uplink/example_test.go index 6c5bcab..8fe1449 100644 --- a/uplink/example_test.go +++ b/uplink/example_test.go @@ -14,7 +14,7 @@ func ExampleClient() { defer cancel() c, err := uplink.NewClient(uplink.ClientOptions{ - Addr: "relay.example.com:443", + Addr: "wss://relay.example.com/uplink/v1", ServerName: "relay.example.com", TLSConfig: &tls.Config{MinVersion: tls.VersionTLS12}, HelloFactory: func() (uplink.ClientHello, error) { @@ -25,7 +25,6 @@ func ExampleClient() { }, nil }, PacketHandler: func(pkt []byte) error { - // Inject pkt into the local WireGuard/TUN receive path. _ = pkt return nil }, @@ -34,7 +33,7 @@ func ExampleClient() { panic(err) } _ = c.Start(ctx) - out := []byte{0x01} // placeholder WG packet bytes + out := []byte{0x01} _ = c.SendPacket(ctx, out) _ = c.Stop(context.Background()) fmt.Println(c.State()) diff --git a/uplink/frame.go b/uplink/frame.go index 82fd26e..693f86c 100644 --- a/uplink/frame.go +++ b/uplink/frame.go @@ -4,6 +4,7 @@ import ( "encoding/binary" "fmt" "io" + "sync" ) // FrameHeader is the 12-byte header (spec §11.1), big-endian. @@ -35,19 +36,42 @@ func decodeFrameHeader(b [frameHeaderSize]byte) FrameHeader { } } +var frameBufPool = sync.Pool{ + New: func() any { + b := make([]byte, 0, frameHeaderSize+DefaultMaxFrameSize) + return &b + }, +} + +// writeFrame emits a frame in a single Write. Writing the header and payload +// separately allows a concurrent writer on the same connection to interleave +// between them, and allows a write deadline to expire mid-frame; either leaves a +// truncated frame that desynchronises the peer's framing for the life of the +// connection. Callers must still serialise writes per connection and must close +// the connection if this returns an error, since a single Write can also flush +// partially. func writeFrame(w io.Writer, h FrameHeader, payload []byte) error { if uint32(len(payload)) != h.PayloadLen { return fmt.Errorf("proxy: payload length mismatch: %w", ErrInvalidFrame) } hdr := encodeFrameHeader(h) - if _, err := w.Write(hdr[:]); err != nil { + if len(payload) == 0 { + _, err := w.Write(hdr[:]) return err } - if len(payload) > 0 { - _, err := w.Write(payload) - return err + + need := frameHeaderSize + len(payload) + bufPtr := frameBufPool.Get().(*[]byte) + buf := (*bufPtr)[:0] + if cap(buf) < need { + buf = make([]byte, 0, need) } - return nil + buf = append(buf, hdr[:]...) + buf = append(buf, payload...) + _, err := w.Write(buf) + *bufPtr = buf[:0] + frameBufPool.Put(bufPtr) + return err } func readFrame(r io.Reader, maxPayload uint32) (FrameHeader, []byte, error) { @@ -71,3 +95,43 @@ func readFrame(r io.Reader, maxPayload uint32) (FrameHeader, []byte, error) { } return h, payload, nil } + +// encodeFrameBytes builds a complete framed message (header + payload) for +// message-oriented transports (e.g. one WebSocket binary message per frame). +// The length prefix is retained as compatibility framing; the message boundary +// also equals the uplink message boundary. +func encodeFrameBytes(h FrameHeader, payload []byte) ([]byte, error) { + if uint32(len(payload)) != h.PayloadLen { + return nil, fmt.Errorf("proxy: payload length mismatch: %w", ErrInvalidFrame) + } + hdr := encodeFrameHeader(h) + out := make([]byte, 0, frameHeaderSize+len(payload)) + out = append(out, hdr[:]...) + out = append(out, payload...) + return out, nil +} + +// decodeFrameBytes parses a complete framed message from a single buffer. +func decodeFrameBytes(msg []byte, maxPayload uint32) (FrameHeader, []byte, error) { + if len(msg) < frameHeaderSize { + return FrameHeader{}, nil, ErrInvalidFrame + } + var raw [frameHeaderSize]byte + copy(raw[:], msg[:frameHeaderSize]) + h := decodeFrameHeader(raw) + if h.Version != ProtocolVersion { + return FrameHeader{}, nil, ErrProtocolVersion + } + if h.PayloadLen > maxPayload { + return FrameHeader{}, nil, ErrInvalidFrame + } + if int(h.PayloadLen) != len(msg)-frameHeaderSize { + return FrameHeader{}, nil, ErrInvalidFrame + } + if h.PayloadLen == 0 { + return h, nil, nil + } + payload := make([]byte, h.PayloadLen) + copy(payload, msg[frameHeaderSize:]) + return h, payload, nil +} diff --git a/uplink/frame_test.go b/uplink/frame_test.go index c55456b..0389f31 100644 --- a/uplink/frame_test.go +++ b/uplink/frame_test.go @@ -36,3 +36,37 @@ func TestReadFrameOversize(t *testing.T) { t.Fatalf("want ErrInvalidFrame, got %v", err) } } + +func TestEncodeDecodeFrameBytes(t *testing.T) { + payload := []byte("hello-wg") + h := FrameHeader{ + Version: ProtocolVersion, + MsgType: MsgData, + SessionID: 7, + PayloadLen: uint32(len(payload)), + } + msg, err := encodeFrameBytes(h, payload) + if err != nil { + t.Fatal(err) + } + rh, rp, err := decodeFrameBytes(msg, DefaultMaxFrameSize) + if err != nil { + t.Fatal(err) + } + if rh.MsgType != MsgData || rh.SessionID != 7 || string(rp) != string(payload) { + t.Fatalf("got %+v %q", rh, rp) + } +} + +func TestDecodeFrameBytesOversize(t *testing.T) { + p := bytes.Repeat([]byte("x"), 100) + h := FrameHeader{Version: ProtocolVersion, MsgType: MsgData, PayloadLen: uint32(len(p))} + msg, err := encodeFrameBytes(h, p) + if err != nil { + t.Fatal(err) + } + _, _, err = decodeFrameBytes(msg, 50) + if err != ErrInvalidFrame { + t.Fatalf("want ErrInvalidFrame, got %v", err) + } +} diff --git a/uplink/integration_test.go b/uplink/integration_test.go index a522c31..9d697fe 100644 --- a/uplink/integration_test.go +++ b/uplink/integration_test.go @@ -8,13 +8,32 @@ import ( "crypto/tls" "crypto/x509" "crypto/x509/pkix" + "encoding/binary" "encoding/pem" + "fmt" "math/big" + "net" "sync" "testing" "time" ) +func testClientURL(t *testing.T, s *Server, scheme string) string { + t.Helper() + addr := s.Addr() + if addr == nil { + t.Fatal("nil server addr") + } + host, port, err := net.SplitHostPort(addr.String()) + if err != nil { + t.Fatal(err) + } + if host == "::" || host == "0.0.0.0" { + host = "127.0.0.1" + } + return fmt.Sprintf("%s://%s/uplink/v1", scheme, net.JoinHostPort(host, port)) +} + func TestClientServerDataRoundTrip(t *testing.T) { srvTLS, cliTLS := testTLSConfigs(t) @@ -35,6 +54,7 @@ func TestClientServerDataRoundTrip(t *testing.T) { s, err := NewServer(ServerOptions{ ListenAddr: "127.0.0.1:0", + TLSMode: TLSModeSelfSigned, TLSConfig: srvTLS, Authenticator: auth, PacketHandler: ph, @@ -51,12 +71,12 @@ func TestClientServerDataRoundTrip(t *testing.T) { if err := s.Start(ctx); err != nil { t.Fatal(err) } - addr := s.Addr().String() + url := testClientURL(t, s, "wss") var inbound [][]byte var inMu sync.Mutex cli, err := NewClient(ClientOptions{ - Addr: addr, + Addr: url, TLSConfig: func() *tls.Config { c := cliTLS.Clone() c.ServerName = "test.local" @@ -72,7 +92,7 @@ func TestClientServerDataRoundTrip(t *testing.T) { return nil }, Logger: noopLogger{}, - KeepAlivePeriod: time.Hour, // avoid ping noise in test + KeepAlivePeriod: time.Hour, }) if err != nil { t.Fatal(err) @@ -115,6 +135,420 @@ func TestClientServerDataRoundTrip(t *testing.T) { _ = s.Stop(context.Background()) } +func TestProxyTLSModePlainWS(t *testing.T) { + var received [][]byte + var mu sync.Mutex + + s, err := NewServer(ServerOptions{ + ListenAddr: "127.0.0.1:0", + TLSMode: TLSModeProxy, + Authenticator: &fnAuthenticator{fn: func(ctx context.Context, hello ClientHello) (*AuthResult, error) { + return &AuthResult{PeerID: "peer-1"}, nil + }}, + PacketHandler: &fnPacketHandler{fn: func(ctx context.Context, peerID string, pkt []byte) error { + mu.Lock() + received = append(received, append([]byte(nil), pkt...)) + mu.Unlock() + return nil + }}, + Logger: noopLogger{}, + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := s.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = s.Stop(context.Background()) }() + + cli, err := NewClient(ClientOptions{ + Addr: testClientURL(t, s, "ws"), + HelloFactory: func() (ClientHello, error) { + return ClientHello{Version: 1, NodeID: "n1", RelayPeerID: "r1", PublicKey: "pk", Proof: "p"}, nil + }, + PacketHandler: func(b []byte) error { return nil }, + Logger: noopLogger{}, + KeepAlivePeriod: time.Hour, + }) + if err != nil { + t.Fatal(err) + } + if err := cli.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = cli.Stop(context.Background()) }() + + deadline := time.After(5 * time.Second) + for cli.State() != StateActive { + select { + case <-deadline: + t.Fatalf("client not active: %s", cli.State()) + case <-time.After(10 * time.Millisecond): + } + } + if err := cli.SendPacket(ctx, []byte("hi")); err != nil { + t.Fatal(err) + } + time.Sleep(50 * time.Millisecond) + mu.Lock() + defer mu.Unlock() + if len(received) != 1 || string(received[0]) != "hi" { + t.Fatalf("got %v", received) + } +} + +func TestDefaultTLSModeSelfSigned(t *testing.T) { + mode, err := ParseTLSMode("") + if err != nil || mode != TLSModeSelfSigned { + t.Fatalf("got %q %v", mode, err) + } +} + +func TestInvalidTLSMode(t *testing.T) { + _, err := NewServer(ServerOptions{ + ListenAddr: "127.0.0.1:0", + TLSMode: "acme", + Authenticator: &fnAuthenticator{fn: func(ctx context.Context, hello ClientHello) (*AuthResult, error) { + return &AuthResult{PeerID: "p"}, nil + }}, + PacketHandler: &fnPacketHandler{fn: func(ctx context.Context, peerID string, pkt []byte) error { + return nil + }}, + }) + if err == nil { + t.Fatal("expected error for invalid TLS mode") + } +} + +func TestWrongRouteRejected(t *testing.T) { + s, err := NewServer(ServerOptions{ + ListenAddr: "127.0.0.1:0", + TLSMode: TLSModeProxy, + Authenticator: &fnAuthenticator{fn: func(ctx context.Context, hello ClientHello) (*AuthResult, error) { + return &AuthResult{PeerID: "p"}, nil + }}, + PacketHandler: &fnPacketHandler{fn: func(ctx context.Context, peerID string, pkt []byte) error { + return nil + }}, + Logger: noopLogger{}, + }) + if err != nil { + t.Fatal(err) + } + ctx := context.Background() + if err := s.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = s.Stop(context.Background()) }() + + addr := s.Addr().String() + _, err = DialWebSocket(ctx, "ws://"+addr+"/wrong", nil, time.Second) + if err == nil { + t.Fatal("expected dial to wrong path to fail") + } +} + +func TestAuthRejected(t *testing.T) { + s, err := NewServer(ServerOptions{ + ListenAddr: "127.0.0.1:0", + TLSMode: TLSModeProxy, + Authenticator: &fnAuthenticator{fn: func(ctx context.Context, hello ClientHello) (*AuthResult, error) { + return nil, errorsNew("nope") + }}, + PacketHandler: &fnPacketHandler{fn: func(ctx context.Context, peerID string, pkt []byte) error { + return nil + }}, + Logger: noopLogger{}, + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + if err := s.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = s.Stop(context.Background()) }() + + cli, err := NewClient(ClientOptions{ + Addr: testClientURL(t, s, "ws"), + HelloFactory: func() (ClientHello, error) { + return ClientHello{Version: 1, NodeID: "n1", RelayPeerID: "r1", PublicKey: "pk", Proof: "p"}, nil + }, + PacketHandler: func(b []byte) error { return nil }, + Logger: noopLogger{}, + KeepAlivePeriod: time.Hour, + ReconnectBackoff: BackoffConfig{Initial: 50 * time.Millisecond, Max: 100 * time.Millisecond, Factor: 1.5}, + }) + if err != nil { + t.Fatal(err) + } + if err := cli.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = cli.Stop(context.Background()) }() + + time.Sleep(300 * time.Millisecond) + if cli.State() == StateActive { + t.Fatal("client should not be active after auth failure") + } +} + +func errorsNew(s string) error { return fmt.Errorf("%s", s) } + +func TestServerStopClosesSessions(t *testing.T) { + srvTLS, cliTLS := testTLSConfigs(t) + + s, err := NewServer(ServerOptions{ + ListenAddr: "127.0.0.1:0", + TLSMode: TLSModeSelfSigned, + TLSConfig: srvTLS, + Authenticator: &fnAuthenticator{fn: func(ctx context.Context, hello ClientHello) (*AuthResult, error) { + return &AuthResult{PeerID: "peer-1"}, nil + }}, + PacketHandler: &fnPacketHandler{fn: func(ctx context.Context, peerID string, pkt []byte) error { + return nil + }}, + Logger: noopLogger{}, + }) + if err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if err := s.Start(ctx); err != nil { + t.Fatal(err) + } + + cli, err := NewClient(ClientOptions{ + Addr: testClientURL(t, s, "wss"), + TLSConfig: func() *tls.Config { + c := cliTLS.Clone() + c.ServerName = "test.local" + return c + }(), + HelloFactory: func() (ClientHello, error) { + return ClientHello{Version: 1, NodeID: "n1", RelayPeerID: "r1", PublicKey: "pk", Proof: "p"}, nil + }, + PacketHandler: func(b []byte) error { return nil }, + Logger: noopLogger{}, + KeepAlivePeriod: time.Hour, + }) + if err != nil { + t.Fatal(err) + } + if err := cli.Start(ctx); err != nil { + t.Fatal(err) + } + + deadline := time.After(5 * time.Second) + for cli.State() != StateActive { + select { + case <-deadline: + t.Fatalf("client not active: %s", cli.State()) + case <-time.After(10 * time.Millisecond): + } + } + + stopCtx, stopCancel := context.WithTimeout(context.Background(), 3*time.Second) + defer stopCancel() + if err := s.Stop(stopCtx); err != nil { + t.Fatalf("Stop: %v", err) + } + + deadline = time.After(3 * time.Second) + for cli.State() == StateActive { + select { + case <-deadline: + t.Fatal("client still Active after Server.Stop; session was not closed") + case <-time.After(10 * time.Millisecond): + } + } + + _ = cli.Stop(context.Background()) +} + +func TestConcurrentSendPacketFramingIntegrity(t *testing.T) { + const ( + senders = 8 + perSender = 150 + wantPackets = senders * perSender + payloadMinSize = 200 + ) + + srvTLS, cliTLS := testTLSConfigs(t) + + var mu sync.Mutex + seen := make(map[int]int, wantPackets) + corrupt := 0 + done := make(chan struct{}) + var closeOnce sync.Once + + ph := &fnPacketHandler{fn: func(ctx context.Context, peerID string, pkt []byte) error { + seq, ok := decodeTestPacket(pkt) + mu.Lock() + if !ok { + corrupt++ + } else { + seen[seq]++ + } + total := len(seen) + corrupt + mu.Unlock() + if total >= wantPackets { + closeOnce.Do(func() { close(done) }) + } + return nil + }} + + s, err := NewServer(ServerOptions{ + ListenAddr: "127.0.0.1:0", + TLSMode: TLSModeSelfSigned, + TLSConfig: srvTLS, + Authenticator: &fnAuthenticator{fn: func(ctx context.Context, hello ClientHello) (*AuthResult, error) { + return &AuthResult{PeerID: "peer-1"}, nil + }}, + PacketHandler: ph, + Logger: noopLogger{}, + }) + if err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := s.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = s.Stop(context.Background()) }() + + cli, err := NewClient(ClientOptions{ + Addr: testClientURL(t, s, "wss"), + TLSConfig: func() *tls.Config { + c := cliTLS.Clone() + c.ServerName = "test.local" + return c + }(), + HelloFactory: func() (ClientHello, error) { + return ClientHello{Version: 1, NodeID: "n1", RelayPeerID: "r1", PublicKey: "pk", Proof: "p"}, nil + }, + PacketHandler: func(b []byte) error { return nil }, + Logger: noopLogger{}, + KeepAlivePeriod: 2 * time.Second, + WriteTimeout: 10 * time.Second, + }) + if err != nil { + t.Fatal(err) + } + if err := cli.Start(ctx); err != nil { + t.Fatal(err) + } + defer func() { _ = cli.Stop(context.Background()) }() + + deadline := time.After(5 * time.Second) + for cli.State() != StateActive { + select { + case <-deadline: + t.Fatalf("client not active: %s", cli.State()) + case <-time.After(10 * time.Millisecond): + } + } + + cli.mu.Lock() + conn := cli.conn + cli.mu.Unlock() + if conn == nil { + t.Fatal("no active connection") + } + + pingStop := make(chan struct{}) + var pingWG sync.WaitGroup + pingWG.Add(1) + go func() { + defer pingWG.Done() + for { + select { + case <-pingStop: + return + default: + } + if err := cli.writeControl(conn, MsgPing, nil); err != nil { + return + } + time.Sleep(100 * time.Microsecond) + } + }() + + var wg sync.WaitGroup + for i := 0; i < senders; i++ { + wg.Add(1) + go func(sender int) { + defer wg.Done() + for j := 0; j < perSender; j++ { + seq := sender*perSender + j + if err := cli.SendPacket(ctx, encodeTestPacket(seq, payloadMinSize+seq%512)); err != nil { + t.Errorf("SendPacket seq=%d: %v", seq, err) + return + } + } + }(i) + } + wg.Wait() + close(pingStop) + pingWG.Wait() + + select { + case <-done: + case <-time.After(10 * time.Second): + } + + mu.Lock() + defer mu.Unlock() + if corrupt != 0 { + t.Errorf("server received %d corrupt packets; frame writes interleaved", corrupt) + } + if len(seen) != wantPackets { + t.Errorf("server received %d distinct packets, want %d", len(seen), wantPackets) + } + for seq, n := range seen { + if n != 1 { + t.Errorf("packet seq=%d received %d times, want 1", seq, n) + } + } +} + +func encodeTestPacket(seq, size int) []byte { + if size < 8 { + size = 8 + } + pkt := make([]byte, size) + binary.BigEndian.PutUint32(pkt[0:4], uint32(seq)) + binary.BigEndian.PutUint32(pkt[4:8], uint32(size)) + for i := 8; i < size; i++ { + pkt[i] = byte(seq % 251) + } + return pkt +} + +func decodeTestPacket(pkt []byte) (int, bool) { + if len(pkt) < 8 { + return 0, false + } + seq := int(binary.BigEndian.Uint32(pkt[0:4])) + size := int(binary.BigEndian.Uint32(pkt[4:8])) + if size != len(pkt) { + return 0, false + } + for i := 8; i < len(pkt); i++ { + if pkt[i] != byte(seq%251) { + return 0, false + } + } + return seq, true +} + type fnPacketHandler struct { fn func(context.Context, string, []byte) error } @@ -144,6 +578,7 @@ func testTLSConfigs(t *testing.T) (server *tls.Config, client *tls.Config) { NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(24 * time.Hour), DNSNames: []string{"test.local"}, + IPAddresses: []net.IP{net.ParseIP("127.0.0.1")}, KeyUsage: x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, } diff --git a/uplink/registry.go b/uplink/registry.go index 76e4391..846eed3 100644 --- a/uplink/registry.go +++ b/uplink/registry.go @@ -36,13 +36,53 @@ func (r *InMemoryRegistry) Get(peerID string) (Session, bool) { return s, ok && s != nil } -// Detach removes a peer from the registry. +// Detach removes a peer from the registry, whichever session is registered. func (r *InMemoryRegistry) Detach(peerID string) { r.mu.Lock() defer r.mu.Unlock() delete(r.m, peerID) } +// DetachSession removes peerID only if sess is still the registered session. +// A reconnecting client attaches its new session before the old session's read +// loop finishes unwinding, so an unconditional Detach from the old session would +// evict the live one and leave the peer looking session-less. +func (r *InMemoryRegistry) DetachSession(peerID string, sess Session) { + r.mu.Lock() + defer r.mu.Unlock() + if cur, ok := r.m[peerID]; ok && cur != sess { + return + } + delete(r.m, peerID) +} + +// PeerIDs returns the peer IDs with a registered session. +func (r *InMemoryRegistry) PeerIDs() []string { + r.mu.RLock() + defer r.mu.RUnlock() + ids := make([]string, 0, len(r.m)) + for id, s := range r.m { + if s != nil { + ids = append(ids, id) + } + } + return ids +} + +// CloseAll closes every registered session and clears the registry. +// Used by Server.Stop so clients drop and re-HELLO after a gateway restart +// (closing the listener alone leaves accepted TLS conns answering PING forever). +func (r *InMemoryRegistry) CloseAll() { + r.mu.Lock() + defer r.mu.Unlock() + for id, s := range r.m { + if c, ok := s.(interface{ Close() error }); ok { + _ = c.Close() + } + delete(r.m, id) + } +} + // Len returns the number of registered sessions (for metrics / tests). func (r *InMemoryRegistry) Len() int { r.mu.RLock() diff --git a/uplink/registry_test.go b/uplink/registry_test.go new file mode 100644 index 0000000..4433266 --- /dev/null +++ b/uplink/registry_test.go @@ -0,0 +1,77 @@ +package uplink + +import ( + "sort" + "testing" +) + +type fakeSession struct { + peerID string + closed bool +} + +func (f *fakeSession) PeerID() string { return f.peerID } +func (f *fakeSession) State() SessionState { return SessionAttached } +func (f *fakeSession) Close() error { + f.closed = true + return nil +} + +func TestInMemoryRegistryPeerIDs(t *testing.T) { + r := NewInMemoryRegistry() + if got := r.PeerIDs(); len(got) != 0 { + t.Fatalf("empty registry returned %v", got) + } + + for _, id := range []string{"a", "b", "c"} { + if err := r.Attach(id, &fakeSession{peerID: id}); err != nil { + t.Fatal(err) + } + } + got := r.PeerIDs() + sort.Strings(got) + if len(got) != 3 || got[0] != "a" || got[1] != "b" || got[2] != "c" { + t.Fatalf("PeerIDs = %v, want [a b c]", got) + } + + r.Detach("b") + got = r.PeerIDs() + sort.Strings(got) + if len(got) != 2 || got[0] != "a" || got[1] != "c" { + t.Fatalf("after Detach PeerIDs = %v, want [a c]", got) + } +} + +// TestDetachSessionKeepsReplacement covers a client reconnect: the new session is +// attached while the old session's read loop is still unwinding, so the old +// session's detach must not evict the live replacement. +func TestDetachSessionKeepsReplacement(t *testing.T) { + r := NewInMemoryRegistry() + old := &fakeSession{peerID: "peer-1"} + if err := r.Attach("peer-1", old); err != nil { + t.Fatal(err) + } + + replacement := &fakeSession{peerID: "peer-1"} + if err := r.Attach("peer-1", replacement); err != nil { + t.Fatal(err) + } + if !old.closed { + t.Error("replaced session was not closed on re-attach") + } + + r.DetachSession("peer-1", old) + + sess, ok := r.Get("peer-1") + if !ok { + t.Fatal("live session evicted by the replaced session's detach") + } + if sess != Session(replacement) { + t.Fatalf("registry holds %v, want the replacement session", sess) + } + + r.DetachSession("peer-1", replacement) + if _, ok := r.Get("peer-1"); ok { + t.Error("session still registered after its own detach") + } +} diff --git a/uplink/server.go b/uplink/server.go index 65da31e..d938c7d 100644 --- a/uplink/server.go +++ b/uplink/server.go @@ -5,28 +5,38 @@ import ( "crypto/tls" "encoding/json" "errors" + "fmt" "io" "net" + "net/http" "sync" "sync/atomic" "time" ) -// ServerOptions configures the TCP/TLS uplink server (spec §9.2). +// ServerOptions configures the WebSocket uplink server. type ServerOptions struct { - ListenAddr string + ListenAddr string + // TLSMode selects TLS termination. Empty defaults to TLSModeSelfSigned. + TLSMode TLSMode + // TLSConfig is required when TLSMode is selfsigned; ignored in proxy mode. TLSConfig *tls.Config Authenticator Authenticator PacketHandler PacketHandler SessionRegistry SessionRegistry Logger Logger Metrics MetricsSink - KeepAlivePeriod time.Duration + KeepAlivePeriod time.Duration // used for app MsgPing fallback / read idle WriteTimeout time.Duration MaxFrameSize int + // PingInterval for WebSocket Ping control frames (default 25s). + PingInterval time.Duration + // PongWait is how long after a Ping to allow for a Pong / read idle slack + // (default 60s). Effective read deadline is PingInterval+PongWait. + PongWait time.Duration } -// Server terminates TLS and manages framed sessions (spec §9.2). +// Server serves /uplink/v1 over WebSocket and manages framed sessions. type Server struct { opts ServerOptions log Logger @@ -35,8 +45,13 @@ type Server struct { maxPayload uint32 keepAlive time.Duration writeTimeout time.Duration + pingInterval time.Duration + pongWait time.Duration + tlsMode TLSMode sessionIDGen atomic.Uint32 + httpServer *http.Server listener net.Listener + serveCtx context.Context closed atomic.Bool shutdown chan struct{} shutdownOnce sync.Once @@ -50,8 +65,15 @@ func NewServer(opts ServerOptions) (*Server, error) { if opts.ListenAddr == "" { return nil, errors.New("proxy: ListenAddr is required") } - if opts.TLSConfig == nil { - return nil, errors.New("proxy: TLSConfig is required") + mode := opts.TLSMode + if mode == "" { + mode = DefaultTLSMode + } + if err := mode.Validate(); err != nil { + return nil, err + } + if mode == TLSModeSelfSigned && opts.TLSConfig == nil { + return nil, errors.New("proxy: TLSConfig is required for selfsigned TLS mode") } if opts.Authenticator == nil { return nil, errors.New("proxy: Authenticator is required") @@ -83,6 +105,14 @@ func NewServer(opts ServerOptions) (*Server, error) { if wt <= 0 { wt = 2 * time.Second } + pi := opts.PingInterval + if pi <= 0 { + pi = defaultWSPingInterval + } + pw := opts.PongWait + if pw <= 0 { + pw = defaultWSPongWait + } return &Server{ opts: opts, log: log, @@ -91,60 +121,150 @@ func NewServer(opts ServerOptions) (*Server, error) { maxPayload: uint32(max), keepAlive: ka, writeTimeout: wt, + pingInterval: pi, + pongWait: pw, + tlsMode: mode, shutdown: make(chan struct{}), }, nil } -// Start listens with TLS and accepts sessions until ctx is cancelled or Stop is called. +// Start binds the listen address and serves WebSocket uplink until ctx is cancelled or Stop. func (s *Server) Start(ctx context.Context) error { s.startMu.Lock() - if s.listener != nil { + if s.httpServer != nil { s.startMu.Unlock() return errors.New("proxy: server already started") } - ln, err := tls.Listen("tcp", s.opts.ListenAddr, s.opts.TLSConfig) + + mux := http.NewServeMux() + mux.HandleFunc(UplinkWSPath, s.handleUplinkHTTP) + + hs := &http.Server{ + Addr: s.opts.ListenAddr, + Handler: mux, + ReadHeaderTimeout: 10 * time.Second, + // Do not set WriteTimeout / ReadTimeout — they apply to upgraded WebSockets. + } + + ln, err := net.Listen("tcp", s.opts.ListenAddr) if err != nil { s.startMu.Unlock() return err } + if s.tlsMode == TLSModeSelfSigned { + ln = tls.NewListener(ln, s.opts.TLSConfig) + } + s.listener = ln + s.httpServer = hs + s.serveCtx = ctx s.startMu.Unlock() + labels := map[string]string{"transport": "wss", "tls_mode": string(s.tlsMode)} + s.log.Info("uplink websocket listener started", + "listen", ln.Addr().String(), + "tls_mode", string(s.tlsMode), + "path", UplinkWSPath, + ) + if s.tlsMode == TLSModeProxy { + s.log.Info("uplink started in reverse-proxy TLS mode", "listen", ln.Addr().String()) + if warnPublicListen(ln.Addr()) { + s.log.Warn("proxy TLS mode listen address is not loopback/private; ensure only a trusted reverse proxy can reach this listener", + "listen", ln.Addr().String()) + } + } + s.metrics.IncCounter("uplink_connections_total", labels) // zero-ish marker for scrape presence + s.serveWG.Add(1) go func() { defer s.serveWG.Done() - for { - conn, err := ln.Accept() - if err != nil { - if s.closed.Load() { - return - } - select { - case <-ctx.Done(): - return - case <-s.shutdown: - return - default: - } - s.log.Warn("accept failed", "err", err) - continue - } - tlsConn, ok := conn.(*tls.Conn) - if !ok { - _ = conn.Close() - continue - } - s.connWG.Add(1) - go func(c *tls.Conn) { - defer s.connWG.Done() - s.handleConnection(ctx, c) - }(tlsConn) + var serveErr error + if s.tlsMode == TLSModeSelfSigned { + serveErr = hs.Serve(ln) + } else { + serveErr = hs.Serve(ln) + } + if serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) && !s.closed.Load() { + s.log.Warn("http serve ended", "err", serveErr) + } + }() + + go func() { + select { + case <-ctx.Done(): + _ = s.Stop(context.Background()) + case <-s.shutdown: } }() + return nil } -// Addr returns the bound listen address (e.g. after Start with ":0"). It is nil before Start or after Stop. +func warnPublicListen(addr net.Addr) bool { + ta, ok := addr.(*net.TCPAddr) + if !ok || ta.IP == nil { + return false + } + if ta.IP.IsUnspecified() { + return true + } + if ta.IP.IsLoopback() || ta.IP.IsPrivate() || ta.IP.IsLinkLocalUnicast() { + return false + } + return true +} + +func (s *Server) handleUplinkHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + clientIP := ClientIPFromRequest(r) + s.log.Info("uplink: websocket upgrade request", + "transport", "wss", + "tls_mode", string(s.tlsMode), + "client_ip", clientIP, + "remote", r.RemoteAddr, + "path", r.URL.Path, + ) + wsConn, err := UpgradeWebSocket(w, r, int(s.maxPayload)+frameHeaderSize) + if err != nil { + s.log.Warn("uplink: websocket upgrade failed", + "transport", "wss", + "client_ip", clientIP, + "remote", r.RemoteAddr, + "err", err, + ) + s.metrics.IncCounter("uplink_connection_errors_total", map[string]string{ + "transport": "wss", "tls_mode": string(s.tlsMode), + }) + return + } + wsConn.pingInterval = s.pingInterval + wsConn.pongWait = s.pongWait + _ = wsConn.SetReadDeadline(time.Now().Add(s.pongWait * 2)) + s.log.Info("uplink: websocket connection established", + "transport", "wss", + "tls_mode", string(s.tlsMode), + "client_ip", clientIP, + "remote", r.RemoteAddr, + "status", wsConn.HandshakeStatus(), + ) + + s.connWG.Add(1) + go func() { + defer s.connWG.Done() + // Do not use r.Context(): it is cancelled when this HTTP handler returns + // after a successful WebSocket upgrade. + base := s.serveCtx + if base == nil { + base = context.Background() + } + s.handleConnection(base, wsConn, clientIP) + }() +} + +// Addr returns the bound listen address (e.g. after Start with ":0"). func (s *Server) Addr() net.Addr { s.startMu.Lock() defer s.startMu.Unlock() @@ -154,18 +274,31 @@ func (s *Server) Addr() net.Addr { return s.listener.Addr() } -// Stop closes the listener and waits for handlers to finish. +// TLSMode returns the configured TLS mode. +func (s *Server) TLSMode() TLSMode { return s.tlsMode } + +// Stop closes the HTTP server, attached sessions, and waits for handlers. func (s *Server) Stop(ctx context.Context) error { s.shutdownOnce.Do(func() { close(s.shutdown) s.closed.Store(true) s.startMu.Lock() + hs := s.httpServer + s.httpServer = nil ln := s.listener s.listener = nil s.startMu.Unlock() + if hs != nil { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + _ = hs.Shutdown(shutdownCtx) + cancel() + } if ln != nil { _ = ln.Close() } + if closer, ok := s.registry.(interface{ CloseAll() }); ok { + closer.CloseAll() + } }) done := make(chan struct{}) @@ -200,41 +333,101 @@ func (s *Server) SendToPeer(ctx context.Context, peerID string, pkt []byte) erro return sw.sendData(ctx, pkt) } -func (s *Server) handleConnection(ctx context.Context, conn *tls.Conn) { +// SessionPeerIDs returns peer IDs with an attached session, if supported. +func (s *Server) SessionPeerIDs() []string { + if reg, ok := s.registry.(interface{ PeerIDs() []string }); ok { + return reg.PeerIDs() + } + return nil +} + +func (s *Server) detach(peerID string, sess Session) { + if reg, ok := s.registry.(interface { + DetachSession(string, Session) + }); ok { + reg.DetachSession(peerID, sess) + return + } + s.registry.Detach(peerID) +} + +func (s *Server) handleConnection(ctx context.Context, conn Conn, clientIP string) { defer conn.Close() + labels := map[string]string{"transport": "wss", "tls_mode": string(s.tlsMode)} s.metrics.IncCounter("proxy_server_connections_total", nil) + s.metrics.IncCounter("uplink_connections_total", labels) + if clientIP == "" { + if ws, ok := conn.(*WebSocketConn); ok { + clientIP = ws.RemoteAddr() + } + } + s.log.Info("uplink: session started", + "transport", "wss", + "tls_mode", string(s.tlsMode), + "client_ip", clientIP, + ) + + sessCtx, sessCancel := context.WithCancel(ctx) + defer sessCancel() - h, payload, err := readFrame(conn, s.maxPayload) + if ws, ok := conn.(*WebSocketConn); ok { + ws.StartPingLoop(sessCtx) + } + + _ = conn.SetReadDeadline(time.Now().Add(30 * time.Second)) + h, payload, err := conn.ReadFrame(s.maxPayload) if err != nil { - s.log.Warn("read hello frame failed", "err", err) + s.log.Warn("uplink: read hello frame failed", + "client_ip", clientIP, + "err", err, + ) return } if h.MsgType != MsgHello { - s.log.Warn("expected MsgHello", "msgType", h.MsgType) + s.log.Warn("uplink: expected MsgHello", + "client_ip", clientIP, + "msgType", h.MsgType, + ) return } var hello ClientHello if err := json.Unmarshal(payload, &hello); err != nil { - s.log.Warn("invalid ClientHello JSON", "err", err) + s.log.Warn("uplink: invalid ClientHello JSON", + "client_ip", clientIP, + "err", err, + ) _ = s.writeError(conn, 0, "invalid_hello", "invalid hello payload") return } - authCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + authCtx, cancel := context.WithTimeout(sessCtx, 30*time.Second) res, err := s.opts.Authenticator.ValidateClientHello(authCtx, hello) cancel() if err != nil || res == nil || res.PeerID == "" { s.metrics.IncCounter("proxy_server_auth_failures_total", nil) + s.metrics.IncCounter("uplink_auth_failures_total", labels) msg := "authentication failed" if err != nil { msg = err.Error() } _ = s.writeError(conn, 0, "auth_failed", msg) - s.log.Warn("auth failed", "node_id", hello.NodeID, "err", err) + s.log.Warn("uplink: authentication failed", + "transport", "wss", + "client_ip", clientIP, + "node_id", hello.NodeID, + "err", err, + ) return } + s.log.Info("uplink: authenticated", + "transport", "wss", + "tls_mode", string(s.tlsMode), + "client_ip", clientIP, + "peer", res.PeerID, + "node_id", hello.NodeID, + ) sid := s.sessionIDGen.Add(1) cs := &connSession{ @@ -243,6 +436,8 @@ func (s *Server) handleConnection(ctx context.Context, conn *tls.Conn) { peerID: res.PeerID, sessionID: sid, state: SessionAuthenticated, + cancel: sessCancel, + clientIP: clientIP, } ackBytes, err := json.Marshal(helloAckWire{SessionID: sid}) @@ -251,55 +446,72 @@ func (s *Server) handleConnection(ctx context.Context, conn *tls.Conn) { return } if err := cs.writeFrameLocked(MsgHelloAck, ackBytes); err != nil { - s.log.Warn("write hello ack", "err", err) + s.log.Warn("write hello ack", "err", err, "client_ip", clientIP) return } cs.setState(SessionAttached) if err := s.registry.Attach(res.PeerID, cs); err != nil { - s.log.Error("session attach", "err", err) + s.log.Error("session attach", "err", err, "client_ip", clientIP) _ = cs.Close() return } s.metrics.IncCounter("proxy_server_sessions_attached_total", nil) if reg, ok := s.registry.(*InMemoryRegistry); ok { s.metrics.SetGauge("proxy_server_active_sessions", float64(reg.Len()), nil) + s.metrics.SetGauge("uplink_connections_active", float64(reg.Len()), labels) } defer func() { - s.registry.Detach(res.PeerID) + s.detach(res.PeerID, cs) cs.setState(SessionClosed) if reg, ok := s.registry.(*InMemoryRegistry); ok { s.metrics.SetGauge("proxy_server_active_sessions", float64(reg.Len()), nil) + s.metrics.SetGauge("uplink_connections_active", float64(reg.Len()), labels) } + s.log.Info("uplink: session cleanup complete", + "transport", "wss", + "client_ip", clientIP, + "peer", res.PeerID, + ) + s.log.Info("uplink: websocket disconnected", + "transport", "wss", + "client_ip", clientIP, + "peer", res.PeerID, + ) }() - cs.readLoop(ctx) + cs.readLoop(sessCtx) } -func (s *Server) writeError(conn net.Conn, sessionID uint32, code, msg string) error { +func (s *Server) writeError(conn Conn, sessionID uint32, code, msg string) error { b, err := json.Marshal(errorPayloadWire{Code: code, Message: msg}) if err != nil { return err } - return writeFrame(conn, FrameHeader{ + _ = conn.SetWriteDeadline(time.Now().Add(s.writeTimeout)) + err = conn.WriteFrame(FrameHeader{ Version: ProtocolVersion, MsgType: MsgError, SessionID: sessionID, PayloadLen: uint32(len(b)), }, b) + _ = conn.SetWriteDeadline(time.Time{}) + return err } -// connSession is the server's session for one TLS connection. +// connSession is the server's session for one WebSocket connection. type connSession struct { server *Server - conn *tls.Conn + conn Conn peerID string sessionID uint32 + clientIP string writeMu sync.Mutex state SessionState stateMu sync.RWMutex closed atomic.Bool + cancel context.CancelFunc } func (c *connSession) PeerID() string { return c.peerID } @@ -320,11 +532,14 @@ func (c *connSession) Close() error { if !c.closed.CompareAndSwap(false, true) { return nil } + if c.cancel != nil { + c.cancel() + } return c.conn.Close() } func (c *connSession) writeFrameLocked(msgType uint8, payload []byte) error { - return writeFrame(c.conn, FrameHeader{ + return c.conn.WriteFrame(FrameHeader{ Version: ProtocolVersion, MsgType: msgType, SessionID: c.sessionID, @@ -334,10 +549,16 @@ func (c *connSession) writeFrameLocked(msgType uint8, payload []byte) error { func (c *connSession) writeFrame(msgType uint8, payload []byte) error { c.writeMu.Lock() - defer c.writeMu.Unlock() _ = c.conn.SetWriteDeadline(time.Now().Add(c.server.writeTimeout)) err := c.writeFrameLocked(msgType, payload) _ = c.conn.SetWriteDeadline(time.Time{}) + c.writeMu.Unlock() + + if err != nil { + c.server.log.Debug("session write failed, closing session", + "peer", c.peerID, "msgType", msgType, "err", err) + _ = c.Close() + } return err } @@ -352,16 +573,44 @@ func (c *connSession) sendData(_ context.Context, pkt []byte) error { } func (c *connSession) readLoop(ctx context.Context) { + idle := c.server.pongWait * 2 + if idle < c.server.keepAlive*3 { + idle = c.server.keepAlive * 3 + } for { if c.closed.Load() { return } - _ = c.conn.SetReadDeadline(time.Now().Add(c.server.keepAlive * 3)) - h, payload, err := readFrame(c.conn, c.server.maxPayload) + select { + case <-ctx.Done(): + _ = c.Close() + return + case <-c.server.shutdown: + _ = c.Close() + return + default: + } + _ = c.conn.SetReadDeadline(time.Now().Add(idle)) + h, payload, err := c.conn.ReadFrame(c.server.maxPayload) if err != nil { if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) { return } + if ne, ok := err.(net.Error); ok && ne.Timeout() { + c.server.log.Info("uplink pong timeout", "peer", c.peerID) + c.server.metrics.IncCounter("uplink_ping_timeout_total", map[string]string{ + "transport": "wss", "tls_mode": string(c.server.tlsMode), + }) + select { + case <-ctx.Done(): + _ = c.Close() + return + case <-c.server.shutdown: + _ = c.Close() + return + default: + } + } c.server.log.Debug("read frame ended", "peer", c.peerID, "err", err) _ = c.Close() return @@ -373,6 +622,8 @@ func (c *connSession) readLoop(ctx context.Context) { c.server.log.Warn("packet handler", "peer", c.peerID, "err", err) } c.server.metrics.IncCounter("proxy_server_frames_in_total", nil) + c.server.metrics.IncCounter("uplink_packets_rx_total", map[string]string{"transport": "wss"}) + c.server.metrics.IncCounter("uplink_bytes_rx_total", map[string]string{"transport": "wss"}) case MsgPing: _ = c.writeFrame(MsgPong, nil) case MsgPong: @@ -388,3 +639,8 @@ func (c *connSession) readLoop(ctx context.Context) { } } } + +// EndpointURL builds a client-facing wss:// URL for the given host and port. +func EndpointURL(host string, port int) string { + return fmt.Sprintf("wss://%s/uplink/v1", net.JoinHostPort(host, fmt.Sprintf("%d", port))) +} diff --git a/uplink/tlsmode.go b/uplink/tlsmode.go new file mode 100644 index 0000000..4fc7af4 --- /dev/null +++ b/uplink/tlsmode.go @@ -0,0 +1,47 @@ +package uplink + +import ( + "errors" + "fmt" + "strings" +) + +// TLSMode controls where TLS is terminated for the uplink HTTP/WebSocket listener. +type TLSMode string + +const ( + // TLSModeSelfSigned: uplink terminates TLS with a persisted self-signed certificate. + TLSModeSelfSigned TLSMode = "selfsigned" + // TLSModeProxy: upstream reverse proxy terminates TLS; uplink serves plain HTTP/WS. + TLSModeProxy TLSMode = "proxy" +) + +// DefaultTLSMode is used when TLS mode is omitted. +const DefaultTLSMode = TLSModeSelfSigned + +// ParseTLSMode validates and normalises a TLS mode string. +// Empty input returns DefaultTLSMode (selfsigned). +func ParseTLSMode(s string) (TLSMode, error) { + m := TLSMode(strings.ToLower(strings.TrimSpace(s))) + if m == "" { + return DefaultTLSMode, nil + } + switch m { + case TLSModeSelfSigned, TLSModeProxy: + return m, nil + default: + return "", fmt.Errorf("unsupported uplink TLS mode: %s", s) + } +} + +// Validate reports whether m is a supported TLS mode. +func (m TLSMode) Validate() error { + switch m { + case TLSModeSelfSigned, TLSModeProxy: + return nil + case "": + return errors.New("unsupported uplink TLS mode: (empty)") + default: + return fmt.Errorf("unsupported uplink TLS mode: %s", m) + } +} diff --git a/uplink/transport.go b/uplink/transport.go new file mode 100644 index 0000000..7c184cb --- /dev/null +++ b/uplink/transport.go @@ -0,0 +1,17 @@ +package uplink + +import "time" + +// Conn is the transport used by uplink session logic. +// Implementations must be safe for one reader and one writer concurrently +// (writer serialisation is the caller's responsibility unless documented otherwise). +type Conn interface { + // ReadFrame reads the next uplink protocol frame. + ReadFrame(maxPayload uint32) (FrameHeader, []byte, error) + // WriteFrame writes one uplink protocol frame. + WriteFrame(h FrameHeader, payload []byte) error + // Close closes the underlying transport. + Close() error + SetReadDeadline(t time.Time) error + SetWriteDeadline(t time.Time) error +} diff --git a/uplink/websocket.go b/uplink/websocket.go new file mode 100644 index 0000000..601f2ab --- /dev/null +++ b/uplink/websocket.go @@ -0,0 +1,274 @@ +package uplink + +import ( + "context" + "crypto/tls" + "fmt" + "net" + "net/http" + "net/url" + "strings" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +const ( + // UplinkWSPath is the versioned WebSocket endpoint for TCP uplink. + UplinkWSPath = "/uplink/v1" + + // DefaultMaxMessageSize caps a single WebSocket binary message (frame header + payload). + DefaultMaxMessageSize = frameHeaderSize + DefaultMaxFrameSize + + defaultWSPingInterval = 25 * time.Second + defaultWSPongWait = 60 * time.Second +) + +// wsUpgrader upgrades HTTP connections to WebSocket. +// CheckOrigin allows missing Origin (machine clients / netclient are not browsers). +// Authentication remains the security boundary after upgrade. +var wsUpgrader = websocket.Upgrader{ + ReadBufferSize: 4096, + WriteBufferSize: 4096, + CheckOrigin: func(r *http.Request) bool { + return true + }, +} + +// WebSocketConn implements Conn over a gorilla WebSocket connection. +// One reader and one writer may run concurrently; WriteFrame is serialised internally. +// +// Each WebSocket binary message carries one complete uplink frame (12-byte header + +// payload). The length prefix inside the message is retained as compatibility framing. +type WebSocketConn struct { + ws *websocket.Conn + writeMu sync.Mutex + maxMsgSize int64 + pingInterval time.Duration + pongWait time.Duration + remote string // peer TCP address (or forwarded client IP when set) + status int // HTTP status of the WebSocket handshake (typically 101) + + pingMu sync.Mutex + pingCancel context.CancelFunc + pingDone chan struct{} +} + +// NewWebSocketConn wraps an upgraded WebSocket connection. +func NewWebSocketConn(ws *websocket.Conn, maxMsgSize int) *WebSocketConn { + if maxMsgSize <= 0 { + maxMsgSize = DefaultMaxMessageSize + } + remote := "" + if ws != nil && ws.RemoteAddr() != nil { + remote = ws.RemoteAddr().String() + } + c := &WebSocketConn{ + ws: ws, + maxMsgSize: int64(maxMsgSize), + pingInterval: defaultWSPingInterval, + pongWait: defaultWSPongWait, + remote: remote, + status: http.StatusSwitchingProtocols, + } + ws.SetReadLimit(c.maxMsgSize) + // Read idle must cover the next Ping interval plus Pong wait; using only + // pongWait tears down idle sessions when pingInterval > pongWait. + _ = ws.SetReadDeadline(time.Now().Add(c.pingInterval + c.pongWait)) + ws.SetPongHandler(func(string) error { + _ = ws.SetReadDeadline(time.Now().Add(c.pingInterval + c.pongWait)) + return nil + }) + return c +} + +// RemoteAddr returns the peer address associated with this connection. +func (c *WebSocketConn) RemoteAddr() string { return c.remote } + +// HandshakeStatus returns the HTTP status from the WebSocket upgrade (typically 101). +func (c *WebSocketConn) HandshakeStatus() int { return c.status } + +// SetRemoteAddr overrides the logged peer address (e.g. client IP behind a reverse proxy). +func (c *WebSocketConn) SetRemoteAddr(addr string) { + if addr != "" { + c.remote = addr + } +} + +// StartPingLoop sends WebSocket Ping frames on interval until ctx is cancelled or Close. +func (c *WebSocketConn) StartPingLoop(ctx context.Context) { + c.pingMu.Lock() + if c.pingCancel != nil { + c.pingMu.Unlock() + return + } + pctx, cancel := context.WithCancel(ctx) + c.pingCancel = cancel + c.pingDone = make(chan struct{}) + c.pingMu.Unlock() + + go func() { + defer close(c.pingDone) + t := time.NewTicker(c.pingInterval) + defer t.Stop() + for { + select { + case <-pctx.Done(): + return + case <-t.C: + c.writeMu.Lock() + err := c.ws.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(c.pongWait)) + c.writeMu.Unlock() + if err != nil { + _ = c.ws.Close() + return + } + } + } + }() +} + +func (c *WebSocketConn) stopPing() { + c.pingMu.Lock() + cancel := c.pingCancel + done := c.pingDone + c.pingCancel = nil + c.pingDone = nil + c.pingMu.Unlock() + if cancel != nil { + cancel() + if done != nil { + <-done + } + } +} + +func (c *WebSocketConn) ReadFrame(maxPayload uint32) (FrameHeader, []byte, error) { + for { + msgType, data, err := c.ws.ReadMessage() + if err != nil { + return FrameHeader{}, nil, err + } + if msgType != websocket.BinaryMessage { + continue // ignore text/control handled by library + } + return decodeFrameBytes(data, maxPayload) + } +} + +func (c *WebSocketConn) WriteFrame(h FrameHeader, payload []byte) error { + msg, err := encodeFrameBytes(h, payload) + if err != nil { + return err + } + c.writeMu.Lock() + defer c.writeMu.Unlock() + return c.ws.WriteMessage(websocket.BinaryMessage, msg) +} + +func (c *WebSocketConn) Close() error { + c.stopPing() + c.writeMu.Lock() + defer c.writeMu.Unlock() + _ = c.ws.WriteControl(websocket.CloseMessage, + websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), + time.Now().Add(time.Second)) + return c.ws.Close() +} + +func (c *WebSocketConn) SetReadDeadline(t time.Time) error { + return c.ws.SetReadDeadline(t) +} + +func (c *WebSocketConn) SetWriteDeadline(t time.Time) error { + return c.ws.SetWriteDeadline(t) +} + +// DialWebSocket dials a WSS (or WS) uplink endpoint and returns a WebSocketConn. +// urlStr should be like wss://gateway.example.com/uplink/v1 (path defaults if omitted). +func DialWebSocket(ctx context.Context, urlStr string, tlsConfig *tls.Config, handshakeTimeout time.Duration) (*WebSocketConn, error) { + u, err := url.Parse(urlStr) + if err != nil { + return nil, fmt.Errorf("proxy: invalid uplink URL: %w", err) + } + switch u.Scheme { + case "wss", "ws", "https", "http": + default: + return nil, fmt.Errorf("proxy: unsupported uplink URL scheme %q", u.Scheme) + } + // Normalise http(s) → ws(s) for the dialer. + switch u.Scheme { + case "https": + u.Scheme = "wss" + case "http": + u.Scheme = "ws" + } + if u.Path == "" || u.Path == "/" { + u.Path = UplinkWSPath + } + + if handshakeTimeout <= 0 { + handshakeTimeout = 15 * time.Second + } + dialer := websocket.Dialer{ + HandshakeTimeout: handshakeTimeout, + TLSClientConfig: tlsConfig, + NetDialContext: (&net.Dialer{Timeout: handshakeTimeout}).DialContext, + } + + ws, resp, err := dialer.DialContext(ctx, u.String(), nil) + if err != nil { + if resp != nil { + _ = resp.Body.Close() + } + return nil, err + } + conn := NewWebSocketConn(ws, DefaultMaxMessageSize) + if resp != nil { + conn.status = resp.StatusCode + _ = resp.Body.Close() + } + return conn, nil +} + +// UpgradeWebSocket upgrades an HTTP request to a WebSocketConn. +func UpgradeWebSocket(w http.ResponseWriter, r *http.Request, maxMsgSize int) (*WebSocketConn, error) { + ws, err := wsUpgrader.Upgrade(w, r, nil) + if err != nil { + return nil, err + } + conn := NewWebSocketConn(ws, maxMsgSize) + conn.status = http.StatusSwitchingProtocols + // Prefer forwarded client IP when behind a reverse proxy; fall back to TCP peer. + if ip := ClientIPFromRequest(r); ip != "" { + conn.SetRemoteAddr(ip) + } + return conn, nil +} + +// ClientIPFromRequest extracts the connecting client IP for logging. +// Prefers X-Forwarded-For (first hop) / X-Real-IP when present (reverse-proxy deployments), +// otherwise uses the TCP RemoteAddr host. Never used for authentication. +func ClientIPFromRequest(r *http.Request) string { + if r == nil { + return "" + } + if xff := r.Header.Get("X-Forwarded-For"); xff != "" { + // First address is the original client when proxies append. + if i := strings.IndexByte(xff, ','); i >= 0 { + xff = xff[:i] + } + if ip := strings.TrimSpace(xff); ip != "" { + return ip + } + } + if xri := strings.TrimSpace(r.Header.Get("X-Real-IP")); xri != "" { + return xri + } + host, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + return r.RemoteAddr + } + return host +}