feat: Wormhole 内网穿透与反向代理网关

Go + Vue 管理台：Agent 隧道、域名反代、IP 安全、访客验证与监控大盘。 Gitea Actions CI 多架构镜像；纯 Go SQLite、无 CGO Docker 构建。 Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-15 04:46:59 +00:00
commit 0d84b07f68
111 changed files with 15770 additions and 0 deletions
@@ -0,0 +1,318 @@
+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)
+}
@@ -0,0 +1,46 @@
+package bridge
+
+import (
+	"crypto/tls"
+	"net"
+	"time"
+
+	"github.com/xtaci/smux"
+)
+
+const (
+	// HeartbeatInterval is the application-level ping interval (agent and server probe).
+	HeartbeatInterval = 30 * time.Second
+	// HeartbeatTimeout evicts a session when no ping/pong for this duration.
+	HeartbeatTimeout = 90 * time.Second
+	// WatchdogInterval is how often the server checks agent heartbeats.
+	WatchdogInterval = 15 * time.Second
+	// ServerProbeAfter triggers a server-initiated ping when agent has been silent.
+	ServerProbeAfter = 45 * time.Second
+)
+
+func SmuxConfig() *smux.Config {
+	return smuxConfig()
+}
+
+func smuxConfig() *smux.Config {
+	cfg := smux.DefaultConfig()
+	cfg.KeepAliveInterval = 10 * time.Second
+	cfg.KeepAliveTimeout = 30 * time.Second
+	return cfg
+}
+
+func SetTCPKeepAlive(conn net.Conn) {
+	for {
+		switch c := conn.(type) {
+		case *net.TCPConn:
+			_ = c.SetKeepAlive(true)
+			_ = c.SetKeepAlivePeriod(30 * time.Second)
+			return
+		case *tls.Conn:
+			conn = c.NetConn()
+		default:
+			return
+		}
+	}
+}
@@ -0,0 +1,74 @@
+package bridge
+
+import (
+	"net"
+	"sync"
+	"time"
+)
+
+// rateLimitedConn limits read/write throughput to bytesPerSec (simple token bucket).
+type rateLimitedConn struct {
+	net.Conn
+	bytesPerSec int64
+	mu          sync.Mutex
+	allowance   float64
+	lastCheck   time.Time
+}
+
+func newRateLimitedConn(conn net.Conn, bytesPerSec int64) net.Conn {
+	if bytesPerSec <= 0 {
+		return conn
+	}
+	return &rateLimitedConn{
+		Conn:        conn,
+		bytesPerSec: bytesPerSec,
+		lastCheck:   time.Now(),
+		allowance:   float64(bytesPerSec),
+	}
+}
+
+func (c *rateLimitedConn) wait(n int) {
+	if c.bytesPerSec <= 0 || n <= 0 {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	now := time.Now()
+	elapsed := now.Sub(c.lastCheck).Seconds()
+	c.lastCheck = now
+	c.allowance += elapsed * float64(c.bytesPerSec)
+	if c.allowance > float64(c.bytesPerSec)*2 {
+		c.allowance = float64(c.bytesPerSec) * 2
+	}
+	for float64(n) > c.allowance {
+		deficit := float64(n) - c.allowance
+		sleep := time.Duration(deficit/float64(c.bytesPerSec)*1e9) * time.Nanosecond
+		if sleep < time.Millisecond {
+			sleep = time.Millisecond
+		}
+		c.mu.Unlock()
+		time.Sleep(sleep)
+		c.mu.Lock()
+		now = time.Now()
+		elapsed = now.Sub(c.lastCheck).Seconds()
+		c.lastCheck = now
+		c.allowance += elapsed * float64(c.bytesPerSec)
+		if c.allowance > float64(c.bytesPerSec)*2 {
+			c.allowance = float64(c.bytesPerSec) * 2
+		}
+	}
+	c.allowance -= float64(n)
+}
+
+func (c *rateLimitedConn) Read(b []byte) (int, error) {
+	n, err := c.Conn.Read(b)
+	if n > 0 {
+		c.wait(n)
+	}
+	return n, err
+}
+
+func (c *rateLimitedConn) Write(b []byte) (int, error) {
+	c.wait(len(b))
+	return c.Conn.Write(b)
+}
@@ -0,0 +1,67 @@
+package bridge
+
+import (
+	"log"
+	"time"
+
+	"github.com/wormhole/wormhole/internal/protocol"
+)
+
+func (b *Bridge) startWatchdog() {
+	go func() {
+		ticker := time.NewTicker(WatchdogInterval)
+		defer ticker.Stop()
+		for range ticker.C {
+			b.runWatchdog()
+		}
+	}()
+}
+
+func (b *Bridge) runWatchdog() {
+	now := time.Now()
+	b.mu.RLock()
+	sessions := make([]*AgentSession, 0, len(b.agents))
+	for _, agent := range b.agents {
+		sessions = append(sessions, agent)
+	}
+	b.mu.RUnlock()
+
+	for _, agent := range sessions {
+		if agent.closed.Load() {
+			continue
+		}
+		lastPing, ok := b.cache.AgentLastPing(agent.ClientID)
+		if !ok {
+			continue
+		}
+		since := now.Sub(lastPing)
+		if since > HeartbeatTimeout {
+			log.Printf("[bridge] client %d heartbeat timeout (%v), closing", agent.ClientID, since.Round(time.Second))
+			agent.Close()
+			continue
+		}
+		if since > ServerProbeAfter {
+			go b.probeAgent(agent)
+		}
+	}
+}
+
+func (b *Bridge) probeAgent(agent *AgentSession) {
+	if agent.closed.Load() {
+		return
+	}
+	stream, err := agent.smux.OpenStream()
+	if err != nil {
+		return
+	}
+	defer stream.Close()
+
+	if err := protocol.WriteMessage(stream, protocol.MsgPing, nil); err != nil {
+		return
+	}
+	msgType, _, err := protocol.ReadMessage(stream)
+	if err != nil || msgType != protocol.MsgPong {
+		return
+	}
+	b.cache.TouchAgent(agent.ClientID)
+}