tunnel/internal/client/bridge.go

package client

import (
	"context"
	"crypto/tls"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coder/websocket"
)

const (
	bridgeSubprotocol = "opentunnel"
	maxPayload        = 32768
	maxConnsPerBridge = 256
)

var (
	errFatalAttach  = errors.New("fatal attach error")
	errBackpressure = errors.New("bridge data queue full")
)

type attachOut struct {
	Type      string   `json:"type"`
	Token     string   `json:"token"`
	Transport string   `json:"transport"`
	Routes    []string `json:"routes"`
	Client    struct {
		Version  string `json:"version"`
		MaxConns int    `json:"max_conns"`
	} `json:"client"`
}

type serverControl struct {
	Type          string   `json:"type"`
	Session       string   `json:"session,omitempty"`
	Routes        []string `json:"routes,omitempty"`
	HeartbeatMs   int      `json:"heartbeat_ms,omitempty"`
	IdleTimeoutMs int      `json:"idle_timeout_ms,omitempty"`
	Conn          *uint32  `json:"conn,omitempty"`
	Peer          string   `json:"peer,omitempty"`
	SNI           string   `json:"sni,omitempty"`
	ALPN          string   `json:"alpn,omitempty"`
	Code          string   `json:"code,omitempty"`
	TimeSent      int64    `json:"time_sent,omitempty"`
	Reason        string   `json:"reason,omitempty"`
}

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

type session struct {
	svc       *Service
	cfg       *Config
	log       *log.Logger
	ws        *websocket.Conn
	routes    map[string]Route
	heartbeat time.Duration
	idle      time.Duration

	control chan outbound
	data    chan outbound
	done    chan struct{}

	closeOnce sync.Once
	mu        sync.Mutex
	conns     map[uint32]*bridgeConn
	lastRecv  atomic.Int64
}

func (s *Service) connect(ctx context.Context, cfg *Config) (*session, error) {
	wsURL, err := s.API.WebSocketURL(cfg.TunnelID)
	if err != nil {
		return nil, err
	}
	dialCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
	ws, _, err := websocket.Dial(dialCtx, wsURL, &websocket.DialOptions{
		HTTPHeader:      http.Header{"Authorization": []string{"Bearer " + cfg.Token}},
		Subprotocols:    []string{bridgeSubprotocol},
		CompressionMode: websocket.CompressionDisabled,
	})
	cancel()
	if err != nil {
		return nil, fmt.Errorf("dial bridge: %w", err)
	}
	ws.SetReadLimit(64 * 1024)

	routes := make([]string, 0, len(cfg.Routes))
	targets := make(map[string]Route, len(cfg.Routes))
	for name, route := range cfg.Routes {
		routes = append(routes, name)
		targets[name] = route
	}
	sort.Strings(routes)
	if len(routes) > 128 {
		ws.Close(websocket.StatusPolicyViolation, "too many routes")
		return nil, fmt.Errorf("too many routes (%d, max 128)", len(routes))
	}

	var attach attachOut
	attach.Type = "attach"
	attach.Token = cfg.Token
	attach.Transport = "ws"
	attach.Routes = routes
	attach.Client.Version = s.Version
	attach.Client.MaxConns = maxConnsPerBridge
	if err := writeJSON(ctx, ws, attach); err != nil {
		ws.Close(websocket.StatusInternalError, "attach write")
		return nil, err
	}

	readCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	typ, raw, err := ws.Read(readCtx)
	cancel()
	if err != nil {
		ws.Close(websocket.StatusInternalError, "attach read")
		return nil, fmt.Errorf("attach: %w", err)
	}
	if typ != websocket.MessageText {
		ws.Close(websocket.StatusProtocolError, "expected attach answer")
		return nil, errors.New("attach: unexpected binary answer")
	}
	var answer serverControl
	if err := json.Unmarshal(raw, &answer); err != nil {
		ws.Close(websocket.StatusProtocolError, "bad attach answer")
		return nil, fmt.Errorf("attach answer: %w", err)
	}
	switch answer.Type {
	case "attached":
	case "attach_error":
		ws.Close(websocket.StatusPolicyViolation, answer.Code)
		if answer.Code == "bad_token" || answer.Code == "cert_not_ready" {
			return nil, fmt.Errorf("%w: %s", errFatalAttach, answer.Code)
		}
		return nil, fmt.Errorf("attach: %s", answer.Code)
	default:
		ws.Close(websocket.StatusProtocolError, "expected attached")
		return nil, fmt.Errorf("attach: unexpected message %q", answer.Type)
	}
	ws.SetReadLimit(4 + maxPayload)

	heartbeat := time.Duration(answer.HeartbeatMs) * time.Millisecond
	if heartbeat <= 0 {
		heartbeat = 15 * time.Second
	}
	idle := time.Duration(answer.IdleTimeoutMs) * time.Millisecond
	if idle <= 0 {
		idle = 45 * time.Second
	}
	sess := &session{
		svc:       s,
		cfg:       cfg,
		log:       s.Log,
		ws:        ws,
		routes:    targets,
		heartbeat: heartbeat,
		idle:      idle,
		control:   make(chan outbound, 64),
		data:      make(chan outbound, 256),
		done:      make(chan struct{}),
		conns:     map[uint32]*bridgeConn{},
	}
	sess.touch()
	go sess.writeLoop()
	go sess.heartbeatLoop()
	go sess.readLoop()
	s.Log.Printf("attached: session %s, %d route(s), heartbeat %s", answer.Session, len(routes), heartbeat)
	return sess, nil
}

func writeJSON(ctx context.Context, ws *websocket.Conn, value any) error {
	data, err := json.Marshal(value)
	if err != nil {
		return err
	}
	writeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	return ws.Write(writeCtx, websocket.MessageText, data)
}

func (s *session) touch() { s.lastRecv.Store(time.Now().UnixMilli()) }

func (s *session) readLoop() {
	ctx := context.Background()
	for {
		typ, raw, err := s.ws.Read(ctx)
		if err != nil {
			s.close("read: " + err.Error())
			return
		}
		s.touch()
		if typ == websocket.MessageText {
			s.handleControl(raw)
			continue
		}
		if len(raw) < 4 {
			continue
		}
		id := binary.BigEndian.Uint32(raw[:4])
		s.mu.Lock()
		bc := s.conns[id]
		s.mu.Unlock()
		if bc == nil {
			continue
		}
		payload := raw[4:]
		select {
		case bc.incoming <- payload:
		case <-bc.done:
		default:
			// The local side is not draining; drop this connection rather
			// than blocking the bridge read loop.
			s.resetConn(id, "upstream_io_error")
		}
	}
}

func (s *session) handleControl(raw []byte) {
	var msg serverControl
	if err := json.Unmarshal(raw, &msg); err != nil {
		return
	}
	switch msg.Type {
	case "ping":
		s.sendControl(serverControl{Type: "pong", TimeSent: msg.TimeSent})
	case "pong":
		// Any received message already refreshed lastRecv.
	case "open":
		if msg.Conn != nil {
			s.handleOpen(*msg.Conn, msg.Peer, msg.SNI, msg.ALPN)
		}
	case "end":
		if msg.Conn != nil {
			s.endFromServer(*msg.Conn)
		}
	case "reset":
		if msg.Conn != nil {
			s.resetFromServer(*msg.Conn)
		}
	case "drain":
		s.log.Printf("server drained the bridge: %s", msg.Reason)
		s.close("server drain")
	case "attached", "attach_error":
		// Handled during connect.
	}
}

func (s *session) handleOpen(id uint32, peer, sni, alpn string) {
	route := routeForSNI(s.cfg.Hostname, sni)
	if route == "" {
		s.sendReset(id, "unknown_route")
		return
	}
	target, ok := s.routes[route]
	if !ok {
		s.sendReset(id, "unknown_route")
		return
	}
	s.mu.Lock()
	if len(s.conns) >= maxConnsPerBridge {
		s.mu.Unlock()
		s.sendReset(id, "too_many_connections")
		return
	}
	bc := newBridgeConn(id, s)
	bc.peer, bc.sni, bc.alpn = peer, sni, alpn
	s.conns[id] = bc
	s.mu.Unlock()
	go s.serveConn(bc, target, peer, sni, alpn)
}

func (s *session) serveConn(bc *bridgeConn, route Route, peer, sni, alpn string) {
	defer func() {
		s.mu.Lock()
		if s.conns[bc.id] == bc {
			delete(s.conns, bc.id)
		}
		s.mu.Unlock()
		bc.fullClose()
	}()

	tlsConn := tls.Server(bc, s.svc.TLSConfig())
	handshakeCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	err := tlsConn.HandshakeContext(handshakeCtx)
	cancel()
	if err != nil {
		s.log.Printf("tls handshake for %s: %v", sni, err)
		s.sendReset(bc.id, "upstream_io_error")
		return
	}

	target, err := net.DialTimeout("tcp", route.Target, 10*time.Second)
	if err != nil {
		s.log.Printf("connect %s: %v", route.Target, err)
		s.sendReset(bc.id, "upstream_connect_failed")
		return
	}
	defer target.Close()
	if route.ProxyProtocol != "" {
		if err := writeProxyHeader(target, route.ProxyProtocol, peer, sni); err != nil {
			s.log.Printf("proxy protocol for %s: %v", sni, err)
			s.sendReset(bc.id, "upstream_io_error")
			return
		}
	}

	// visitor -> target. On TLS EOF, half-close the target.
	visitorDone := make(chan error, 1)
	go func() {
		_, err := io.Copy(target, tlsConn)
		if tcp, ok := target.(*net.TCPConn); ok {
			tcp.CloseWrite()
		}
		visitorDone <- err
	}()
	// target -> visitor. When the target ends, send TLS close_notify and an
	// `end` control so the server half-closes the visitor.
	targetDone := make(chan error, 1)
	go func() {
		_, err := io.Copy(tlsConn, target)
		_ = tlsConn.CloseWrite()
		s.sendEnd(bc.id)
		targetDone <- err
	}()

	var firstErr error
	for done := 0; done < 2; {
		select {
		case err := <-visitorDone:
			done++
			if err != nil && !isClosedError(err) {
				firstErr = err
			}
			visitorDone = nil
		case err := <-targetDone:
			done++
			if err != nil && !isClosedError(err) {
				firstErr = err
			}
			targetDone = nil
		case <-bc.done:
			s.sendReset(bc.id, "io_error")
			return
		case <-time.After(60 * time.Second):
			s.sendReset(bc.id, "idle_timeout")
			return
		}
	}
	if firstErr != nil {
		s.sendReset(bc.id, "upstream_io_error")
	}
}

func (s *session) endFromServer(id uint32) {
	s.mu.Lock()
	bc := s.conns[id]
	s.mu.Unlock()
	if bc != nil {
		bc.remoteEnd()
	}
}

func (s *session) resetFromServer(id uint32) {
	s.mu.Lock()
	bc := s.conns[id]
	delete(s.conns, id)
	s.mu.Unlock()
	if bc != nil {
		bc.remoteReset()
	}
}

func (s *session) resetConn(id uint32, code string) {
	s.mu.Lock()
	bc := s.conns[id]
	delete(s.conns, id)
	s.mu.Unlock()
	s.sendReset(id, code)
	if bc != nil {
		bc.remoteReset()
	}
}

func (s *session) sendReset(id uint32, code string) {
	s.sendControl(serverControl{Type: "reset", Conn: &id, Code: code})
}

func (s *session) sendControl(msg serverControl) {
	data, err := json.Marshal(msg)
	if err != nil {
		return
	}
	select {
	case s.control <- outbound{typ: websocket.MessageText, data: data}:
	case <-s.done:
	default:
		s.close("control queue full")
	}
}

func (s *session) sendData(frame []byte) error {
	select {
	case <-s.done:
		return io.ErrClosedPipe
	default:
	}
	select {
	case s.data <- outbound{typ: websocket.MessageBinary, data: frame}:
		return nil
	case <-s.done:
		return io.ErrClosedPipe
	default:
		return errBackpressure
	}
}

// sendEnd queues a half-close after any already-queued data frames. It must
// not use the priority control queue: an `end` overtaking the final response
// bytes truncates the visitor's connection.
func (s *session) sendEnd(id uint32) {
	data, err := json.Marshal(serverControl{Type: "end", Conn: &id})
	if err != nil {
		return
	}
	select {
	case s.data <- outbound{typ: websocket.MessageText, data: data}:
	case <-s.done:
	default:
		s.close("data queue full")
	}
}

func (s *session) writeLoop() {
	for {
		// Give control messages priority over queued visitor data.
		select {
		case <-s.done:
			return
		default:
		}
		select {
		case <-s.done:
			return
		case msg := <-s.control:
			if err := s.write(msg); err != nil {
				s.close("write: " + err.Error())
				return
			}
			continue
		default:
		}
		select {
		case <-s.done:
			return
		case msg := <-s.control:
			if err := s.write(msg); err != nil {
				s.close("write: " + err.Error())
				return
			}
		case frame := <-s.data:
			if err := s.write(frame); err != nil {
				s.close("write: " + err.Error())
				return
			}
		}
	}
}

func (s *session) write(msg outbound) error {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	return s.ws.Write(ctx, msg.typ, msg.data)
}

func (s *session) heartbeatLoop() {
	ticker := time.NewTicker(s.heartbeat)
	defer ticker.Stop()
	for {
		select {
		case <-s.done:
			return
		case now := <-ticker.C:
			if now.UnixMilli()-s.lastRecv.Load() > s.idle.Milliseconds() {
				s.close("idle timeout")
				return
			}
			s.sendControl(serverControl{Type: "ping", TimeSent: now.UnixMilli()})
		}
	}
}

func (s *session) close(reason string) {
	s.closeOnce.Do(func() {
		close(s.done)
		s.ws.Close(websocket.StatusGoingAway, reason)
		s.mu.Lock()
		conns := make([]*bridgeConn, 0, len(s.conns))
		for _, bc := range s.conns {
			conns = append(conns, bc)
		}
		s.conns = map[uint32]*bridgeConn{}
		s.mu.Unlock()
		for _, bc := range conns {
			bc.remoteReset()
		}
	})
}

// bridgeConn is a net.Conn fed by bridge frames.
type bridgeConn struct {
	id       uint32
	session  *session
	incoming chan []byte
	buf      []byte
	peer     string
	sni      string
	alpn     string

	mu      sync.Mutex
	readErr error
	readEnd chan struct{}
	done    chan struct{}
	once    sync.Once
}

func newBridgeConn(id uint32, s *session) *bridgeConn {
	return &bridgeConn{
		id:       id,
		session:  s,
		incoming: make(chan []byte, 64),
		readEnd:  make(chan struct{}),
		done:     make(chan struct{}),
	}
}

func (c *bridgeConn) Read(p []byte) (int, error) {
	for len(c.buf) == 0 {
		select {
		case buf := <-c.incoming:
			c.buf = buf
		case <-c.readEnd:
			c.mu.Lock()
			err := c.readErr
			c.mu.Unlock()
			if err != nil {
				return 0, err
			}
			return 0, io.EOF
		}
	}
	n := copy(p, c.buf)
	c.buf = c.buf[n:]
	return n, nil
}

func (c *bridgeConn) Write(p []byte) (int, error) {
	select {
	case <-c.done:
		return 0, io.ErrClosedPipe
	default:
	}
	frame := make([]byte, 4+len(p))
	binary.BigEndian.PutUint32(frame, c.id)
	copy(frame[4:], p)
	if err := c.session.sendData(frame); err != nil {
		c.session.sendReset(c.id, "upstream_io_error")
		return 0, err
	}
	return len(p), nil
}

func (c *bridgeConn) Close() error {
	c.remoteReset()
	return nil
}

func (c *bridgeConn) LocalAddr() net.Addr              { return bridgeAddr("bridge") }
func (c *bridgeConn) RemoteAddr() net.Addr             { return bridgeAddr(c.peer) }
func (c *bridgeConn) SetDeadline(time.Time) error      { return nil }
func (c *bridgeConn) SetReadDeadline(time.Time) error  { return nil }
func (c *bridgeConn) SetWriteDeadline(time.Time) error { return nil }

func (c *bridgeConn) remoteEnd() {
	c.mu.Lock()
	select {
	case <-c.readEnd:
	default:
		close(c.readEnd)
	}
	c.mu.Unlock()
}

func (c *bridgeConn) remoteReset() {
	c.mu.Lock()
	if c.readErr == nil {
		c.readErr = io.ErrClosedPipe
	}
	select {
	case <-c.readEnd:
	default:
		close(c.readEnd)
	}
	c.mu.Unlock()
	c.once.Do(func() { close(c.done) })
}

func (c *bridgeConn) fullClose() {
	c.remoteReset()
}

type bridgeAddr string

func (a bridgeAddr) Network() string { return "bridge" }
func (a bridgeAddr) String() string  { return string(a) }

func routeForSNI(hostname, sni string) string {
	sni = strings.ToLower(strings.TrimSuffix(sni, "."))
	hostname = strings.ToLower(hostname)
	if sni == hostname {
		return "@"
	}
	suffix := "." + hostname
	if !strings.HasSuffix(sni, suffix) {
		return ""
	}
	prefix := strings.TrimSuffix(sni, suffix)
	if strings.Contains(prefix, ".") || !ValidRouteName(prefix) {
		return ""
	}
	return prefix
}

func isClosedError(err error) bool {
	return err == nil || errors.Is(err, io.EOF) || errors.Is(err, io.ErrClosedPipe) || errors.Is(err, net.ErrClosed)
}