init: 自 zomaintain/backend/zonat 平移独立成库; module path 改为 git.zeroonesoft.cn/golib/zonat; 内部 import 同步重写

2026-09-19 18:55:43 +08:00
commit 984041f1a0
20 changed files with 2427 additions and 0 deletions
@@ -0,0 +1,137 @@
+package node
+
+import (
+	"net"
+	"sync"
+	"sync/atomic"
+
+	"git.zeroonesoft.cn/golib/zonat/agent"
+
+	"github.com/xtaci/smux"
+)
+
+// Session 一个 agent 连接的会话：smux 服务端 + 控制流处理。
+// 数据流由节点侧隧道监听主动 OpenStream；agent 主动打开的第一条流是控制流，
+// 用于注册/注销隧道与心跳。
+type Session struct {
+	agentId string
+	node    *Node
+	session *smux.Session
+
+	control   net.Conn   // agent 打开的控制流
+	controlMu sync.Mutex // 控制流写锁（handleControl 与 pushClosed 并发写）
+	dead      atomic.Bool
+}
+
+func (s *Session) isDead() bool { return s.dead.Load() }
+
+func (s *Session) close() {
+	s.dead.Store(true)
+	s.session.Close()
+}
+
+// OpenStream 建立一条到 agent 的数据流
+func (s *Session) OpenStream() (net.Conn, error) {
+	return s.session.OpenStream()
+}
+
+// loop 接受 agent 打开的控制流，直到会话结束
+func (s *Session) loop() {
+	for {
+		stream, err := s.session.AcceptStream()
+		if err != nil {
+			s.close()
+			return
+		}
+		// agent 打开的第一条流是控制流，pushClosed 依赖它；
+		// 赋值与写入共用 controlMu，避免与推送协程竞争
+		s.controlMu.Lock()
+		if s.control == nil {
+			s.control = stream
+		}
+		s.controlMu.Unlock()
+		go s.handleControl(stream)
+	}
+}
+
+// handleControl 处理控制流上的帧：注册/注销隧道、心跳。
+// 单条控制流内的帧顺序处理；应答与节点推送共用 controlMu 串行写入。
+func (s *Session) handleControl(stream net.Conn) {
+	defer stream.Close()
+	for {
+		f, err := agent.ReadFrame(stream)
+		if err != nil {
+			return
+		}
+		switch f.Cmd() {
+		case agent.CmdRegisterTunnel:
+			t := agent.Tunnel{}
+			ret := agent.Ret{Code: 1, Msg: "注册失败"}
+			if err := f.Unmarshal(&t); err != nil {
+				ret.Msg = err.Error()
+			} else if nt, err := s.node.registerTunnel(s, t); err != nil {
+				ret.Msg = err.Error()
+			} else {
+				ret = agent.Ret{Code: 0, Msg: "注册成功", Port: nt.ListenPort}
+			}
+			s.reply(stream, f, ret)
+		case agent.CmdUnregisterTunnel:
+			t := agent.Tunnel{}
+			ret := agent.Ret{Code: 1, Msg: "注销失败"}
+			if err := f.Unmarshal(&t); err != nil {
+				ret.Msg = err.Error()
+			} else if s.node.unregisterTunnel(t.Id) {
+				ret = agent.Ret{Code: 0, Msg: "注销成功"}
+			} else {
+				ret.Msg = "隧道不存在"
+			}
+			s.reply(stream, f, ret)
+		case agent.CmdPing:
+			s.writeControl(func() error {
+				return agent.WriteFrame(stream, agent.NewFrame(agent.FrameVersion, agent.CmdPing, f.StreamID()))
+			})
+		case agent.CmdLogin, agent.CmdTarget, agent.CmdTunnelClosed:
+			// 不该出现在 agent→节点方向，忽略
+		default:
+			s.reply(stream, f, agent.Ret{Code: 1, Msg: "未知命令"})
+		}
+	}
+}
+
+func (s *Session) reply(stream net.Conn, req agent.Frame, ret agent.Ret) {
+	s.writeControl(func() error {
+		f := agent.NewFrame(agent.FrameVersion, req.Cmd(), req.StreamID())
+		if err := f.Marshal(ret); err != nil {
+			return err
+		}
+		return agent.WriteFrame(stream, f)
+	})
+}
+
+// pushClosed 节点主动通知 agent：隧道已被删除（TTL 过期等）
+func (s *Session) pushClosed(tunnelId string) {
+	if s.control == nil {
+		return
+	}
+	s.writeControl(func() error {
+		f := agent.NewFrame(agent.FrameVersion, agent.CmdTunnelClosed, 0)
+		if err := f.Marshal(agent.Tunnel{Id: tunnelId}); err != nil {
+			return err
+		}
+		return agent.WriteFrame(s.control, f)
+	})
+}
+
+func (s *Session) writeControl(fn func() error) {
+	s.controlMu.Lock()
+	defer s.controlMu.Unlock()
+	if s.dead.Load() {
+		return
+	}
+	// 写失败意味着会话已死，由读循环感知并收尾
+	_ = fn()
+}
+
+func smuxServer(conn net.Conn) (*smux.Session, error) {
+	return smux.Server(conn, agent.SmuxConfig())
+}