138 lines
3.6 KiB
Go
138 lines
3.6 KiB
Go
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())
|
|
}
|