Wormhole/internal/bridge/bridge.go

package bridge

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/xtaci/smux"
	"github.com/wormhole/wormhole/internal/cache"
	"github.com/wormhole/wormhole/internal/protocol"
	"github.com/wormhole/wormhole/internal/store"
)

type AgentSession struct {
	ClientID    uint
	Version     string
	conn        net.Conn
	smux        *smux.Session
	mu          sync.Mutex
	closed      atomic.Bool
	activeConns atomic.Int32
}

type Bridge struct {
	cache    *cache.Manager
	store    *store.Store
	mu       sync.RWMutex
	agents   map[uint]*AgentSession
	listener net.Listener
}

func New(c *cache.Manager, st *store.Store) *Bridge {
	return &Bridge{
		cache:  c,
		store:  st,
		agents: make(map[uint]*AgentSession),
	}
}

func (b *Bridge) Start(addr string, certFile, keyFile string) error {
	var ln net.Listener
	var err error
	if certFile != "" && keyFile != "" {
		cert, err := tls.LoadX509KeyPair(certFile, keyFile)
		if err != nil {
			return fmt.Errorf("load bridge tls cert: %w", err)
		}
		ln, err = tls.Listen("tcp", addr, &tls.Config{
			Certificates: []tls.Certificate{cert},
			MinVersion:   tls.VersionTLS12,
		})
		if err != nil {
			return err
		}
		log.Printf("[bridge] listening on %s (TLS)", addr)
	} else {
		ln, err = net.Listen("tcp", addr)
		if err != nil {
			return err
		}
		log.Printf("[bridge] listening on %s (plain TCP, configure TLS cert for production)", addr)
	}
	b.listener = ln
	go b.acceptLoop()
	b.startWatchdog()
	return nil
}

func (b *Bridge) Shutdown() error {
	if b.listener != nil {
		return b.listener.Close()
	}
	return nil
}

func (b *Bridge) acceptLoop() {
	for {
		conn, err := b.listener.Accept()
		if err != nil {
			return
		}
		go b.handleConn(conn)
	}
}

func (b *Bridge) validateClient(client *store.Client) error {
	if !client.Status {
		return fmt.Errorf("client disabled")
	}
	if client.ExpireAt != nil && time.Now().After(*client.ExpireAt) {
		return fmt.Errorf("client expired")
	}
	return nil
}

func (b *Bridge) handleConn(conn net.Conn) {
	defer conn.Close()
	SetTCPKeepAlive(conn)
	msgType, body, err := protocol.ReadMessage(conn)
	if err != nil {
		log.Printf("[bridge] read register failed: %v", err)
		return
	}
	if msgType != protocol.MsgRegister {
		return
	}
	reg, err := protocol.ParseRegister(body)
	if err != nil {
		return
	}
	client, ok := b.cache.GetClientByKey(reg.VerifyKey)
	if !ok {
		_ = protocol.WriteMessage(conn, protocol.MsgRegisterAck, protocol.RegisterAckPayload{
			OK:      false,
			Message: "invalid verify key",
		})
		return
	}
	if err := b.validateClient(client); err != nil {
		_ = protocol.WriteMessage(conn, protocol.MsgRegisterAck, protocol.RegisterAckPayload{
			OK:      false,
			Message: err.Error(),
		})
		return
	}
	if err := protocol.WriteMessage(conn, protocol.MsgRegisterAck, protocol.RegisterAckPayload{
		OK:       true,
		ClientID: client.ID,
		Message:  "registered",
	}); err != nil {
		return
	}

	session, err := smux.Client(conn, smuxConfig())
	if err != nil {
		log.Printf("[bridge] smux client failed: %v", err)
		return
	}

	agent := &AgentSession{
		ClientID: client.ID,
		Version:  reg.Version,
		conn:     conn,
		smux:     session,
	}
	b.mu.Lock()
	if old, exists := b.agents[client.ID]; exists {
		old.Close()
	}
	b.agents[client.ID] = agent
	b.mu.Unlock()

	b.cache.SetClientOnline(client.ID, reg.Version)
	log.Printf("[bridge] client %d online (version=%s)", client.ID, reg.Version)

	b.controlLoop(agent)

	b.mu.Lock()
	delete(b.agents, client.ID)
	b.mu.Unlock()
	b.cache.SetClientOffline(client.ID)
	log.Printf("[bridge] client %d offline", client.ID)
}

func (b *Bridge) controlLoop(agent *AgentSession) {
	for {
		stream, err := agent.smux.AcceptStream()
		if err != nil {
			agent.Close()
			return
		}
		go b.handleControlStream(agent, stream)
	}
}

func (b *Bridge) handleControlStream(agent *AgentSession, stream *smux.Stream) {
	defer stream.Close()
	msgType, body, err := protocol.ReadMessage(stream)
	if err != nil {
		return
	}
	switch msgType {
	case protocol.MsgPing:
		b.cache.TouchAgent(agent.ClientID)
		_ = protocol.WriteMessage(stream, protocol.MsgPong, nil)
	case protocol.MsgClose:
		agent.Close()
	}
	_ = body
}

func (a *AgentSession) Close() {
	if a.closed.Swap(true) {
		return
	}
	if a.smux != nil {
		_ = a.smux.Close()
	}
	if a.conn != nil {
		_ = a.conn.Close()
	}
}

func (b *Bridge) OpenStream(clientID uint, meta *protocol.LinkMeta) (net.Conn, error) {
	client, err := b.store.GetClientByID(clientID)
	if err != nil {
		return nil, fmt.Errorf("client %d not found", clientID)
	}
	if err := b.validateClient(client); err != nil {
		return nil, err
	}

	b.mu.RLock()
	agent, ok := b.agents[clientID]
	b.mu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("client %d offline", clientID)
	}
	if agent.closed.Load() {
		return nil, fmt.Errorf("client %d session closed", clientID)
	}
	if client.MaxConn > 0 && int(agent.activeConns.Load()) >= client.MaxConn {
		return nil, fmt.Errorf("client %d max connections reached", clientID)
	}

	stream, err := agent.smux.OpenStream()
	if err != nil {
		return nil, err
	}
	if err := protocol.WriteMessage(stream, protocol.MsgOpenStream, meta); err != nil {
		stream.Close()
		return nil, err
	}
	msgType, _, err := protocol.ReadMessage(stream)
	if err != nil {
		stream.Close()
		return nil, err
	}
	if msgType != protocol.MsgStreamReady {
		stream.Close()
		return nil, fmt.Errorf("stream not ready")
	}

	agent.activeConns.Add(1)
	var wrapped net.Conn = &trackedConn{
		Conn:   stream,
		onClose: func() { agent.activeConns.Add(-1) },
	}
	if client.RateLimit > 0 {
		wrapped = newRateLimitedConn(wrapped, client.RateLimit*1024)
	}
	return wrapped, nil
}

func (b *Bridge) IsOnline(clientID uint) bool {
	b.mu.RLock()
	defer b.mu.RUnlock()
	_, ok := b.agents[clientID]
	return ok
}

func (b *Bridge) OnlineCount() int {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return len(b.agents)
}

type trackedConn struct {
	net.Conn
	onClose func()
	once    sync.Once
}

func (c *trackedConn) Close() error {
	c.once.Do(c.onClose)
	return c.Conn.Close()
}

func CopyBoth(a, b net.Conn) (inlet, export int64) {
	var inVal, outVal int64
	done := make(chan struct{}, 2)
	go func() {
		n, _ := io.Copy(b, a)
		atomic.AddInt64(&outVal, n)
		done <- struct{}{}
	}()
	go func() {
		n, _ := io.Copy(a, b)
		atomic.AddInt64(&inVal, n)
		done <- struct{}{}
	}()
	<-done
	<-done
	return atomic.LoadInt64(&inVal), atomic.LoadInt64(&outVal)
}

func CopyBothIO(a io.ReadWriteCloser, b io.ReadWriteCloser) (inlet, export int64) {
	var inVal, outVal int64
	done := make(chan struct{}, 2)
	go func() {
		n, _ := io.Copy(b, a)
		atomic.AddInt64(&outVal, n)
		done <- struct{}{}
	}()
	go func() {
		n, _ := io.Copy(a, b)
		atomic.AddInt64(&inVal, n)
		done <- struct{}{}
	}()
	<-done
	<-done
	return atomic.LoadInt64(&inVal), atomic.LoadInt64(&outVal)
}