Privatae
119 lines
3.5 KiB
Go
119 lines
3.5 KiB
Go
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()
|
|
}
|
|
}()
|
|
}
|
|
}
|