Go + Vue 管理台：Agent 隧道、域名反代、IP 安全、访客验证与监控大盘。 Gitea Actions CI 多架构镜像；纯 Go SQLite、无 CGO Docker 构建。 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
renjue
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user