bastion 0.53.2

client-go/egress_relay.go

@@ -0,0 +1,118 @@
+package main
+
+import (
+	"context"
+	"crypto/tls"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const egressWorkerCount = 4
+
+// wsConn adapts a worker WebSocket into a net.Conn so the existing handleEgressConn drives
+// it unchanged. The host frames the proxied client stream as binary messages: 'D'+payload =
+// data, 'E' = end-of-stream. Reads surface the host's bytes; writes go back as 'D' frames;
+// Close/CloseWrite send 'E'.
+type wsConn struct {
+	ws   *websocket.Conn
+	rbuf []byte
+	wmu  sync.Mutex
+}
+
+func (c *wsConn) Read(p []byte) (int, error) {
+	for len(c.rbuf) == 0 {
+		mt, data, err := c.ws.ReadMessage()
+		if err != nil {
+			return 0, err
+		}
+		if mt != websocket.BinaryMessage || len(data) == 0 {
+			continue
+		}
+		switch data[0] {
+		case 'D':
+			c.rbuf = data[1:]
+		case 'E':
+			return 0, io.EOF
+		}
+	}
+	n := copy(p, c.rbuf)
+	c.rbuf = c.rbuf[n:]
+	return n, nil
+}
+
+func (c *wsConn) Write(p []byte) (int, error) {
+	c.wmu.Lock()
+	defer c.wmu.Unlock()
+	if err := c.ws.WriteMessage(websocket.BinaryMessage, append([]byte{'D'}, p...)); err != nil {
+		return 0, err
+	}
+	return len(p), nil
+}
+
+// CloseWrite signals end-of-stream to the host (half-close) without tearing the WS down,
+// so the splice's half-close pattern works.
+func (c *wsConn) CloseWrite() error {
+	c.wmu.Lock()
+	defer c.wmu.Unlock()
+	return c.ws.WriteMessage(websocket.BinaryMessage, []byte{'E'})
+}
+
+func (c *wsConn) Close() error {
+	_ = c.CloseWrite()
+	return c.ws.Close()
+}
+
+func (c *wsConn) LocalAddr() net.Addr  { return c.ws.LocalAddr() }
+func (c *wsConn) RemoteAddr() net.Addr { return c.ws.RemoteAddr() }
+func (c *wsConn) SetDeadline(t time.Time) error {
+	_ = c.ws.SetReadDeadline(t)
+	return c.ws.SetWriteDeadline(t)
+}
+func (c *wsConn) SetReadDeadline(t time.Time) error  { return c.ws.SetReadDeadline(t) }
+func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(t) }
+
+// runEgressWorkers opens a pool of worker WebSockets TO the host (the host can't reach the
+// bastion's tunnel IP) and serves one egress job per worker, reconnecting after each. Each
+// worker authenticates with the bastion's bearer, then blocks until the host dispatches a
+// request; handleEgressConn then dials it out the USER's connection (dialOut) and relays.
+//
+// pctx is the pairing's context (run.go): when the anchor revokes a pairing the
+// workers must exit, not re-dial a closed tunnel every 3s for the process's life.
+func runEgressWorkers(pctx context.Context, self string, wsURL string, tlsCfg *tls.Config,
+	dialHost func(ctx context.Context) (net.Conn, error), bearer string,
+	dialOut egressDialer, n int) {
+	dialer := &websocket.Dialer{
+		NetDialContext:   func(ctx context.Context, network, addr string) (net.Conn, error) { return dialHost(ctx) },
+		TLSClientConfig:  tlsCfg,
+		HandshakeTimeout: 20 * time.Second,
+	}
+	for i := 0; i < n; i++ {
+		go func() {
+			for pctx.Err() == nil {
+				ctx, cancel := context.WithTimeout(pctx, 25*time.Second)
+				ws, _, err := dialer.DialContext(ctx, wsURL, http.Header{})
+				cancel()
+				if err != nil {
+					sleepCtx(pctx, 3*time.Second)
+					continue
+				}
+				if err := ws.WriteMessage(websocket.TextMessage, []byte("BEARER "+bearer)); err != nil {
+					_ = ws.Close()
+					sleepCtx(pctx, time.Second)
+					continue
+				}
+				stopKA := armWSKeepalive(ws) // ping/read-deadline: a dead tunnel unblocks the park
+				unpark := context.AfterFunc(pctx, func() { _ = ws.Close() })
+				handleEgressConn(&wsConn{ws: ws}, dialOut, self) // blocks until a job, serves it, then returns
+				unpark()
+				stopKA()
+				_ = ws.Close()
+			}
+		}()
+	}
+}