223 lines
6.5 KiB
Go
223 lines
6.5 KiB
Go
// Package agent 轻量穿透被控端（zonat client 的裁剪版），设计为**可嵌入公开包**：
|
||
// 每条隧道自带节点地址与登录令牌，由独立连接完成
|
||
// 登录节点 → smux 会话 → 控制流注册自身 → 数据流拨号本地 TCP 目标，
|
||
// 因此一个 agent 的多条隧道可以分布在**多个不同节点**。
|
||
// Agent 本身只是隧道的注册表与生命周期管理器，不持有网络连接、不做登录。
|
||
// 相比 zonat client 裁掉：KCP 拨号、UDP 代理、普通(非复用)连接模式、
|
||
// 用户级通道分发、服务安装等本机管理逻辑。
|
||
// 永久隧道（TTLSec=0）断线各自自动重连重注册；临时隧道（TTLSec>0）
|
||
// 一次性，连接断开或节点 TTL 过期推送后即消失，与 zonat 临时通道语义一致。
|
||
//
|
||
// 日志：包内不依赖任何具体日志实现，宿主项目通过注入 Logger 接口
|
||
// 把 agent 日志纳入自己的日志体系（格式/级别/输出/文件路径归一）；
|
||
// 不注入时使用仅标准库的兜底实现（stderr，Info 级）。
|
||
// 注意：日志中的 file:line 归属调用点。宿主以 import 方式引用本包并加
|
||
// -trimpath 构建时，路径为模块形式 zonat@版本/agent/agent.go（跨机器稳定）；
|
||
// 以整目录拷贝方式放进宿主项目时，路径在宿主项目内，宿主现有 rootDir
|
||
// 削前缀规则原样生效。两种方式都不需要在包内做任何路径处理。
|
||
package agent
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"log"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
defaultBackoff = 5 * time.Second
|
||
defaultPing = 30 * time.Second
|
||
)
|
||
|
||
// Logger agent 日志接口。logrus.Logger 的 Debugf/Infof/Warnf/Errorf 签名
|
||
// 与此完全一致，可直接注入；slog 等其它实现写一个几行的适配器即可。
|
||
type Logger interface {
|
||
Debugf(format string, args ...any)
|
||
Infof(format string, args ...any)
|
||
Warnf(format string, args ...any)
|
||
Errorf(format string, args ...any)
|
||
}
|
||
|
||
// DefaultLogger 未注入时的兜底日志：标准错误，Info 级，仅标准库实现。
|
||
type DefaultLogger struct{}
|
||
|
||
func (DefaultLogger) Debugf(string, ...any) {}
|
||
|
||
func (DefaultLogger) Infof(format string, args ...any) {
|
||
log.Printf("INFO %s", fmt.Sprintf(format, args...))
|
||
}
|
||
|
||
func (DefaultLogger) Warnf(format string, args ...any) {
|
||
log.Printf("WARN %s", fmt.Sprintf(format, args...))
|
||
}
|
||
|
||
func (DefaultLogger) Errorf(format string, args ...any) {
|
||
log.Printf("ERROR %s", fmt.Sprintf(format, args...))
|
||
}
|
||
|
||
// Agent 被控端：隧道注册表。每条 Tunnel 携带自己的 NodeAddr+Token，
|
||
// 由独立 worker（tunnelConn）连接登录各自节点；Agent 不再有集中连接与登录。
|
||
type Agent struct {
|
||
AgentId string
|
||
Backoff time.Duration // 永久隧道重连退避（默认 5s）
|
||
PingEvery time.Duration // 各隧道控制流心跳间隔（默认 30s）
|
||
|
||
// Logger 日志注入点。需要纳入宿主统一日志体系时设置
|
||
// （logrus 直接注入 logrus.StandardLogger() 即可）；未设置用 DefaultLogger。
|
||
Logger Logger
|
||
|
||
// OnTunnelPort 隧道端口确定/变更回调：port>0 已注册，port<0 已关闭。
|
||
// 节点随机分配端口时，上层靠它拿到实际端口。多隧道场景下可能从
|
||
// 不同 worker 并发回调，宿主需自行保证线程安全。
|
||
OnTunnelPort func(id string, port int)
|
||
|
||
mu sync.Mutex
|
||
closed bool
|
||
tunnels map[string]*tunnelConn
|
||
|
||
done chan struct{} // Run 的阻塞点，Close 时关闭
|
||
}
|
||
|
||
// New 创建被控端。节点地址与令牌随各隧道下发（Tunnel.NodeAddr/Token）。
|
||
func New(agentId string) *Agent {
|
||
return &Agent{
|
||
AgentId: agentId,
|
||
Backoff: defaultBackoff,
|
||
PingEvery: defaultPing,
|
||
Logger: DefaultLogger{},
|
||
tunnels: make(map[string]*tunnelConn),
|
||
done: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// logger 取当前日志实现（零值 Agent 未注入时兜底）
|
||
func (a *Agent) logger() Logger {
|
||
a.mu.Lock()
|
||
l := a.Logger
|
||
a.mu.Unlock()
|
||
if l == nil {
|
||
return DefaultLogger{}
|
||
}
|
||
return l
|
||
}
|
||
|
||
// Run 阻塞直到 Close。隧道的连接在 RegisterTunnel 时即自行驱动，
|
||
// 这里仅保留"go a.Run()"式宿主用法的阻塞点。
|
||
func (a *Agent) Run() error {
|
||
<-a.done
|
||
return nil
|
||
}
|
||
|
||
// Close 停止 agent：断开并停止全部隧道 worker
|
||
func (a *Agent) Close() {
|
||
a.mu.Lock()
|
||
if a.closed {
|
||
a.mu.Unlock()
|
||
return
|
||
}
|
||
a.closed = true
|
||
workers := make([]*tunnelConn, 0, len(a.tunnels))
|
||
for _, w := range a.tunnels {
|
||
workers = append(workers, w)
|
||
}
|
||
close(a.done)
|
||
a.mu.Unlock()
|
||
for _, w := range workers {
|
||
w.stop()
|
||
}
|
||
}
|
||
|
||
func (a *Agent) isClosed() bool {
|
||
a.mu.Lock()
|
||
defer a.mu.Unlock()
|
||
return a.closed
|
||
}
|
||
|
||
// Connected 指定隧道是否已连接其节点
|
||
func (a *Agent) Connected(id string) bool {
|
||
a.mu.Lock()
|
||
w := a.tunnels[id]
|
||
a.mu.Unlock()
|
||
if w == nil {
|
||
return false
|
||
}
|
||
w.mu.Lock()
|
||
s := w.session
|
||
w.mu.Unlock()
|
||
return s != nil && !s.IsClosed()
|
||
}
|
||
|
||
// RegisterTunnel 注册隧道：同步等待首次"连接+登录+注册"完成并返回节点分配端口，
|
||
// 之后由后台 worker 维持——永久隧道（TTLSec=0）断线自动重连重注册，
|
||
// 临时隧道（TTLSec>0）一次性。同 ID 重复注册视为更新配置（停旧 worker 起新）。
|
||
func (a *Agent) RegisterTunnel(t Tunnel) (int, error) {
|
||
if t.Id == "" {
|
||
return 0, errors.New("隧道ID不能为空")
|
||
}
|
||
if t.NodeAddr == "" || t.Token == "" {
|
||
return 0, fmt.Errorf("隧道缺少节点地址或令牌 id=%s", t.Id)
|
||
}
|
||
a.mu.Lock()
|
||
if a.closed {
|
||
a.mu.Unlock()
|
||
return 0, errors.New("agent 已关闭")
|
||
}
|
||
old := a.tunnels[t.Id]
|
||
w := newTunnelConn(a, t)
|
||
a.tunnels[t.Id] = w
|
||
a.mu.Unlock()
|
||
if old != nil {
|
||
old.stop()
|
||
}
|
||
|
||
go w.run()
|
||
select {
|
||
case <-w.first:
|
||
return w.firstPort, w.firstErr
|
||
case <-time.After(registerWait):
|
||
return 0, fmt.Errorf("注册超时(节点 %s)，永久隧道将在后台继续重试", t.NodeAddr)
|
||
}
|
||
}
|
||
|
||
// UnregisterTunnel 注销隧道：尽力经该隧道的控制流通知节点删除，随后停止本地 worker
|
||
func (a *Agent) UnregisterTunnel(id string) error {
|
||
a.mu.Lock()
|
||
w := a.tunnels[id]
|
||
a.mu.Unlock()
|
||
if w == nil {
|
||
return errors.New("隧道不存在:" + id)
|
||
}
|
||
err := w.unregister()
|
||
w.stop()
|
||
a.removeTunnelState(id)
|
||
return err
|
||
}
|
||
|
||
// notifyPort 隧道注册成功（含重连后重注册）时回调
|
||
func (a *Agent) notifyPort(id string, port int) {
|
||
a.mu.Lock()
|
||
cb := a.OnTunnelPort
|
||
a.mu.Unlock()
|
||
if cb != nil {
|
||
cb(id, port)
|
||
}
|
||
}
|
||
|
||
func (a *Agent) removeTunnelState(id string) {
|
||
a.mu.Lock()
|
||
delete(a.tunnels, id)
|
||
cb := a.OnTunnelPort
|
||
a.mu.Unlock()
|
||
if cb != nil {
|
||
cb(id, -1)
|
||
}
|
||
}
|
||
|
||
func errText(err error) string {
|
||
if err == nil {
|
||
return ""
|
||
}
|
||
return err.Error()
|
||
}
|