305 lines
6.9 KiB
Go
305 lines
6.9 KiB
Go
// tunnel.go 一条隧道的节点连接 worker：每条隧道独立完成
|
||
// 连接→登录（携带自身 Token）→smux→控制流注册自身→数据流服务，
|
||
// 因此天然支持不同隧道连接不同节点。
|
||
// 永久隧道断线按 Backoff 重连重注册；临时隧道一次性退出。
|
||
package agent
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/xtaci/smux"
|
||
)
|
||
|
||
const (
|
||
loginTimeout = 10 * time.Second // 登录握手超时（含拨号）
|
||
registerWait = 30 * time.Second // RegisterTunnel 等待首次注册完成的上限
|
||
)
|
||
|
||
// errTunnelClosed 节点推送 CmdTunnelClosed：隧道已被节点删除（TTL 过期等）
|
||
var errTunnelClosed = errors.New("隧道已被节点关闭")
|
||
|
||
// tunnelConn 一条隧道的全部运行态：期望配置 + 到该隧道节点的桥接连接
|
||
type tunnelConn struct {
|
||
ag *Agent
|
||
t Tunnel
|
||
|
||
stopCh chan struct{}
|
||
stopOnce sync.Once
|
||
|
||
firstOnce sync.Once
|
||
first chan struct{} // 首次"连接+登录+注册"尝试结束后关闭
|
||
firstPort int
|
||
firstErr error
|
||
|
||
everRegistered bool // 本次注册成功过至少一次（run 协程独占读写）
|
||
|
||
mu sync.Mutex
|
||
conn net.Conn
|
||
session *smux.Session
|
||
control net.Conn
|
||
|
||
ctrlMu sync.Mutex // 控制流"写请求+读应答"串行化
|
||
}
|
||
|
||
func newTunnelConn(ag *Agent, t Tunnel) *tunnelConn {
|
||
return &tunnelConn{
|
||
ag: ag,
|
||
t: t,
|
||
stopCh: make(chan struct{}),
|
||
first: make(chan struct{}),
|
||
}
|
||
}
|
||
|
||
// markFirst 记录首次注册尝试的结果；channel close 保证读取方看到字段写入
|
||
func (w *tunnelConn) markFirst(port int, err error) {
|
||
w.firstOnce.Do(func() {
|
||
w.firstPort, w.firstErr = port, err
|
||
close(w.first)
|
||
})
|
||
}
|
||
|
||
// run 重连主循环，阻塞直到停止/一次性结束
|
||
func (w *tunnelConn) run() {
|
||
for {
|
||
port, err := w.connect()
|
||
w.markFirst(port, err)
|
||
if err == nil {
|
||
err = w.serve()
|
||
}
|
||
w.teardown()
|
||
if w.isStopped() || w.ag.isClosed() {
|
||
return
|
||
}
|
||
if errors.Is(err, errTunnelClosed) {
|
||
w.ag.removeTunnelState(w.t.Id)
|
||
return
|
||
}
|
||
if w.t.TTLSec > 0 {
|
||
// 临时隧道一次性：连接断开即结束（节点侧随会话移除）
|
||
if w.everRegistered {
|
||
w.ag.removeTunnelState(w.t.Id)
|
||
}
|
||
return
|
||
}
|
||
w.ag.logger().Warnf("agent 隧道连接断开 id=%s node=%s err=%s backoff=%s",
|
||
w.t.Id, w.t.NodeAddr, errText(err), w.ag.Backoff)
|
||
select {
|
||
case <-time.After(w.ag.Backoff):
|
||
case <-w.stopCh:
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
// connect 连接并登录该隧道的节点，经控制流注册自身，返回节点分配端口
|
||
func (w *tunnelConn) connect() (int, error) {
|
||
conn, err := net.DialTimeout("tcp", w.t.NodeAddr, loginTimeout)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
|
||
// 登录（同 zonat：CmdLogin 帧 + Ret 应答）
|
||
lf := NewFrame(FrameVersion, CmdLogin, 0)
|
||
if err := lf.Marshal(Login{AgentId: w.ag.AgentId, Token: w.t.Token, Version: Version}); err != nil {
|
||
conn.Close()
|
||
return 0, err
|
||
}
|
||
if err := WriteFrame(conn, lf); err != nil {
|
||
conn.Close()
|
||
return 0, err
|
||
}
|
||
conn.SetReadDeadline(time.Now().Add(loginTimeout))
|
||
rf, err := ReadFrame(conn)
|
||
if err != nil {
|
||
conn.Close()
|
||
return 0, err
|
||
}
|
||
ret := Ret{}
|
||
if err := rf.Unmarshal(&ret); err != nil {
|
||
conn.Close()
|
||
return 0, err
|
||
}
|
||
if ret.Code != 0 {
|
||
conn.Close()
|
||
return 0, fmt.Errorf("登录失败:%s", ret.Msg)
|
||
}
|
||
conn.SetReadDeadline(time.Time{})
|
||
|
||
sess, err := smux.Client(conn, SmuxConfig())
|
||
if err != nil {
|
||
conn.Close()
|
||
return 0, err
|
||
}
|
||
control, err := sess.OpenStream()
|
||
if err != nil {
|
||
sess.Close()
|
||
return 0, err
|
||
}
|
||
|
||
w.mu.Lock()
|
||
w.conn, w.session, w.control = conn, sess, control
|
||
w.mu.Unlock()
|
||
|
||
// 注册自身；凭据只用于登录，不随注册/数据帧回发
|
||
wire := w.t
|
||
wire.NodeAddr, wire.Token = "", ""
|
||
f, err := w.callControl(func(control net.Conn) error {
|
||
rf := NewFrame(FrameVersion, CmdRegisterTunnel, 0)
|
||
if err := rf.Marshal(wire); err != nil {
|
||
return err
|
||
}
|
||
return WriteFrame(control, rf)
|
||
}, CmdRegisterTunnel)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
rret := Ret{}
|
||
if err := f.Unmarshal(&rret); err != nil {
|
||
return 0, err
|
||
}
|
||
if rret.Code != 0 {
|
||
return 0, errors.New(rret.Msg)
|
||
}
|
||
w.everRegistered = true
|
||
w.ag.notifyPort(w.t.Id, rret.Port)
|
||
w.ag.logger().Infof("agent 隧道已连接节点 id=%s node=%s listen=%d",
|
||
w.t.Id, w.t.NodeAddr, rret.Port)
|
||
if w.ag.PingEvery > 0 {
|
||
go w.pingLoop()
|
||
}
|
||
return rret.Port, nil
|
||
}
|
||
|
||
// serve 数据流循环（阻塞；出错即断线，由 run 重连）
|
||
func (w *tunnelConn) serve() error {
|
||
w.mu.Lock()
|
||
sess := w.session
|
||
w.mu.Unlock()
|
||
if sess == nil {
|
||
return errors.New("未连接节点")
|
||
}
|
||
for {
|
||
stream, err := sess.AcceptStream()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
go w.ag.handleData(stream)
|
||
}
|
||
}
|
||
|
||
func (w *tunnelConn) pingLoop() {
|
||
ticker := time.NewTicker(w.ag.PingEvery)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-w.stopCh:
|
||
return
|
||
case <-ticker.C:
|
||
}
|
||
_, err := w.callControl(func(control net.Conn) error {
|
||
return WriteFrame(control, NewFrame(FrameVersion, CmdPing, 0))
|
||
}, CmdPing)
|
||
if err != nil {
|
||
if errors.Is(err, errTunnelClosed) {
|
||
// 节点已删除本隧道：杀掉会话让 serve 退出，run 收尾
|
||
w.teardown()
|
||
}
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
// callControl 串行化控制流交互。等待应答期间收到本隧道的关闭推送即返回 errTunnelClosed。
|
||
func (w *tunnelConn) callControl(write func(control net.Conn) error, wantCmd byte) (Frame, error) {
|
||
w.ctrlMu.Lock()
|
||
defer w.ctrlMu.Unlock()
|
||
|
||
w.mu.Lock()
|
||
control := w.control
|
||
w.mu.Unlock()
|
||
if control == nil {
|
||
return Frame{}, errors.New("未连接节点")
|
||
}
|
||
if err := write(control); err != nil {
|
||
return Frame{}, err
|
||
}
|
||
for {
|
||
f, err := ReadFrame(control)
|
||
if err != nil {
|
||
return Frame{}, err
|
||
}
|
||
if f.Cmd() == CmdTunnelClosed {
|
||
t := Tunnel{}
|
||
_ = f.Unmarshal(&t)
|
||
if t.Id == w.t.Id {
|
||
return Frame{}, errTunnelClosed
|
||
}
|
||
continue
|
||
}
|
||
if f.Cmd() == wantCmd {
|
||
return f, nil
|
||
}
|
||
// 其它命令忽略
|
||
}
|
||
}
|
||
|
||
// unregister 经控制流注销本隧道。传输失败视为成功：
|
||
// 连接已断时节点会随会话移除隧道，无需再通知。
|
||
func (w *tunnelConn) unregister() error {
|
||
w.mu.Lock()
|
||
connected := w.control != nil
|
||
w.mu.Unlock()
|
||
if !connected {
|
||
return nil
|
||
}
|
||
f, err := w.callControl(func(control net.Conn) error {
|
||
rf := NewFrame(FrameVersion, CmdUnregisterTunnel, 0)
|
||
if err := rf.Marshal(Tunnel{Id: w.t.Id}); err != nil {
|
||
return err
|
||
}
|
||
return WriteFrame(control, rf)
|
||
}, CmdUnregisterTunnel)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
ret := Ret{}
|
||
if err := f.Unmarshal(&ret); err != nil {
|
||
return nil
|
||
}
|
||
if ret.Code != 0 {
|
||
return errors.New(ret.Msg)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (w *tunnelConn) stop() {
|
||
w.stopOnce.Do(func() { close(w.stopCh) })
|
||
w.teardown()
|
||
}
|
||
|
||
func (w *tunnelConn) isStopped() bool {
|
||
select {
|
||
case <-w.stopCh:
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
func (w *tunnelConn) teardown() {
|
||
w.mu.Lock()
|
||
sess, conn := w.session, w.conn
|
||
w.session, w.conn, w.control = nil, nil, nil
|
||
w.mu.Unlock()
|
||
if sess != nil {
|
||
sess.Close()
|
||
}
|
||
if conn != nil {
|
||
conn.Close()
|
||
}
|
||
}
|