fix(rpc): bind io streams to target agents

Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
2026-05-21 02:01:54 +00:00
co-authored by naiba/CloudCode
parent 79c06d0f95
commit 8add559ec8
7 changed files with 99 additions and 19 deletions
+1 -1
View File
@@ -46,7 +46,7 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
return nil, err
}
rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c))
rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c), server.ID)
fmData, _ := json.Marshal(&model.TaskFM{
StreamID: streamId,
@@ -50,7 +50,7 @@ func TestTerminalStreamRejectsForeignMember(t *testing.T) {
gin.SetMode(gin.TestMode)
ensureLocalizerForStreamTests(t)
rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler()
rpc.NezhaHandlerSingleton.CreateStream("alice-terminal", 100)
rpc.NezhaHandlerSingleton.CreateStream("alice-terminal", 100, 1)
r := gin.New()
r.Use(func(c *gin.Context) {
@@ -76,7 +76,7 @@ func TestFMStreamRejectsForeignMember(t *testing.T) {
gin.SetMode(gin.TestMode)
ensureLocalizerForStreamTests(t)
rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler()
rpc.NezhaHandlerSingleton.CreateStream("alice-fm", 100)
rpc.NezhaHandlerSingleton.CreateStream("alice-fm", 100, 1)
r := gin.New()
r.Use(func(c *gin.Context) {
@@ -160,4 +160,3 @@ func TestWriteOauth2StateCookieIsHttpOnly(t *testing.T) {
t.Fatalf("nz-o2s must be HttpOnly to prevent XSS reading OAuth state, got %q", header)
}
}
+1 -1
View File
@@ -44,7 +44,7 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
return nil, err
}
rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c))
rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c), server.ID)
terminalData, _ := json.Marshal(&model.TerminalTask{
StreamID: streamId,
+5 -2
View File
@@ -140,8 +140,11 @@ func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
// NAT streams are anonymous HTTP-facing tunnels; they are NOT reachable
// via /ws/terminal or /ws/file (which check stream ownership), so the
// creator user ID does not need to identify a real user.
rpcService.NezhaHandlerSingleton.CreateStream(streamId, 0)
// creator user ID does not need to identify a real user. The targetServerID
// IS required though — the receiving agent must prove it is the server the
// NAT config addressed, otherwise any agent that snoops the streamId can
// answer NAT traffic on behalf of an unrelated host.
rpcService.NezhaHandlerSingleton.CreateStream(streamId, 0, server.ID)
defer rpcService.NezhaHandlerSingleton.CloseStream(streamId)
taskData, err := json.Marshal(model.TaskNAT{
+22 -1
View File
@@ -12,6 +12,7 @@ import (
type ioStreamContext struct {
creatorUserID uint64
targetServerID uint64
userIo io.ReadWriteCloser
agentIo io.ReadWriteCloser
userIoConnectCh chan struct{}
@@ -32,17 +33,37 @@ var bufPool = sync.Pool{
},
}
func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64) {
func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) {
s.ioStreamMutex.Lock()
defer s.ioStreamMutex.Unlock()
s.ioStreams[streamId] = &ioStreamContext{
creatorUserID: creatorUserID,
targetServerID: targetServerID,
userIoConnectCh: make(chan struct{}),
agentIoConnectCh: make(chan struct{}),
}
}
// IsStreamAuthorizedForAgent reports whether the connecting agent is the
// server the dashboard selected when CreateStream was called. Without this
// check any authenticated agent that learns an active streamId — via
// task-stream observation, leaked logs, or a shared global agent secret —
// can race in via IOStream() and serve a terminal / fm / NAT session that
// was addressed to a different server, turning the channel into a
// session-hijack RCE primitive. This is the agent-side dual of
// IsStreamAuthorizedForUser.
func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID uint64) bool {
s.ioStreamMutex.RLock()
defer s.ioStreamMutex.RUnlock()
ctx, ok := s.ioStreams[streamId]
if !ok {
return false
}
return ctx.targetServerID != 0 && ctx.targetServerID == agentServerID
}
// IsStreamAuthorizedForUser checks whether the requesting user may attach to
// the stream. A stream is reachable only by its creator or by an admin; any
// other authenticated user must be rejected. Unknown streams are always
+53 -7
View File
@@ -12,7 +12,7 @@ func TestIOStream(t *testing.T) {
const testStreamID = "ffffffff-ffff-ffff-ffff-ffffffffffff"
handler.CreateStream(testStreamID, 0)
handler.CreateStream(testStreamID, 0, 0)
userIo, agentIo := newPipeReadWriter(), newPipeReadWriter()
defer func() {
userIo.Close()
@@ -108,7 +108,7 @@ func newPipeReadWriter() io.ReadWriteCloser {
func TestStreamOwnershipReturnsCreatorUserID(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100)
h.CreateStream("alice-stream", 100, 0)
creator, found := h.StreamOwnership("alice-stream")
if !found {
@@ -128,8 +128,8 @@ func TestStreamOwnershipReturnsNotFoundForUnknownID(t *testing.T) {
func TestStreamOwnershipPreservesPerStreamCreator(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100)
h.CreateStream("bob-stream", 200)
h.CreateStream("alice-stream", 100, 0)
h.CreateStream("bob-stream", 200, 0)
aliceCreator, _ := h.StreamOwnership("alice-stream")
bobCreator, _ := h.StreamOwnership("bob-stream")
@@ -141,7 +141,7 @@ func TestStreamOwnershipPreservesPerStreamCreator(t *testing.T) {
func TestIsStreamAuthorizedForUserAllowsCreator(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100)
h.CreateStream("alice-stream", 100, 0)
if !h.IsStreamAuthorizedForUser("alice-stream", 100, false) {
t.Fatalf("creator must be authorized to attach to their own stream")
@@ -150,7 +150,7 @@ func TestIsStreamAuthorizedForUserAllowsCreator(t *testing.T) {
func TestIsStreamAuthorizedForUserDeniesForeignMember(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100)
h.CreateStream("alice-stream", 100, 0)
if h.IsStreamAuthorizedForUser("alice-stream", 200, false) {
t.Fatalf("foreign member must not be authorized — session hijack would be possible")
@@ -159,7 +159,7 @@ func TestIsStreamAuthorizedForUserDeniesForeignMember(t *testing.T) {
func TestIsStreamAuthorizedForUserAllowsAdmin(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("alice-stream", 100)
h.CreateStream("alice-stream", 100, 0)
if !h.IsStreamAuthorizedForUser("alice-stream", 999, true) {
t.Fatalf("admin must be authorized to attach regardless of creator")
@@ -197,6 +197,52 @@ func TestIsValidIOStreamMagicRejectsShortData(t *testing.T) {
}
}
// Agent-side stream authorization is the dual of IsStreamAuthorizedForUser:
// only the server the dashboard selected when CreateStream was called may
// attach via IOStream(). Without it, any authenticated agent that learns an
// active streamId (task-stream observation, leaked logs) can race in and
// serve a terminal/fm/NAT session originally addressed to a different
// server — a session-hijack RCE intermediation primitive.
func TestIsStreamAuthorizedForAgentAllowsBoundServer(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("terminal-for-server-100", 1, 100)
if !h.IsStreamAuthorizedForAgent("terminal-for-server-100", 100) {
t.Fatalf("the bound target server must be authorized to attach")
}
}
func TestIsStreamAuthorizedForAgentDeniesForeignServer(t *testing.T) {
h := NewNezhaHandler()
h.CreateStream("terminal-for-server-100", 1, 100)
if h.IsStreamAuthorizedForAgent("terminal-for-server-100", 200) {
t.Fatalf("a foreign agent must not be able to attach — session hijack would be possible")
}
}
func TestIsStreamAuthorizedForAgentDeniesUnboundStream(t *testing.T) {
h := NewNezhaHandler()
// targetServerID == 0 means the stream was created without a bound agent
// — no agent should be allowed to attach.
h.CreateStream("unbound-stream", 1, 0)
if h.IsStreamAuthorizedForAgent("unbound-stream", 100) {
t.Fatalf("unbound stream must not authorize any agent")
}
if h.IsStreamAuthorizedForAgent("unbound-stream", 0) {
t.Fatalf("unbound stream must not authorize a zero clientID either")
}
}
func TestIsStreamAuthorizedForAgentDeniesUnknownStreamID(t *testing.T) {
h := NewNezhaHandler()
if h.IsStreamAuthorizedForAgent("nonexistent", 100) {
t.Fatalf("unknown stream id must not authorize any agent")
}
}
func TestIsValidIOStreamMagicRejectsPartialOrWrongMagic(t *testing.T) {
// Each case has at least one byte that does NOT match the magic. The
// previous && short-circuit bug let cases like {0xff, 0, 0, 0} pass
+15 -4
View File
@@ -57,7 +57,8 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err
case model.TaskTypeCommand:
// 处理上报的计划任务
cr, _ := singleton.CronShared.Get(result.GetId())
if cr != nil {
// 任务结果 ID 来自 agent,必须确认该 cron 本应派发给当前 reporter。
if singleton.CanReportCronResult(cr, server) {
// 保存当前服务器状态信息
var curServer model.Server
copier.Copy(&curServer, server)
@@ -208,7 +209,8 @@ func (s *NezhaHandler) ReportSystemInfo2(c context.Context, r *pb.Host) (*pb.Uin
}
func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
if _, err := s.Auth.Check(stream.Context()); err != nil {
clientID, err := s.Auth.Check(stream.Context())
if err != nil {
return err
}
id, err := stream.Recv()
@@ -222,6 +224,17 @@ func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
return fmt.Errorf("invalid stream id")
}
streamId := string(id.Data[4:])
// agent 侧归属校验:只有 createTerminal / createFM / ServeNAT 选定的目标 server
// 才能接管该 stream。漏掉这一步等同于把 terminal / fm / NAT 会话向所有合法 agent
// 开放(任何获得 streamId 的 agent 都能抢答),构成 session-hijack RCE 中介。
// 这是 commit 6661d6auser 侧归属校验)的对偶补丁。先校验后启 keepalive,
// 避免未授权 agent 触发悬空 goroutine 持续向其发心跳。
if !s.IsStreamAuthorizedForAgent(streamId, clientID) {
return fmt.Errorf("stream not authorized for agent")
}
go func() {
for {
if err := stream.Send(&pb.IOStreamData{Data: []byte{}}); err != nil {
@@ -232,8 +245,6 @@ func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error {
}
}()
streamId := string(id.Data[4:])
if _, err := s.GetStream(streamId); err != nil {
return err
}