zonat/agent/tunnel.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()
	}
}