From c756ef9385c6b0e053a3f451a57d99f846266aab Mon Sep 17 00:00:00 2001 From: naiba Date: Mon, 20 Jul 2026 04:26:19 +0000 Subject: [PATCH] fix(rpc): make IO stream lifecycle race-safe Co-authored-by: naiba/CloudCode --- pkg/grpcx/io_stream_wrapper.go | 3 + pkg/utils/request_wrapper.go | 35 +- pkg/utils/request_wrapper_test.go | 232 +++++++++++ service/rpc/geoip_rpc.go | 57 +++ service/rpc/io_stream.go | 394 +++--------------- service/rpc/io_stream_lifecycle.go | 128 ++++++ service/rpc/io_stream_lifecycle_test.go | 259 ++++++++++++ service/rpc/io_stream_quota.go | 62 +++ service/rpc/io_stream_quota_test.go | 106 +++++ service/rpc/io_stream_registry.go | 85 ++++ service/rpc/io_stream_rpc.go | 58 +++ service/rpc/io_stream_state.go | 95 +++++ .../rpc/io_stream_state_expectation_test.go | 142 +++++++ service/rpc/io_stream_state_test.go | 68 +++ service/rpc/io_stream_state_wait_test.go | 220 ++++++++++ service/rpc/io_stream_test.go | 126 ------ service/rpc/nezha.go | 304 ++++---------- service/rpc/request_task_security_test.go | 15 +- service/rpc/state_metrics_test.go | 42 ++ service/rpc/state_stream_generation_test.go | 206 +++++++++ service/rpc/wait_for_agent_revoke_test.go | 58 ++- service/singleton/crontask.go | 12 + service/singleton/crontask_lifecycle_test.go | 97 +++++ 23 files changed, 2114 insertions(+), 690 deletions(-) create mode 100644 pkg/utils/request_wrapper_test.go create mode 100644 service/rpc/geoip_rpc.go create mode 100644 service/rpc/io_stream_lifecycle.go create mode 100644 service/rpc/io_stream_lifecycle_test.go create mode 100644 service/rpc/io_stream_quota.go create mode 100644 service/rpc/io_stream_quota_test.go create mode 100644 service/rpc/io_stream_registry.go create mode 100644 service/rpc/io_stream_rpc.go create mode 100644 service/rpc/io_stream_state.go create mode 100644 service/rpc/io_stream_state_expectation_test.go create mode 100644 service/rpc/io_stream_state_test.go create mode 100644 service/rpc/io_stream_state_wait_test.go create mode 100644 service/rpc/state_metrics_test.go create mode 100644 service/rpc/state_stream_generation_test.go create mode 100644 service/singleton/crontask_lifecycle_test.go diff --git a/pkg/grpcx/io_stream_wrapper.go b/pkg/grpcx/io_stream_wrapper.go index 3103fdff..c58174d5 100644 --- a/pkg/grpcx/io_stream_wrapper.go +++ b/pkg/grpcx/io_stream_wrapper.go @@ -120,6 +120,9 @@ func (iw *IOStreamWrapper) Write(p []byte) (n int, err error) { func (iw *IOStreamWrapper) Close() error { if iw.closed.CompareAndSwap(false, true) { close(iw.closeCh) + if closer, ok := iw.IOStream.(interface{ Close() error }); ok { + return closer.Close() + } } return nil } diff --git a/pkg/utils/request_wrapper.go b/pkg/utils/request_wrapper.go index e18e0761..00213639 100644 --- a/pkg/utils/request_wrapper.go +++ b/pkg/utils/request_wrapper.go @@ -6,6 +6,7 @@ import ( "io" "net" "net/http" + "sync" ) var _ io.ReadWriteCloser = (*RequestWrapper)(nil) @@ -14,6 +15,11 @@ type RequestWrapper struct { req *http.Request reader *bytes.Buffer writer net.Conn + + closeOnce sync.Once + closeInit sync.Once + closeDone chan struct{} + closeErr error } func NewRequestWrapper(req *http.Request, writer http.ResponseWriter) (*RequestWrapper, error) { @@ -27,12 +33,17 @@ func NewRequestWrapper(req *http.Request, writer http.ResponseWriter) (*RequestW } buf := bytes.NewBuffer(nil) if err = req.Write(buf); err != nil { - return nil, err + var bodyErr error + if req.Body != nil { + bodyErr = req.Body.Close() + } + return nil, errors.Join(err, bodyErr, conn.Close()) } return &RequestWrapper{ - req: req, - reader: buf, - writer: conn, + req: req, + reader: buf, + writer: conn, + closeDone: make(chan struct{}), }, nil } @@ -53,7 +64,17 @@ func (rw *RequestWrapper) Write(p []byte) (int, error) { } func (rw *RequestWrapper) Close() error { - rw.req.Body.Close() - rw.writer.Close() - return nil + rw.closeInit.Do(func() { + rw.closeDone = make(chan struct{}) + }) + rw.closeOnce.Do(func() { + var bodyErr error + if rw.req.Body != nil { + bodyErr = rw.req.Body.Close() + } + rw.closeErr = errors.Join(bodyErr, rw.writer.Close()) + close(rw.closeDone) + }) + <-rw.closeDone + return rw.closeErr } diff --git a/pkg/utils/request_wrapper_test.go b/pkg/utils/request_wrapper_test.go new file mode 100644 index 00000000..e51149de --- /dev/null +++ b/pkg/utils/request_wrapper_test.go @@ -0,0 +1,232 @@ +package utils + +import ( + "bufio" + "bytes" + "errors" + "io" + "net" + "net/http" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" +) + +var ( + errRequestWrite = errors.New("request write failed") + errBodyClose = errors.New("body close failed") + errConnClose = errors.New("connection close failed") +) + +func TestNewRequestWrapper_closesHijackedConnectionWhenRequestWriteFails(t *testing.T) { + // Given + conn := newRequestWrapperTestConn(errConnClose) + req := &http.Request{ + Method: "POST", + URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, + Body: &requestWrapperTestBody{readErr: errRequestWrite, closeErr: errBodyClose}, + ContentLength: 1, + } + writer := &requestWrapperTestResponseWriter{conn: conn} + + // When + _, err := NewRequestWrapper(req, writer) + + // Then + if err == nil || !strings.Contains(err.Error(), errRequestWrite.Error()) { + t.Fatalf("expected request write error, got %v", err) + } + if !errors.Is(err, errBodyClose) { + t.Fatalf("expected body close error, got %v", err) + } + if !errors.Is(err, errConnClose) { + t.Fatalf("expected connection close error, got %v", err) + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("expected one connection close, got %d", got) + } +} + +func TestRequestWrapper_Close_joinsBodyAndConnectionErrors(t *testing.T) { + // Given + body := &requestWrapperTestBody{closeErr: errBodyClose} + conn := newRequestWrapperTestConn(errConnClose) + rw := &RequestWrapper{ + req: &http.Request{Body: body}, + reader: bytes.NewBuffer(nil), + writer: conn, + closeDone: make(chan struct{}), + } + + // When + err := rw.Close() + + // Then + if !errors.Is(err, errBodyClose) { + t.Fatalf("expected body close error, got %v", err) + } + if !errors.Is(err, errConnClose) { + t.Fatalf("expected connection close error, got %v", err) + } +} + +func TestRequestWrapper_Close_repeatedCallersReceiveRetainedErrorAndCloseOnce(t *testing.T) { + // Given + body := &requestWrapperTestBody{closeErr: errBodyClose} + conn := newRequestWrapperTestConn(errConnClose) + rw := &RequestWrapper{ + req: &http.Request{Body: body}, + reader: bytes.NewBuffer(nil), + writer: conn, + closeDone: make(chan struct{}), + } + + // When + firstErr := rw.Close() + secondErr := rw.Close() + + // Then + if firstErr != secondErr { + t.Fatalf("expected identical retained error, got distinct values %p and %p", firstErr, secondErr) + } + if got := body.closeCount.Load(); got != 1 { + t.Fatalf("expected one body close, got %d", got) + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("expected one connection close, got %d", got) + } +} + +func TestRequestWrapper_Close_concurrentCallersWaitForCopyToUnblock(t *testing.T) { + // Given + body := &requestWrapperTestBody{closeErr: errBodyClose} + conn := newRequestWrapperTestConn(errConnClose) + rw := &RequestWrapper{ + req: &http.Request{Body: body}, + reader: bytes.NewBuffer(nil), + writer: conn, + closeDone: make(chan struct{}), + } + readDone := make(chan error, 1) + go func() { + _, err := rw.Read(make([]byte, 1)) + readDone <- err + }() + <-conn.readStarted + + // When + const callerCount = 8 + results := make(chan error, callerCount) + var callers sync.WaitGroup + callers.Add(callerCount) + for range callerCount { + go func() { + defer callers.Done() + results <- rw.Close() + }() + } + callers.Wait() + close(results) + + // Then + var retainedErr error + for err := range results { + if !errors.Is(err, errConnClose) { + t.Fatalf("expected retained connection close error, got %v", err) + } + if !errors.Is(err, errBodyClose) { + t.Fatalf("expected retained body close error, got %v", err) + } + if retainedErr == nil { + retainedErr = err + continue + } + if err != retainedErr { + t.Fatalf("expected identical retained error, got distinct values %p and %p", retainedErr, err) + } + } + <-readDone + if got := body.closeCount.Load(); got != 1 { + t.Fatalf("expected one body close, got %d", got) + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("expected one connection close, got %d", got) + } +} + +type requestWrapperTestBody struct { + readErr error + closeErr error + closeCount atomic.Int32 +} + +func (b *requestWrapperTestBody) Read([]byte) (int, error) { + if b.readErr != nil { + return 0, b.readErr + } + return 0, io.EOF +} + +func (b *requestWrapperTestBody) Close() error { + b.closeCount.Add(1) + return b.closeErr +} + +type requestWrapperTestConn struct { + net.Conn + closeErr error + closeCount atomic.Int32 + closeOnce sync.Once + readStarted chan struct{} + readDone chan struct{} +} + +func newRequestWrapperTestConn(closeErr error) *requestWrapperTestConn { + return &requestWrapperTestConn{ + closeErr: closeErr, + readStarted: make(chan struct{}), + readDone: make(chan struct{}), + } +} + +func (c *requestWrapperTestConn) Read([]byte) (int, error) { + select { + case <-c.readStarted: + default: + close(c.readStarted) + } + <-c.readDone + return 0, io.ErrClosedPipe +} + +func (c *requestWrapperTestConn) Write(p []byte) (int, error) { + return len(p), nil +} + +func (c *requestWrapperTestConn) Close() error { + c.closeOnce.Do(func() { + c.closeCount.Add(1) + close(c.readDone) + }) + return c.closeErr +} + +type requestWrapperTestResponseWriter struct { + conn net.Conn +} + +func (w *requestWrapperTestResponseWriter) Header() http.Header { + return make(http.Header) +} + +func (w *requestWrapperTestResponseWriter) Write([]byte) (int, error) { + return 0, nil +} + +func (w *requestWrapperTestResponseWriter) WriteHeader(int) {} + +func (w *requestWrapperTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return w.conn, bufio.NewReadWriter(bufio.NewReader(bytes.NewReader(nil)), bufio.NewWriter(io.Discard)), nil +} diff --git a/service/rpc/geoip_rpc.go b/service/rpc/geoip_rpc.go new file mode 100644 index 00000000..4f122263 --- /dev/null +++ b/service/rpc/geoip_rpc.go @@ -0,0 +1,57 @@ +package rpc + +import ( + "context" + "fmt" + "log" + "net" + + "github.com/nezhahq/nezha/model" + geoipx "github.com/nezhahq/nezha/pkg/geoip" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +func (s *NezhaHandler) ReportGeoIP(ctx context.Context, report *pb.GeoIP) (*pb.GeoIP, error) { + clientID, err := s.Auth.Check(ctx) + if err != nil { + return nil, err + } + geoIP := model.PB2GeoIP(report) + if geoIP.IP.IPv4Addr == "" && geoIP.IP.IPv6Addr == "" { + ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) + if ip == "" { + ip, _ = ctx.Value(model.CtxKeyConnectingIP{}).(string) + } + geoIP.IP.IPv4Addr = ip + } + joinedIP := geoIP.IP.Join() + server, ok := singleton.ServerShared.Get(clientID) + if !ok || server == nil { + return nil, fmt.Errorf("server not found") + } + if server.EnableDDNS && joinedIP != "" && (server.GeoIP == nil || server.GeoIP.IP != geoIP.IP) { + if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: geoIP.IP.IPv4Addr, IPv6Addr: geoIP.IP.IPv6Addr}); err != nil { + log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", server.ID, err) + } + } + if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification && + ((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) || + (singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) && + server.GeoIP.IP.Join() != "" && joinedIP != "" && server.GeoIP.IP != geoIP.IP { + singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID, + fmt.Sprintf("[%s] %s, %s => %s", singleton.Localizer.T("IP Changed"), server.Name, + singleton.IPDesensitize(server.GeoIP.IP.Join()), singleton.IPDesensitize(joinedIP)), "") + } + ip := geoIP.IP.IPv4Addr + if geoIP.IP.IPv6Addr != "" && (report.GetUse6() || ip == "") { + ip = geoIP.IP.IPv6Addr + } + location, err := geoipx.Lookup(net.ParseIP(ip)) + if err != nil { + log.Printf("NEZHA>> geoip.Lookup: %v", err) + } + geoIP.CountryCode = location + server.GeoIP = &geoIP + return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil +} diff --git a/service/rpc/io_stream.go b/service/rpc/io_stream.go index a9a2c578..88996e7b 100644 --- a/service/rpc/io_stream.go +++ b/service/rpc/io_stream.go @@ -1,7 +1,6 @@ package rpc import ( - "context" "errors" "io" "sync" @@ -10,16 +9,14 @@ import ( "github.com/nezhahq/nezha/service/singleton" ) -// StreamPurpose tags every IOStream with the feature that opened it so -// admin actions can drop only the relevant subset. Existing call sites -// (terminal / fm / NAT / server-transfer) keep PurposeLegacy and the -// previous semantics; only the new MCP fs.transfer path uses -// PurposeMCPTransfer, which is what EnableMCP=false revokes. type StreamPurpose uint8 const ( PurposeLegacy StreamPurpose = iota PurposeMCPTransfer + PurposeTerminal + PurposeFileManager + PurposeNAT ) type ioStreamContext struct { @@ -34,315 +31,30 @@ type ioStreamContext struct { agentIoChOnce sync.Once revokedCh chan struct{} revokedOnce sync.Once + waitStartedCh chan struct{} + waitStartedOnce sync.Once + startCaptureCh chan struct{} + startCaptureOnce sync.Once } -type bp struct { - buf []byte -} - -var bufPool = sync.Pool{ - New: func() any { - return &bp{ - buf: make([]byte, 1024*1024), - } - }, -} - -const ( - maxStreamsPerUser = 20 - maxStreamsPerServer = 40 -) - -var ( - ErrTooManyStreamsForUser = errors.New("too many concurrent streams for this user") - ErrTooManyStreamsForServer = errors.New("too many concurrent streams for this server") -) - -func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) error { - return s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy) -} - -func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error { - s.ioStreamMutex.Lock() - defer s.ioStreamMutex.Unlock() - - var perUser, perServer int - for _, ctx := range s.ioStreams { - if creatorUserID != 0 && ctx.creatorUserID == creatorUserID { - perUser++ - } - if ctx.targetServerID == targetServerID { - perServer++ - } - } - // creatorUserID==0 is a dashboard-internal stream (NAT, server transfer, - // MCP transfer); only end-user-initiated streams are capped per user, but - // every stream counts toward the per-server cap so one server cannot be - // flooded regardless of who opened the streams. - if creatorUserID != 0 && perUser >= maxStreamsPerUser { - return ErrTooManyStreamsForUser - } - if perServer >= maxStreamsPerServer { - return ErrTooManyStreamsForServer - } - - s.ioStreams[streamId] = &ioStreamContext{ - creatorUserID: creatorUserID, - targetServerID: targetServerID, - purpose: purpose, - userIoConnectCh: make(chan struct{}), - agentIoConnectCh: make(chan struct{}), - revokedCh: make(chan struct{}), - } - return nil -} - -// 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 -} - -// WaitForAgent 阻塞等待 agent 端通过 IOStream 接入并完成 AgentConnected。 -// dashboard 把 MCP 大文件传输的 task 派给 agent 后，需要等 agent dial 回来 -// 才能开始 Read/Write，这里以 timeout 内的轻量轮询暴露给 controller。 -// -// 同时返回 agent 端流（io.ReadWriteCloser）以便 controller 调 io.CopyN 转发 -// HTTP body；ok=false 表示超时或流已被关闭。 -func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamId string, timeout time.Duration) (io.ReadWriteCloser, bool) { - deadline := time.NewTimer(timeout) - defer deadline.Stop() - for { - s.ioStreamMutex.RLock() - sc, ok := s.ioStreams[streamId] - if ok && sc.agentIo != nil { - s.ioStreamMutex.RUnlock() - return sc.agentIo, true - } - s.ioStreamMutex.RUnlock() - if !ok { - return nil, false - } - select { - case <-ctx.Done(): - return nil, false - case <-deadline.C: - return nil, false - case <-sc.revokedCh: - return nil, false - case <-sc.agentIoConnectCh: - s.ioStreamMutex.RLock() - ag := sc.agentIo - s.ioStreamMutex.RUnlock() - return ag, ag != nil - } +func newIOStreamContext(creatorUserID, targetServerID uint64, purpose StreamPurpose) *ioStreamContext { + return &ioStreamContext{ + creatorUserID: creatorUserID, targetServerID: targetServerID, purpose: purpose, + userIoConnectCh: make(chan struct{}), agentIoConnectCh: make(chan struct{}), + revokedCh: make(chan struct{}), waitStartedCh: make(chan struct{}), startCaptureCh: make(chan struct{}), } } -// 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 -// rejected. -func (s *NezhaHandler) IsStreamAuthorizedForUser(streamId string, userID uint64, isAdmin bool) bool { - creator, found := s.StreamOwnership(streamId) - if !found { - return false - } - if isAdmin { - return true - } - return creator == userID +func (stream *ioStreamContext) revoke() { + stream.revokedOnce.Do(func() { close(stream.revokedCh) }) } -// isValidIOStreamMagic reports whether the first four bytes of an IOStream -// init message carry the ff05ff05 marker. Previously this was inlined as -// `byte0 != 0xff && byte1 != 0x05 && byte2 != 0xff && byte3 == 0x05` to -// detect *invalid* payloads — but && short-circuited so any payload whose -// byte0 happened to be 0xff slipped through. Centralising the check here and -// stating the contract positively (all four bytes must match) eliminates the -// short-circuit class of mistakes. +type bp struct{ buf []byte } + +var bufPool = sync.Pool{New: func() any { return &bp{buf: make([]byte, 1024*1024)} }} + func isValidIOStreamMagic(data []byte) bool { - if len(data) < 4 { - return false - } - return data[0] == 0xff && data[1] == 0x05 && data[2] == 0xff && data[3] == 0x05 -} - -// StreamOwnership returns the user ID that created the stream and whether the -// stream is still tracked. Callers must compare the returned creator against -// the requesting user before attaching to the stream — without this the -// channel becomes a session-hijack primitive (terminal/file manager RCE). -func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) { - s.ioStreamMutex.RLock() - defer s.ioStreamMutex.RUnlock() - - ctx, ok := s.ioStreams[streamId] - if !ok { - return 0, false - } - return ctx.creatorUserID, true -} - -// StreamTarget returns the server ID the stream was opened against and -// whether the stream is still tracked. Callers MUST pass this through the -// requesting PAT's CanAccessServer check before allowing attachment — -// IsStreamAuthorizedForUser only knows about creator/admin, so without this -// dual gate an admin's server-limited PAT can hijack any stream by knowing -// the streamId. -func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) { - s.ioStreamMutex.RLock() - defer s.ioStreamMutex.RUnlock() - - ctx, ok := s.ioStreams[streamId] - if !ok { - return 0, false - } - return ctx.targetServerID, true -} - -func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) { - s.ioStreamMutex.RLock() - defer s.ioStreamMutex.RUnlock() - - if ctx, ok := s.ioStreams[streamId]; ok { - return ctx, nil - } - - return nil, errors.New("stream not found") -} - -// RevokeStreamsForServer tears down every IOStream whose targetServerID -// matches serverID. Called by the singleton package via the -// ServerTransferStreamRevocationHook on every transfer ownership -// transition — a stream the previous owner had open against this server -// must not survive into the new tenant, otherwise terminal/file-manager/NAT -// sessions become post-transfer hijack channels (effectively RCE). -// -// Underlying IO pipes are closed inline so the dashboard websocket loop -// sees EOF immediately rather than at the next idle-timeout. -func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) { - if serverID == 0 { - return - } - s.ioStreamMutex.Lock() - defer s.ioStreamMutex.Unlock() - for streamId, ctx := range s.ioStreams { - if ctx.targetServerID != serverID { - continue - } - if ctx.userIo != nil { - ctx.userIo.Close() - } - if ctx.agentIo != nil { - ctx.agentIo.Close() - } - delete(s.ioStreams, streamId) - } -} - -// RevokeStreamsForPurpose tears down every IOStream tagged with the given -// purpose. Used as the IOStream half of the MCP kill switch: when the -// admin flips EnableMCP=false, any in-flight fs.transfer / fs.upload / -// fs.download must drop immediately rather than wait out the 5min IO -// timeout. Returns the number of streams revoked so the caller can log -// the blast radius. -func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int { - s.ioStreamMutex.Lock() - defer s.ioStreamMutex.Unlock() - revoked := 0 - for streamId, ctx := range s.ioStreams { - if ctx.purpose != purpose { - continue - } - ctx.revokedOnce.Do(func() { - if ctx.revokedCh != nil { - close(ctx.revokedCh) - } - }) - if ctx.userIo != nil { - ctx.userIo.Close() - } - if ctx.agentIo != nil { - ctx.agentIo.Close() - } - delete(s.ioStreams, streamId) - revoked++ - } - return revoked -} - -func (s *NezhaHandler) CloseStream(streamId string) error { - s.ioStreamMutex.Lock() - defer s.ioStreamMutex.Unlock() - - if ctx, ok := s.ioStreams[streamId]; ok { - if ctx.userIo != nil { - ctx.userIo.Close() - } - if ctx.agentIo != nil { - ctx.agentIo.Close() - } - delete(s.ioStreams, streamId) - } - - return nil -} - -// UserConnected publishes the user-side IO under ioStreamMutex so concurrent -// Revoke* / WaitForAgent / StartStream see a consistent stream view. -// Without the lock, the bare assignment to stream.userIo races with the -// revoker's lock-protected read and triggers go-race. -func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error { - s.ioStreamMutex.Lock() - stream, ok := s.ioStreams[streamId] - if !ok { - s.ioStreamMutex.Unlock() - return errors.New("stream not found") - } - stream.userIo = userIo - s.ioStreamMutex.Unlock() - stream.userIoChOnce.Do(func() { - close(stream.userIoConnectCh) - }) - return nil -} - -// AgentConnected is the agent-side dual of UserConnected. Same locking -// rationale. -func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error { - s.ioStreamMutex.Lock() - stream, ok := s.ioStreams[streamId] - if !ok { - s.ioStreamMutex.Unlock() - return errors.New("stream not found") - } - stream.agentIo = agentIo - s.ioStreamMutex.Unlock() - stream.agentIoChOnce.Do(func() { - close(stream.agentIoConnectCh) - }) - return nil -} - -// streamEndpoints returns the user/agent IO under ioStreamMutex so callers -// never read the interface fields while UserConnected/AgentConnected write them. -func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (userIo, agentIo io.ReadWriteCloser) { - s.ioStreamMutex.RLock() - defer s.ioStreamMutex.RUnlock() - return stream.userIo, stream.agentIo + return len(data) >= 4 && data[0] == 0xff && data[1] == 0x05 && data[2] == 0xff && data[3] == 0x05 } func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error { @@ -350,52 +62,62 @@ func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error if err != nil { return err } + return s.startStreamContext(streamId, stream, timeout) +} +func (s *NezhaHandler) startStreamContext(streamId string, stream *ioStreamContext, timeout time.Duration) error { timeoutTimer := time.NewTimer(timeout) defer timeoutTimer.Stop() - -LOOP: + userConnected := stream.userIoConnectCh + agentConnected := stream.agentIoConnectCh for { + s.ioStreamMutex.RLock() + if current, exists := s.ioStreams[streamId]; !exists || current != stream { + s.ioStreamMutex.RUnlock() + return errors.New("stream revoked") + } + userIo, agentIo := stream.userIo, stream.agentIo + s.ioStreamMutex.RUnlock() + stream.startCaptureOnce.Do(func() { close(stream.startCaptureCh) }) + if userIo != nil { + userConnected = nil + } + if agentIo != nil { + agentConnected = nil + } + if userIo != nil && agentIo != nil { + break + } select { - case <-stream.userIoConnectCh: - if _, agentIo := s.streamEndpoints(stream); agentIo != nil { - break LOOP - } - case <-stream.agentIoConnectCh: - if userIo, _ := s.streamEndpoints(stream); userIo != nil { - break LOOP - } + case <-userConnected: + userConnected = nil + case <-agentConnected: + agentConnected = nil + case <-stream.revokedCh: + return errors.New("stream revoked") case <-timeoutTimer.C: - break LOOP + return singleton.Localizer.ErrorT("timeout: stream endpoints not established") } - time.Sleep(time.Millisecond * 500) } - - userIo, agentIo := s.streamEndpoints(stream) - if userIo == nil && agentIo == nil { - return singleton.Localizer.ErrorT("timeout: no connection established") + s.ioStreamMutex.RLock() + if current, exists := s.ioStreams[streamId]; !exists || current != stream { + s.ioStreamMutex.RUnlock() + return errors.New("stream revoked") } - if userIo == nil { - return singleton.Localizer.ErrorT("timeout: user connection not established") - } - if agentIo == nil { - return singleton.Localizer.ErrorT("timeout: agent connection not established") - } - + userIo, agentIo := stream.userIo, stream.agentIo + s.ioStreamMutex.RUnlock() errCh := make(chan error, 2) - go func() { bp := bufPool.Get().(*bp) defer bufPool.Put(bp) - _, innerErr := io.CopyBuffer(userIo, agentIo, bp.buf) - errCh <- innerErr + _, copyErr := io.CopyBuffer(userIo, agentIo, bp.buf) + errCh <- copyErr }() go func() { bp := bufPool.Get().(*bp) defer bufPool.Put(bp) - _, innerErr := io.CopyBuffer(agentIo, userIo, bp.buf) - errCh <- innerErr + _, copyErr := io.CopyBuffer(agentIo, userIo, bp.buf) + errCh <- copyErr }() - return <-errCh } diff --git a/service/rpc/io_stream_lifecycle.go b/service/rpc/io_stream_lifecycle.go new file mode 100644 index 00000000..4e158167 --- /dev/null +++ b/service/rpc/io_stream_lifecycle.go @@ -0,0 +1,128 @@ +package rpc + +import ( + "context" + "errors" + "io" + "time" +) + +type ioStreamDetach struct { + stream *ioStreamContext + endpoints []io.ReadWriteCloser +} + +func detachStreamLocked(streamID string, retainedStream *ioStreamContext, streams map[string]*ioStreamContext) (ioStreamDetach, bool) { + current, live := streams[streamID] + if streamID == "" || !live || current != retainedStream { + return ioStreamDetach{}, false + } + retainedStream.revoke() + endpoints := make([]io.ReadWriteCloser, 0, 2) + if retainedStream.userIo != nil { + endpoints = append(endpoints, retainedStream.userIo) + } + if retainedStream.agentIo != nil { + endpoints = append(endpoints, retainedStream.agentIo) + } + delete(streams, streamID) + return ioStreamDetach{stream: retainedStream, endpoints: endpoints}, true +} + +func (s *NezhaHandler) detachExactStream(streamID string, retainedStream *ioStreamContext) error { + s.ioStreamMutex.Lock() + detached, ok := detachStreamLocked(streamID, retainedStream, s.ioStreams) + if !ok { + s.ioStreamMutex.Unlock() + return nil + } + s.publishIOStreamStateChangeLocked() + s.ioStreamMutex.Unlock() + + var closeErrors []error + for _, endpoint := range detached.endpoints { + if err := endpoint.Close(); err != nil { + closeErrors = append(closeErrors, err) + } + } + return errors.Join(closeErrors...) +} + +func (s *NezhaHandler) detachStreams(shouldDetach func(*ioStreamContext) bool) (int, error) { + s.ioStreamMutex.Lock() + detached := make([]ioStreamDetach, 0) + for streamID, stream := range s.ioStreams { + if !shouldDetach(stream) { + continue + } + item, ok := detachStreamLocked(streamID, stream, s.ioStreams) + if ok { + detached = append(detached, item) + } + } + if len(detached) > 0 { + s.publishIOStreamStateChangeLocked() + } + s.ioStreamMutex.Unlock() + + var closeErrors []error + for _, item := range detached { + // Registry publication must precede endpoint Close so Close implementations may reenter safely. + for _, endpoint := range item.endpoints { + if err := endpoint.Close(); err != nil { + closeErrors = append(closeErrors, err) + } + } + } + return len(detached), errors.Join(closeErrors...) +} + +func (s *NezhaHandler) CloseStream(streamID string) error { + _, err := s.detachStreams(func(stream *ioStreamContext) bool { + return stream != nil && streamID != "" && stream == s.ioStreams[streamID] + }) + return err +} + +func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamID string, timeout time.Duration) (io.ReadWriteCloser, bool) { + deadline := time.NewTimer(timeout) + defer deadline.Stop() + for { + s.ioStreamMutex.RLock() + stream, ok := s.ioStreams[streamID] + if ok && stream.agentIo != nil { + agentIo := stream.agentIo + s.ioStreamMutex.RUnlock() + return agentIo, true + } + if !ok { + s.ioStreamMutex.RUnlock() + return nil, false + } + revokedCh := stream.revokedCh + agentIoConnectCh := stream.agentIoConnectCh + stream.waitStartedOnce.Do(func() { close(stream.waitStartedCh) }) + s.ioStreamMutex.RUnlock() + select { + case <-ctx.Done(): + return nil, false + case <-deadline.C: + return nil, false + case <-revokedCh: + return nil, false + case <-agentIoConnectCh: + } + } +} + +func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) { + if serverID == 0 { + return + } + _, _ = s.detachStreams(func(stream *ioStreamContext) bool { return stream.targetServerID == serverID }) +} + +func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int { + revoked, _ := s.detachStreams(func(stream *ioStreamContext) bool { return stream.purpose == purpose }) + return revoked +} diff --git a/service/rpc/io_stream_lifecycle_test.go b/service/rpc/io_stream_lifecycle_test.go new file mode 100644 index 00000000..879bebed --- /dev/null +++ b/service/rpc/io_stream_lifecycle_test.go @@ -0,0 +1,259 @@ +package rpc + +import ( + "context" + "errors" + "io" + "testing" + "time" +) + +type lifecycleRWC struct { + closed chan struct{} +} + +type reenteringErrorRWC struct { + handler *NezhaHandler + streamID string + err error +} + +func (stream *reenteringErrorRWC) Read([]byte) (int, error) { return 0, io.EOF } +func (stream *reenteringErrorRWC) Write(data []byte) (int, error) { return len(data), nil } +func (stream *reenteringErrorRWC) Close() error { + if _, ok := stream.handler.StreamOwnership(stream.streamID); ok { + return errors.Join(stream.err, errors.New("stream remained registered during endpoint close")) + } + return stream.err +} + +func newLifecycleRWC() *lifecycleRWC { + return &lifecycleRWC{closed: make(chan struct{})} +} + +func (stream *lifecycleRWC) Read([]byte) (int, error) { return 0, io.EOF } +func (stream *lifecycleRWC) Write(data []byte) (int, error) { return len(data), nil } +func (stream *lifecycleRWC) Close() error { + select { + case <-stream.closed: + default: + close(stream.closed) + } + return nil +} + +func TestIOStreamValidCreateAttachCloseLifecycle(t *testing.T) { + handler := NewNezhaHandler() + user := newLifecycleRWC() + agent := newLifecycleRWC() + if err := handler.CreateStream("valid-lifecycle", 11, 22); err != nil { + t.Fatalf("Given a new stream, CreateStream failed: %v", err) + } + if err := handler.UserConnected("valid-lifecycle", user); err != nil { + t.Fatalf("Given a tracked stream, UserConnected failed: %v", err) + } + if err := handler.AgentConnected("valid-lifecycle", agent); err != nil { + t.Fatalf("Given a tracked stream, AgentConnected failed: %v", err) + } + if _, ok := handler.StreamOwnership("valid-lifecycle"); !ok { + t.Fatal("Then a valid attached stream must remain tracked") + } + if err := handler.CloseStream("valid-lifecycle"); err != nil { + t.Fatalf("When closing the valid stream, CloseStream failed: %v", err) + } + if _, ok := handler.StreamOwnership("valid-lifecycle"); ok { + t.Fatal("Then CloseStream must remove the tracked stream") + } +} + +func TestCreateStreamKeepsExistingStreamWhenIDIsDuplicated(t *testing.T) { + handler := NewNezhaHandler() + original := newLifecycleRWC() + if err := handler.CreateStream("duplicate-id", 11, 22); err != nil { + t.Fatalf("Given a new stream ID, CreateStream failed: %v", err) + } + if err := handler.AgentConnected("duplicate-id", original); err != nil { + t.Fatalf("Given a live stream, AgentConnected failed: %v", err) + } + + err := handler.CreateStream("duplicate-id", 33, 44) + if !errors.Is(err, ErrStreamAlreadyExists) { + t.Fatalf("When reusing a live ID, expected ErrStreamAlreadyExists, got %v", err) + } + owner, found := handler.StreamOwnership("duplicate-id") + if !found || owner != 11 { + t.Fatalf("Then the original stream ownership must remain, found=%v owner=%d", found, owner) + } + select { + case <-original.closed: + t.Fatal("Then duplicate creation must not close the original endpoint") + default: + } +} + +func TestAgentConnectedRejectsDuplicateEndpointWithoutReplacingLiveRelay(t *testing.T) { + handler := NewNezhaHandler() + first := newLifecycleRWC() + second := newLifecycleRWC() + if err := handler.CreateStream("agent-once", 11, 22); err != nil { + t.Fatalf("Given a new stream, CreateStream failed: %v", err) + } + if err := handler.AgentConnected("agent-once", first); err != nil { + t.Fatalf("Given no agent endpoint, AgentConnected failed: %v", err) + } + if err := handler.AgentConnected("agent-once", second); !errors.Is(err, ErrAgentStreamAlreadyConnected) { + t.Fatalf("When attaching a second agent endpoint, expected ErrAgentStreamAlreadyConnected, got %v", err) + } + endpoints, err := handler.GetStream("agent-once") + if err != nil || endpoints.agentIo != first { + t.Fatalf("Then the first endpoint must remain attached, err=%v", err) + } + select { + case <-second.closed: + default: + t.Fatal("Then the rejected duplicate endpoint must be closed") + } +} + +func TestCloseStreamWakesWaitForAgentAndAllowsSlotReuse(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("wait-close", 11, 22); err != nil { + t.Fatalf("Given a pending stream, CreateStream failed: %v", err) + } + stream, err := handler.GetStream("wait-close") + if err != nil { + t.Fatalf("Given a created stream, GetStream failed: %v", err) + } + result := make(chan bool, 1) + go func() { + _, ok := handler.WaitForAgent(context.Background(), "wait-close", time.Minute) + result <- ok + }() + select { + case <-stream.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not enter its blocking select") + } + + if err := handler.CloseStream("wait-close"); err != nil { + t.Fatalf("When closing a pending stream, CloseStream failed: %v", err) + } + select { + case ok := <-result: + if ok { + t.Fatal("Then WaitForAgent must report no attached agent") + } + case <-time.After(time.Second): + t.Fatal("Then CloseStream must wake WaitForAgent") + } + if err := handler.CreateStream("wait-close-reused", 11, 22); err != nil { + t.Fatalf("Then the released user/server slot must be reusable: %v", err) + } +} + +func TestRevokeStreamsForPurposeWakesWaitForAgentAndIsRepeatable(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStreamWithPurpose("revoke-wait", 0, 22, PurposeMCPTransfer); err != nil { + t.Fatalf("Given a pending MCP stream, CreateStream failed: %v", err) + } + stream, err := handler.GetStream("revoke-wait") + if err != nil { + t.Fatalf("Given a created stream, GetStream failed: %v", err) + } + result := make(chan bool, 1) + go func() { + _, ok := handler.WaitForAgent(context.Background(), "revoke-wait", time.Minute) + result <- ok + }() + select { + case <-stream.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not enter its blocking select") + } + + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 { + t.Fatalf("When revoking the purpose, expected one stream, got %d", revoked) + } + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 { + t.Fatalf("When repeating revocation, expected zero streams, got %d", revoked) + } + select { + case ok := <-result: + if ok { + t.Fatal("Then WaitForAgent must report no attached agent") + } + case <-time.After(time.Second): + t.Fatal("Then revocation must wake WaitForAgent") + } +} + +func TestCloseStreamDetachesBeforeReenteringEndpointCloseAndJoinsErrors(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("close-errors", 1, 1); err != nil { + t.Fatal(err) + } + firstErr := errors.New("first close error") + secondErr := errors.New("second close error") + if err := handler.UserConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: firstErr}); err != nil { + t.Fatal(err) + } + if err := handler.AgentConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: secondErr}); err != nil { + t.Fatal(err) + } + err := handler.CloseStream("close-errors") + if !errors.Is(err, firstErr) || !errors.Is(err, secondErr) { + t.Fatalf("close errors were not joined: %v", err) + } +} + +func TestStartStreamReturnsImmediatelyWhenRevoked(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("start-revoked", 1, 1); err != nil { + t.Fatal(err) + } + result := make(chan error, 1) + go func() { result <- handler.StartStream("start-revoked", time.Minute) }() + if revoked := handler.RevokeStreamsForPurpose(PurposeLegacy); revoked != 1 { + t.Fatalf("revoked streams: %d", revoked) + } + select { + case err := <-result: + if err == nil { + t.Fatal("revoked StartStream must return an error") + } + case <-time.After(time.Second): + t.Fatal("StartStream did not wake on revoke") + } +} + +func TestConcurrentCloseAndRevokePublishOneGeneration(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("single-generation", 1, 1); err != nil { + t.Fatal(err) + } + start := handler.SnapshotIOStreamState() + closeDone := make(chan struct{}) + revokeDone := make(chan struct{}) + go func() { + _ = handler.CloseStream("single-generation") + close(closeDone) + }() + go func() { + handler.RevokeStreamsForPurpose(PurposeLegacy) + close(revokeDone) + }() + select { + case <-closeDone: + case <-time.After(time.Second): + t.Fatal("CloseStream did not complete") + } + select { + case <-revokeDone: + case <-time.After(time.Second): + t.Fatal("RevokeStreamsForPurpose did not complete") + } + state := handler.SnapshotIOStreamState() + if state.Count != 0 || state.Generation != start.Generation+1 { + t.Fatalf("single detach publication: start=%+v final=%+v", start, state) + } +} diff --git a/service/rpc/io_stream_quota.go b/service/rpc/io_stream_quota.go new file mode 100644 index 00000000..dd5571cc --- /dev/null +++ b/service/rpc/io_stream_quota.go @@ -0,0 +1,62 @@ +package rpc + +import "errors" + +const ( + maxStreamsPerUser = 20 + maxStreamsPerServer = 40 +) + +var ( + ErrTooManyStreamsForUser = errors.New("too many concurrent streams for this user") + ErrTooManyStreamsForServer = errors.New("too many concurrent streams for this server") + ErrStreamAlreadyExists = errors.New("stream already exists") +) + +func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) error { + return s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy) +} + +func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + return s.createStreamLocked(streamId, creatorUserID, targetServerID, purpose) +} + +func (s *NezhaHandler) createStreamLocked(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error { + if _, exists := s.ioStreams[streamId]; exists { + // Stream IDs identify live relay ownership; never overwrite one or orphan its endpoint. + return ErrStreamAlreadyExists + } + + var perUser, perServer int + for _, ctx := range s.ioStreams { + if creatorUserID != 0 && ctx.creatorUserID == creatorUserID { + perUser++ + } + if ctx.targetServerID == targetServerID { + perServer++ + } + } + // creatorUserID==0 is a dashboard-internal stream (NAT, server transfer, + // MCP transfer); only end-user-initiated streams are capped per user, but + // every stream counts toward the per-server cap so one server cannot be + // flooded regardless of who opened the streams. + if creatorUserID != 0 && perUser >= maxStreamsPerUser { + return ErrTooManyStreamsForUser + } + if perServer >= maxStreamsPerServer { + return ErrTooManyStreamsForServer + } + + s.ioStreams[streamId] = newIOStreamContext(creatorUserID, targetServerID, purpose) + s.publishIOStreamStateChangeLocked() + return nil +} + +// StreamCount reports the registry size under the same lock used by lifecycle mutations. +func (s *NezhaHandler) StreamCount() int { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + return len(s.ioStreams) +} diff --git a/service/rpc/io_stream_quota_test.go b/service/rpc/io_stream_quota_test.go new file mode 100644 index 00000000..2f2ca8fb --- /dev/null +++ b/service/rpc/io_stream_quota_test.go @@ -0,0 +1,106 @@ +package rpc + +import ( + "errors" + "fmt" + "testing" +) + +func TestCreateStreamExactUserBoundary(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("quota-user-%d", i), 1, uint64(i+1)); err != nil { + t.Fatalf("20th user stream must succeed: %v", err) + } + } + if err := h.CreateStream("quota-user-21", 1, 100); !errors.Is(err, ErrTooManyStreamsForUser) { + t.Fatalf("21st user stream must be rejected: %v", err) + } +} + +func TestCreateStreamNormalUserEverydayUseSucceeds(t *testing.T) { + h := NewNezhaHandler() + if err := h.CreateStream("term", 7, 1); err != nil { + t.Fatal(err) + } + if err := h.CreateStream("fm", 7, 1); err != nil { + t.Fatal(err) + } +} + +func TestCreateStreamNormalUsersAreIndependent(t *testing.T) { + h := NewNezhaHandler() + for userID := uint64(1); userID <= 5; userID++ { + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("independent-%d-%d", userID, i), userID, 100+userID); err != nil { + t.Fatalf("user %d stream %d: %v", userID, i, err) + } + } + } +} + +func TestCreateStreamExemptsInternalStreamsFromPerUserCap(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerUser*3; i++ { + if err := h.CreateStream(fmt.Sprintf("internal-user-%d", i), 0, uint64(i+1)); err != nil { + t.Fatal(err) + } + } +} + +func TestCreateStreamInternalStreamsStillCountTowardPerServerCap(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("internal-server-%d", i), 0, 9); err != nil { + t.Fatal(err) + } + } + if err := h.CreateStream("internal-server-over", 0, 9); !errors.Is(err, ErrTooManyStreamsForServer) { + t.Fatal(err) + } +} + +func TestCreateStreamExactServerBoundary(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("quota-server-%d", i), uint64(i+1), 2); err != nil { + t.Fatalf("40th server stream must succeed: %v", err) + } + } + if err := h.CreateStream("quota-server-41", 100, 2); !errors.Is(err, ErrTooManyStreamsForServer) { + t.Fatalf("41st server stream must be rejected: %v", err) + } +} + +func TestCreateStreamReleasesUserAndServerSlots(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("reuse-user-%d", i), 1, uint64(i+10)); err != nil { + t.Fatalf("user setup stream %d failed: %v", i, err) + } + } + if !errors.Is(h.CreateStream("reuse-user-over", 1, 100), ErrTooManyStreamsForUser) { + t.Fatal("user cap was not enforced") + } + if err := h.CloseStream("reuse-user-0"); err != nil { + t.Fatal(err) + } + if err := h.CreateStream("reuse-user-new", 1, 101); err != nil { + t.Fatalf("closed user slot must be reusable: %v", err) + } + + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("reuse-server-%d", i), uint64(i+2), 2); err != nil { + t.Fatalf("server setup stream %d failed: %v", i, err) + } + } + if !errors.Is(h.CreateStream("reuse-server-over", 100, 2), ErrTooManyStreamsForServer) { + t.Fatal("server cap was not enforced") + } + if err := h.CloseStream("reuse-server-0"); err != nil { + t.Fatal(err) + } + if err := h.CreateStream("reuse-server-new", 100, 2); err != nil { + t.Fatalf("closed server slot must be reusable: %v", err) + } +} diff --git a/service/rpc/io_stream_registry.go b/service/rpc/io_stream_registry.go new file mode 100644 index 00000000..cc70c7fe --- /dev/null +++ b/service/rpc/io_stream_registry.go @@ -0,0 +1,85 @@ +package rpc + +import ( + "errors" + "io" +) + +var ErrAgentStreamAlreadyConnected = errors.New("agent stream already connected") + +func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID uint64) bool { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + ctx, ok := s.ioStreams[streamId] + return ok && ctx.targetServerID != 0 && ctx.targetServerID == agentServerID +} + +func (s *NezhaHandler) IsStreamAuthorizedForUser(streamId string, userID uint64, isAdmin bool) bool { + creator, found := s.StreamOwnership(streamId) + return found && (isAdmin || creator == userID) +} + +func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + ctx, ok := s.ioStreams[streamId] + if !ok { + return 0, false + } + return ctx.creatorUserID, true +} + +func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + ctx, ok := s.ioStreams[streamId] + if !ok { + return 0, false + } + return ctx.targetServerID, true +} + +func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + if ctx, ok := s.ioStreams[streamId]; ok { + return ctx, nil + } + return nil, errors.New("stream not found") +} + +func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error { + s.ioStreamMutex.Lock() + stream, ok := s.ioStreams[streamId] + if !ok { + s.ioStreamMutex.Unlock() + return errors.New("stream not found") + } + stream.userIo = userIo + s.ioStreamMutex.Unlock() + stream.userIoChOnce.Do(func() { close(stream.userIoConnectCh) }) + return nil +} + +func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error { + s.ioStreamMutex.Lock() + stream, ok := s.ioStreams[streamId] + if !ok { + s.ioStreamMutex.Unlock() + return errors.Join(errors.New("stream not found"), agentIo.Close()) + } + if stream.agentIo != nil { + s.ioStreamMutex.Unlock() + return errors.Join(ErrAgentStreamAlreadyConnected, agentIo.Close()) + } + stream.agentIo = agentIo + s.ioStreamMutex.Unlock() + stream.agentIoChOnce.Do(func() { close(stream.agentIoConnectCh) }) + return nil +} + +func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (io.ReadWriteCloser, io.ReadWriteCloser) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + return stream.userIo, stream.agentIo +} diff --git a/service/rpc/io_stream_rpc.go b/service/rpc/io_stream_rpc.go new file mode 100644 index 00000000..3d878fbe --- /dev/null +++ b/service/rpc/io_stream_rpc.go @@ -0,0 +1,58 @@ +package rpc + +import ( + "fmt" + "log" + "time" + + "github.com/nezhahq/nezha/pkg/grpcx" + pb "github.com/nezhahq/nezha/proto" +) + +func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error { + clientID, err := s.Auth.Check(stream.Context()) + if err != nil { + return err + } + id, err := stream.Recv() + if err != nil { + return err + } + if id == nil || !isValidIOStreamMagic(id.Data) { + return fmt.Errorf("invalid stream id") + } + streamID := string(id.Data[4:]) + if !s.IsStreamAuthorizedForAgent(streamID, clientID) { + return fmt.Errorf("stream not authorized for agent") + } + if _, err := s.GetStream(streamID); err != nil { + return err + } + wrapper := grpcx.NewIOStreamWrapper(stream) + keepaliveDone := make(chan struct{}) + go func() { + defer close(keepaliveDone) + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-wrapper.Context().Done(): + return + case <-wrapper.Done(): + return + case <-ticker.C: + if err := wrapper.SendKeepalive(); err != nil { + log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err) + return + } + } + } + }() + if err := s.AgentConnected(streamID, wrapper); err != nil { + _ = wrapper.Close() + return err + } + wrapper.Wait() + <-keepaliveDone + return nil +} diff --git a/service/rpc/io_stream_state.go b/service/rpc/io_stream_state.go new file mode 100644 index 00000000..4c79676a --- /dev/null +++ b/service/rpc/io_stream_state.go @@ -0,0 +1,95 @@ +package rpc + +import ( + "context" + "errors" +) + +var ErrInvalidIOStreamStateExpectation = errors.New("invalid IOStream state expectation") + +type IOStreamState struct { + Count int `json:"count"` + Generation uint64 `json:"generation"` +} + +type IOStreamStateExpectation struct { + // A pointer distinguishes an omitted count from an explicit zero count. + ExpectedCount *int `json:"expected_count,omitempty"` + PresentStreamID string `json:"present_stream_id,omitempty"` + AbsentStreamID string `json:"absent_stream_id,omitempty"` +} + +func ExpectedIOStreamCount(count int) *int { + return &count +} + +func (s IOStreamStateExpectation) validate() error { + if s.ExpectedCount == nil && s.PresentStreamID == "" && s.AbsentStreamID == "" { + return ErrInvalidIOStreamStateExpectation + } + if s.ExpectedCount != nil && *s.ExpectedCount < 0 { + return ErrInvalidIOStreamStateExpectation + } + if s.PresentStreamID != "" && s.PresentStreamID == s.AbsentStreamID { + return ErrInvalidIOStreamStateExpectation + } + return nil +} + +func (s *NezhaHandler) SnapshotIOStreamState() IOStreamState { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + return s.snapshotIOStreamStateLocked() +} + +func (s *NezhaHandler) snapshotIOStreamStateLocked() IOStreamState { + return IOStreamState{Count: len(s.ioStreams), Generation: s.ioStreamGeneration} +} + +func (s *NezhaHandler) ioStreamStateExpectationSatisfiedLocked(expectation IOStreamStateExpectation) bool { + if expectation.ExpectedCount != nil && len(s.ioStreams) != *expectation.ExpectedCount { + return false + } + if expectation.PresentStreamID != "" { + if _, exists := s.ioStreams[expectation.PresentStreamID]; !exists { + return false + } + } + if expectation.AbsentStreamID != "" { + if _, exists := s.ioStreams[expectation.AbsentStreamID]; exists { + return false + } + } + return true +} + +func (s *NezhaHandler) WaitForIOStreamState(ctx context.Context, expectation IOStreamStateExpectation) (IOStreamState, error) { + if err := expectation.validate(); err != nil { + return IOStreamState{}, err + } + for { + s.ioStreamMutex.RLock() + notify := s.ioStreamNotify + state := s.snapshotIOStreamStateLocked() + satisfied := s.ioStreamStateExpectationSatisfiedLocked(expectation) + observer := s.ioStreamWaitLockedHook + s.ioStreamMutex.RUnlock() + if observer != nil { + observer() + } + if satisfied { + return state, nil + } + select { + case <-ctx.Done(): + return IOStreamState{}, ctx.Err() + case <-notify: + } + } +} + +func (s *NezhaHandler) publishIOStreamStateChangeLocked() { + s.ioStreamGeneration++ + close(s.ioStreamNotify) + s.ioStreamNotify = make(chan struct{}) +} diff --git a/service/rpc/io_stream_state_expectation_test.go b/service/rpc/io_stream_state_expectation_test.go new file mode 100644 index 00000000..7344883d --- /dev/null +++ b/service/rpc/io_stream_state_expectation_test.go @@ -0,0 +1,142 @@ +package rpc + +import ( + "context" + "errors" + "strings" + "testing" +) + +func TestWaitForIOStreamStateRejectsZeroValueExpectation(t *testing.T) { + handler := NewNezhaHandler() + if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("zero-value expectation error: %v", err) + } +} + +func TestWaitForIOStreamStateAcceptsExplicitZeroCount(t *testing.T) { + handler := NewNezhaHandler() + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)}) + if err != nil { + t.Fatal(err) + } + if state != (IOStreamState{}) { + t.Fatalf("explicit zero state: %+v", state) + } +} + +func TestWaitForIOStreamStateAcceptsPresentOnlyExpectation(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("present-only", 1, 1); err != nil { + t.Fatal(err) + } + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "present-only"}) + if err != nil { + t.Fatal(err) + } + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("present-only state: %+v", state) + } +} + +func TestWaitForIOStreamStateRejectsSamePresentAndAbsentID(t *testing.T) { + handler := NewNezhaHandler() + _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "same", AbsentStreamID: "same"}) + if !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("same identity expectation error: %v", err) + } +} + +func TestWaitForIOStreamStateRequiresAllSpecifiedConditions(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("present", 1, 1); err != nil { + t.Fatal(err) + } + if err := handler.CreateStream("other", 1, 2); err != nil { + t.Fatal(err) + } + if err := handler.CreateStream("absent", 1, 3); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ + ExpectedCount: ExpectedIOStreamCount(2), + PresentStreamID: "present", + AbsentStreamID: "absent", + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("combined expectation cancellation: %v", err) + } +} + +func TestWaitForIOStreamStateAbsenceOnlyIgnoresUnrelatedStreams(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("unrelated", 1, 1); err != nil { + t.Fatal(err) + } + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{AbsentStreamID: "absent"}) + if err != nil { + t.Fatal(err) + } + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("absence-only state: %+v", state) + } +} + +func TestWaitForIOStreamStateRejectsNegativeCountWithoutPrivateID(t *testing.T) { + handler := NewNezhaHandler() + _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1), AbsentStreamID: "private-stream-id"}) + if !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("negative expectation error: %v", err) + } + if err != nil && strings.Contains(err.Error(), "private-stream-id") { + t.Fatalf("private stream ID leaked: %v", err) + } +} + +func TestWaitForIOStreamStateRequiresCombinedCountAndAbsence(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("present", 1, 1); err != nil { + t.Fatal(err) + } + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1), AbsentStreamID: "absent"}) + if err != nil { + t.Fatal(err) + } + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("combined expectation state: %+v", state) + } +} + +func TestWaitForIOStreamStateCancellationRemainsValidForUnsatisfiedExpectation(t *testing.T) { + handler := NewNezhaHandler() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) { + t.Fatalf("cancel error: %v", err) + } +} + +func TestWaitForIOStreamStateAlreadySatisfiedReturnsSnapshot(t *testing.T) { + handler := NewNezhaHandler() + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)}) + if err != nil { + t.Fatal(err) + } + if state != (IOStreamState{}) { + t.Fatalf("already satisfied state: %+v", state) + } +} + +func TestWaitForIOStreamStateRejectsInvalidAndHonorsCancellation(t *testing.T) { + handler := NewNezhaHandler() + if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1)}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("invalid count error: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) { + t.Fatalf("cancel error: %v", err) + } +} diff --git a/service/rpc/io_stream_state_test.go b/service/rpc/io_stream_state_test.go new file mode 100644 index 00000000..da0153e2 --- /dev/null +++ b/service/rpc/io_stream_state_test.go @@ -0,0 +1,68 @@ +package rpc + +import ( + "errors" + "testing" +) + +func TestIOStreamStateSnapshotAndGeneration(t *testing.T) { + handler := NewNezhaHandler() + initial := handler.SnapshotIOStreamState() + if initial.Count != 0 || initial.Generation != 0 { + t.Fatalf("unexpected initial state: %+v", initial) + } + if err := handler.CreateStream("state-stream", 1, 1); err != nil { + t.Fatal(err) + } + created := handler.SnapshotIOStreamState() + if created.Count != 1 || created.Generation != 1 { + t.Fatalf("unexpected created state: %+v", created) + } + if err := handler.CreateStream("state-stream", 2, 2); !errors.Is(err, ErrStreamAlreadyExists) { + t.Fatalf("duplicate create error: %v", err) + } + if got := handler.SnapshotIOStreamState(); got != created { + t.Fatalf("duplicate create changed state: %+v", got) + } + if err := handler.CloseStream("unknown"); err != nil { + t.Fatal(err) + } + if got := handler.SnapshotIOStreamState(); got != created { + t.Fatalf("unknown close changed state: %+v", got) + } +} + +func TestIOStreamStateRevocationPublishesOncePerBatch(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStreamWithPurpose("purpose-a", 0, 1, PurposeMCPTransfer); err != nil { + t.Fatal(err) + } + if err := handler.CreateStreamWithPurpose("purpose-b", 0, 1, PurposeMCPTransfer); err != nil { + t.Fatal(err) + } + if err := handler.CreateStream("server-a", 0, 2); err != nil { + t.Fatal(err) + } + before := handler.SnapshotIOStreamState() + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 2 { + t.Fatalf("revoked purpose streams: %d", revoked) + } + afterPurpose := handler.SnapshotIOStreamState() + if afterPurpose.Generation != before.Generation+1 || afterPurpose.Count != 1 { + t.Fatalf("purpose revocation state: before=%+v after=%+v", before, afterPurpose) + } + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 { + t.Fatalf("repeat purpose revocation: %d", revoked) + } + if got := handler.SnapshotIOStreamState(); got != afterPurpose { + t.Fatalf("empty purpose revocation changed state: %+v", got) + } + handler.RevokeStreamsForServer(2) + if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 || got.Count != 0 { + t.Fatalf("server revocation state: %+v", got) + } + handler.RevokeStreamsForServer(2) + if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 { + t.Fatalf("empty server revocation changed generation: %+v", got) + } +} diff --git a/service/rpc/io_stream_state_wait_test.go b/service/rpc/io_stream_state_wait_test.go new file mode 100644 index 00000000..825e5598 --- /dev/null +++ b/service/rpc/io_stream_state_wait_test.go @@ -0,0 +1,220 @@ +package rpc + +import ( + "context" + "sync" + "testing" + "time" +) + +func TestWaitForIOStreamStateWakesOnCloseAndAbsence(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("wait-state", 1, 1); err != nil { + t.Fatal(err) + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "wait-state"}) + if err != nil { + t.Errorf("wait failed: %v", err) + return + } + result <- state + }() + if err := handler.CloseStream("wait-state"); err != nil { + t.Fatal(err) + } + state := <-result + if state.Count != 0 || state.Generation != 2 { + t.Fatalf("unexpected waited state: %+v", state) + } +} + +func TestWaitForIOStreamStateCreateWakeUsesCapturedNotification(t *testing.T) { + handler := NewNezhaHandler() + waitReady := make(chan struct{}) + handler.ioStreamWaitLockedHook = func() { + select { + case <-waitReady: + default: + close(waitReady) + } + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err == nil { + result <- state + } + }() + select { + case <-waitReady: + case <-time.After(time.Second): + t.Fatal("waiter did not capture its notification channel") + } + if err := handler.CreateStream("create-wake", 1, 1); err != nil { + t.Fatal(err) + } + select { + case state := <-result: + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("unexpected created state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("create did not wake waiter") + } +} + +func TestWaitForIOStreamStateCloseWakeUsesCapturedNotification(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("close-wake", 1, 1); err != nil { + t.Fatal(err) + } + waitReady := make(chan struct{}) + handler.ioStreamWaitLockedHook = func() { + select { + case <-waitReady: + default: + close(waitReady) + } + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "close-wake"}) + if err == nil { + result <- state + } + }() + select { + case <-waitReady: + case <-time.After(time.Second): + t.Fatal("waiter did not capture its notification channel") + } + if err := handler.CloseStream("close-wake"); err != nil { + t.Fatal(err) + } + select { + case state := <-result: + if state.Count != 0 || state.Generation != 2 { + t.Fatalf("unexpected closed state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("close did not wake waiter") + } +} + +func TestWaitForIOStreamStateDoesNotMissMutationBetweenSnapshotAndWait(t *testing.T) { + handler := NewNezhaHandler() + hookCalled := make(chan struct{}) + mutationDone := make(chan error, 1) + var hookOnce sync.Once + handler.ioStreamWaitLockedHook = func() { + hookOnce.Do(func() { + close(hookCalled) + go func() { + mutationDone <- handler.CreateStream("lost-wakeup", 1, 1) + }() + }) + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err == nil { + result <- state + } + }() + select { + case <-hookCalled: + case <-time.After(time.Second): + t.Fatal("waiter did not reach deterministic mutation seam") + } + select { + case err := <-mutationDone: + if err != nil { + t.Fatal(err) + } + case <-time.After(time.Second): + t.Fatal("mutation did not complete") + } + select { + case state := <-result: + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("unexpected mutation state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("waiter missed mutation published during wait setup") + } +} + +func TestWaitForIOStreamStateConcurrentCreateCloseWaiters(t *testing.T) { + handler := NewNezhaHandler() + created := make(chan IOStreamState, 1) + closed := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err == nil { + created <- state + } + }() + if err := handler.CreateStream("concurrent", 1, 1); err != nil { + t.Fatal(err) + } + select { + case state := <-created: + if state.Count != 1 { + t.Fatalf("created waiter state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("created waiter did not wake") + } + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "concurrent"}) + if err == nil { + closed <- state + } + }() + if err := handler.CloseStream("concurrent"); err != nil { + t.Fatal(err) + } + select { + case state := <-closed: + if state.Count != 0 { + t.Fatalf("closed waiter state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("closed waiter did not wake") + } +} + +func TestWaitForIOStreamStateDoesNotAcceptUnrelatedSameCountForPresentID(t *testing.T) { + handler := NewNezhaHandler() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + go func() { + _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ + ExpectedCount: ExpectedIOStreamCount(1), + PresentStreamID: "wanted", + }) + result <- err + }() + if err := handler.CreateStream("unrelated", 1, 1); err != nil { + t.Fatal(err) + } + select { + case err := <-result: + if err == nil { + t.Fatal("same-count unrelated stream satisfied identity expectation") + } + default: + } + cancel() + select { + case err := <-result: + if err == nil { + t.Fatal("identity waiter unexpectedly succeeded") + } + case <-time.After(time.Second): + t.Fatal("identity waiter did not observe cancellation") + } +} diff --git a/service/rpc/io_stream_test.go b/service/rpc/io_stream_test.go index 55927363..f329da99 100644 --- a/service/rpc/io_stream_test.go +++ b/service/rpc/io_stream_test.go @@ -1,8 +1,6 @@ package rpc import ( - "errors" - "fmt" "io" "reflect" "testing" @@ -100,130 +98,6 @@ func TestIOStream(t *testing.T) { }) } -// The WebSocket stream endpoints (terminal / fm) were unbounded: an -// authenticated member could open thousands of streams, each spawning -// goroutines, a 1 MiB buffer, and an agent-side PTY, exhausting dashboard and -// agent resources (GHSA-jg62-j5h6-8mpq). CreateStream now caps concurrent -// streams per user and per server. These tests pin the caps and the -// dashboard-internal (uid==0) exemption. - -// Baseline: a normal operator opening a terminal and a file-manager session -// against one server (the everyday case) must always succeed — the cap exists -// to stop floods, not to interfere with ordinary use. -func TestCreateStreamNormalUserEverydayUseSucceeds(t *testing.T) { - h := NewNezhaHandler() - const uid, serverID = uint64(7), uint64(1) - - if err := h.CreateStream("term", uid, serverID); err != nil { - t.Fatalf("opening a terminal must succeed for a normal user, got %v", err) - } - if err := h.CreateStream("fm", uid, serverID); err != nil { - t.Fatalf("opening a file manager alongside a terminal must succeed, got %v", err) - } -} - -// Several normal users working at the same time must not interfere: one user's -// streams do not consume another user's per-user budget. -func TestCreateStreamNormalUsersAreIndependent(t *testing.T) { - h := NewNezhaHandler() - - for u := uint64(1); u <= 5; u++ { - for i := 0; i < maxStreamsPerUser; i++ { - id := fmt.Sprintf("u%d-s%d", u, i) - if err := h.CreateStream(id, u, 100+u); err != nil { - t.Fatalf("user %d stream %d must succeed; per-user budgets must be independent, got %v", u, i, err) - } - } - } -} - -func TestCreateStreamEnforcesPerUserCap(t *testing.T) { - h := NewNezhaHandler() - const uid = uint64(42) - - for i := 0; i < maxStreamsPerUser; i++ { - if err := h.CreateStream(fmt.Sprintf("u-%d", i), uid, uint64(i)); err != nil { - t.Fatalf("stream %d within the per-user cap must succeed, got %v", i, err) - } - } - - err := h.CreateStream("u-over", uid, 9999) - if !errors.Is(err, ErrTooManyStreamsForUser) { - t.Fatalf("the (maxStreamsPerUser+1)-th stream must be rejected with ErrTooManyStreamsForUser, got %v", err) - } -} - -func TestCreateStreamEnforcesPerServerCap(t *testing.T) { - h := NewNezhaHandler() - const serverID = uint64(7) - - for i := 0; i < maxStreamsPerServer; i++ { - if err := h.CreateStream(fmt.Sprintf("s-%d", i), uint64(i+1), serverID); err != nil { - t.Fatalf("stream %d within the per-server cap must succeed, got %v", i, err) - } - } - - err := h.CreateStream("s-over", 99999, serverID) - if !errors.Is(err, ErrTooManyStreamsForServer) { - t.Fatalf("the (maxStreamsPerServer+1)-th stream to one server must be rejected with ErrTooManyStreamsForServer, got %v", err) - } -} - -// Dashboard-internal streams (NAT, server transfer, MCP transfer) pass -// creatorUserID==0. They must NOT be capped per user, or those features would -// throttle themselves; but they must still count toward the per-server cap so -// no single server can be flooded regardless of the originating path. -func TestCreateStreamExemptsInternalStreamsFromPerUserCap(t *testing.T) { - h := NewNezhaHandler() - - for i := 0; i < maxStreamsPerUser*3; i++ { - if err := h.CreateStream(fmt.Sprintf("internal-%d", i), 0, uint64(i)); err != nil { - t.Fatalf("internal stream %d (uid==0) must never hit the per-user cap, got %v", i, err) - } - } -} - -func TestCreateStreamInternalStreamsStillCountTowardPerServerCap(t *testing.T) { - h := NewNezhaHandler() - const serverID = uint64(3) - - for i := 0; i < maxStreamsPerServer; i++ { - if err := h.CreateStream(fmt.Sprintf("internal-s-%d", i), 0, serverID); err != nil { - t.Fatalf("internal stream %d within the per-server cap must succeed, got %v", i, err) - } - } - - err := h.CreateStream("internal-s-over", 0, serverID) - if !errors.Is(err, ErrTooManyStreamsForServer) { - t.Fatalf("internal streams must still be subject to the per-server cap, got %v", err) - } -} - -// Closing a stream must free its slot so a user who hit the cap can open new -// streams after old ones end — otherwise normal churn would permanently lock -// a user out. -func TestCreateStreamFreesSlotAfterClose(t *testing.T) { - h := NewNezhaHandler() - const uid = uint64(55) - - for i := 0; i < maxStreamsPerUser; i++ { - if err := h.CreateStream(fmt.Sprintf("c-%d", i), uid, 1); err != nil { - t.Fatalf("setup stream %d must succeed, got %v", i, err) - } - } - if err := h.CreateStream("c-over", uid, 1); !errors.Is(err, ErrTooManyStreamsForUser) { - t.Fatalf("expected per-user cap to be hit, got %v", err) - } - - if err := h.CloseStream("c-0"); err != nil { - t.Fatalf("CloseStream failed: %v", err) - } - - if err := h.CreateStream("c-after-close", uid, 1); err != nil { - t.Fatalf("after closing one stream the user must be able to open another, got %v", err) - } -} - func newPipeReadWriter() io.ReadWriteCloser { r, w := io.Pipe() return struct { diff --git a/service/rpc/nezha.go b/service/rpc/nezha.go index c006ee49..5ad70c67 100644 --- a/service/rpc/nezha.go +++ b/service/rpc/nezha.go @@ -5,13 +5,10 @@ import ( "errors" "fmt" "log" - "net" "sync" "time" "github.com/jinzhu/copier" - geoipx "github.com/nezhahq/nezha/pkg/geoip" - "github.com/nezhahq/nezha/pkg/grpcx" "github.com/nezhahq/nezha/pkg/tsdb" "github.com/nezhahq/nezha/model" @@ -24,17 +21,36 @@ var _ pb.NezhaServiceServer = (*NezhaHandler)(nil) var NezhaHandlerSingleton *NezhaHandler type NezhaHandler struct { - Auth *authHandler - ioStreams map[string]*ioStreamContext - ioStreamMutex *sync.RWMutex + Auth *authHandler + ioStreams map[string]*ioStreamContext + ioStreamMutex *sync.RWMutex + ioStreamGeneration uint64 + ioStreamNotify chan struct{} + ioStreamWaitLockedHook func() + // Capability authorization and exact stream deletion share ioStreamMutex to avoid TOCTOU. + agentCompatCapabilities agentCompatCapabilityState +} + +type serverMetricsWriter func(*tsdb.ServerMetrics) error + +var writeServerMetrics serverMetricsWriter = writeServerMetricsToTSDB + +func writeServerMetricsToTSDB(metrics *tsdb.ServerMetrics) error { + if !singleton.TSDBEnabled() { + return nil + } + return singleton.TSDBShared.WriteServerMetrics(metrics) } func NewNezhaHandler() *NezhaHandler { - return &NezhaHandler{ - Auth: &authHandler{}, - ioStreamMutex: new(sync.RWMutex), - ioStreams: make(map[string]*ioStreamContext), + handler := &NezhaHandler{ + Auth: &authHandler{}, + ioStreamMutex: new(sync.RWMutex), + ioStreams: make(map[string]*ioStreamContext), + ioStreamNotify: make(chan struct{}), } + handler.initializeAgentCompatCapabilities() + return handler } // attachRequestTaskStream resolves the server for clientID and publishes the @@ -161,84 +177,91 @@ func (s *NezhaHandler) ReportSystemState(stream pb.NezhaService_ReportSystemStat if err != nil { return err } + server, ok := singleton.ServerShared.Get(clientID) + if !ok || server == nil { + return errors.New("server not found") + } + lease := server.AttachStateStream(stream) + defer lease.Clear() var state *pb.State + var stateCount uint64 for { state, err = stream.Recv() if err != nil { log.Printf("NEZHA>> ReportSystemState error: %v, clientID: %d\n", err, clientID) return err } + stateCount++ innerState := model.PB2State(state) - server, ok := singleton.ServerShared.Get(clientID) - if !ok || server == nil { - return errors.New("server not found") - } - - server.LastActive = time.Now() - server.State = &innerState - - if singleton.TSDBEnabled() { - maxTemp := 0.0 - for _, t := range innerState.Temperatures { - if t.Temperature > maxTemp { - maxTemp = t.Temperature + lastActive := time.Now() + accepted := lease.UpdateStateWithSideEffect(&innerState, lastActive, func() error { + { + maxTemp := 0.0 + for _, t := range innerState.Temperatures { + if t.Temperature > maxTemp { + maxTemp = t.Temperature + } + } + maxGPU := 0.0 + for _, g := range innerState.GPU { + if g > maxGPU { + maxGPU = g + } + } + if err := writeServerMetrics(&tsdb.ServerMetrics{ + ServerID: clientID, + Timestamp: lastActive, + CPU: innerState.CPU, + MemUsed: innerState.MemUsed, + SwapUsed: innerState.SwapUsed, + DiskUsed: innerState.DiskUsed, + NetInSpeed: innerState.NetInSpeed, + NetOutSpeed: innerState.NetOutSpeed, + NetInTransfer: innerState.NetInTransfer, + NetOutTransfer: innerState.NetOutTransfer, + Load1: innerState.Load1, + Load5: innerState.Load5, + Load15: innerState.Load15, + TCPConnCount: innerState.TcpConnCount, + UDPConnCount: innerState.UdpConnCount, + ProcessCount: innerState.ProcessCount, + Temperature: maxTemp, + Uptime: innerState.Uptime, + GPU: maxGPU, + }); err != nil { + log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err) } } - maxGPU := 0.0 - for _, g := range innerState.GPU { - if g > maxGPU { - maxGPU = g - } - } - if err := singleton.TSDBShared.WriteServerMetrics(&tsdb.ServerMetrics{ - ServerID: clientID, - Timestamp: time.Now(), - CPU: innerState.CPU, - MemUsed: innerState.MemUsed, - SwapUsed: innerState.SwapUsed, - DiskUsed: innerState.DiskUsed, - NetInSpeed: innerState.NetInSpeed, - NetOutSpeed: innerState.NetOutSpeed, - NetInTransfer: innerState.NetInTransfer, - NetOutTransfer: innerState.NetOutTransfer, - Load1: innerState.Load1, - Load5: innerState.Load5, - Load15: innerState.Load15, - TCPConnCount: innerState.TcpConnCount, - UDPConnCount: innerState.UdpConnCount, - ProcessCount: innerState.ProcessCount, - Temperature: maxTemp, - Uptime: innerState.Uptime, - GPU: maxGPU, - }); err != nil { - log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err) - } + return nil + }) + if !accepted { + return errors.New("state stream superseded") } - // 应对 dashboard / agent 重启的情况，如果从未记录过，先打点，等到小时时间点时入库 - if server.PrevTransferInSnapshot == 0 || server.PrevTransferOutSnapshot == 0 { - server.PrevTransferInSnapshot = state.NetInTransfer - server.PrevTransferOutSnapshot = state.NetOutTransfer + if err := notifyStateReceived(clientID, server.UUID, lease.Generation(), stateCount); err != nil { + return err + } + if err := notifyReceiptAccepted(clientID, server.UUID, lease.Generation(), stateCount); err != nil { + return err } - if err = stream.Send(&pb.Receipt{Proced: true}); err != nil { return err } } } -func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error { +func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) (model.HostReportResult, error) { var clientID uint64 var err error if clientID, err = s.Auth.Check(c); err != nil { - return err + return model.HostReportResult{}, err } host := model.PB2Host(r) server, ok := singleton.ServerShared.Get(clientID) if !ok || server == nil { - return errors.New("server not found") + return model.HostReportResult{}, errors.New("server not found") } /** @@ -246,170 +269,23 @@ func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error { * 当 agent 重启时，bootTime 变大，agent 端会先上报 host 信息，然后上报 state 信息 * 这时可以借助上报顺序的空档，立即记录停机前的数据并重置 Prev* 数据，并由接下来的 state 方法重新赋值 */ - if !server.LastActive.IsZero() && host.BootTime > server.Host.BootTime { - singleton.RecordTransferHourlyUsage(server) - server.PrevTransferInSnapshot = 0 - server.PrevTransferOutSnapshot = 0 - } - - server.Host = &host - return nil + return server.RuntimeHandle().ApplyHostReport(&host, time.Now(), singleton.PersistTransfer) } func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Receipt, error) { - if err := s.onReportSystemInfo(c, r); err != nil { + if _, err := s.onReportSystemInfo(c, r); err != nil { return nil, err } return &pb.Receipt{Proced: true}, nil } func (s *NezhaHandler) ReportSystemInfo2(c context.Context, r *pb.Host) (*pb.Uint64Receipt, error) { - if err := s.onReportSystemInfo(c, r); err != nil { + result, err := s.onReportSystemInfo(c, r) + if err != nil { + return nil, err + } + if err := notifyInfo2(result.ServerID, result.UUID); err != nil { return nil, err } return &pb.Uint64Receipt{Data: singleton.DashboardBootTime}, nil } - -func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error { - clientID, err := s.Auth.Check(stream.Context()) - if err != nil { - return err - } - id, err := stream.Recv() - if err != nil { - return err - } - - // ff05ff05 是 Nezha 的魔数，用于标识流 ID。校验由 isValidIOStreamMagic 完成， - // 历史 inline 检查曾因 && 短路放过几乎全部非魔数 payload (byte0==0xff 即通过)。 - if id == nil || !isValidIOStreamMagic(id.Data) { - 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 6661d6a（user 侧归属校验）的对偶补丁。先校验后启 keepalive， - // 避免未授权 agent 触发悬空 goroutine 持续向其发心跳。 - if !s.IsStreamAuthorizedForAgent(streamId, clientID) { - return fmt.Errorf("stream not authorized for agent") - } - - if _, err := s.GetStream(streamId); err != nil { - return err - } - iw := grpcx.NewIOStreamWrapper(stream) - - // Keepalive MUST go through the wrapper so it shares the same sendMu as - // MCP fs.transfer / terminal / fm Writers. Calling stream.Send directly - // here used to race those Writers — grpc-go forbids concurrent SendMsg - // on the same stream. The wrapper's sendMu is the dashboard-side dual - // of agent/cmd/agent/mcp_fs_transfer.go's serialIOStreamSender. - keepaliveDone := make(chan struct{}) - go func() { - defer close(keepaliveDone) - ticker := time.NewTicker(time.Second * 30) - defer ticker.Stop() - for { - select { - case <-iw.Context().Done(): - return - case <-iw.Done(): - // 业务侧（CloseStream / RevokeStreamsForPurpose）调过 - // iw.Close()。即便底层 gRPC stream context 尚未取消，也 - // 必须立刻收手——否则要再等一整个 30s tick，handler 在 - // iw.Wait() 之后又得多等一拍 keepaliveDone 才能返回。 - return - case <-ticker.C: - if err := iw.SendKeepalive(); err != nil { - log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err) - return - } - } - } - }() - - if err := s.AgentConnected(streamId, iw); err != nil { - return err - } - iw.Wait() - <-keepaliveDone - return nil -} - -func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, error) { - var clientID uint64 - var err error - if clientID, err = s.Auth.Check(c); err != nil { - return nil, err - } - - geoip := model.PB2GeoIP(r) - use6 := r.GetUse6() - - if geoip.IP.IPv4Addr == "" && geoip.IP.IPv6Addr == "" { - ip, _ := c.Value(model.CtxKeyRealIP{}).(string) - if ip == "" { - ip, _ = c.Value(model.CtxKeyConnectingIP{}).(string) - } - geoip.IP.IPv4Addr = ip - } - - joinedIP := geoip.IP.Join() - - server, ok := singleton.ServerShared.Get(clientID) - if !ok || server == nil { - return nil, fmt.Errorf("server not found") - } - - // 检查并更新DDNS - if server.EnableDDNS && joinedIP != "" && - (server.GeoIP == nil || server.GeoIP.IP != geoip.IP) { - ipv4 := geoip.IP.IPv4Addr - ipv6 := geoip.IP.IPv6Addr - - if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: ipv4, IPv6Addr: ipv6}); err != nil { - log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", err, server.ID) - } - } - - // 发送IP变动通知 - if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification && - ((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) || - (singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) && - server.GeoIP.IP.Join() != "" && - joinedIP != "" && - server.GeoIP.IP != geoip.IP { - - singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID, - fmt.Sprintf( - "[%s] %s, %s => %s", - singleton.Localizer.T("IP Changed"), - server.Name, singleton.IPDesensitize(server.GeoIP.IP.Join()), - singleton.IPDesensitize(joinedIP), - ), - "") - } - - // 根据内置数据库查询 IP 地理位置 - var ip string - if geoip.IP.IPv6Addr != "" && (use6 || geoip.IP.IPv4Addr == "") { - ip = geoip.IP.IPv6Addr - } else { - ip = geoip.IP.IPv4Addr - } - - netIP := net.ParseIP(ip) - location, err := geoipx.Lookup(netIP) - if err != nil { - log.Printf("NEZHA>> geoip.Lookup: %v", err) - } - geoip.CountryCode = location - - // 将地区码写入到 Host - server.GeoIP = &geoip - - return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil -} diff --git a/service/rpc/request_task_security_test.go b/service/rpc/request_task_security_test.go index 2ae3d997..3a86c4e2 100644 --- a/service/rpc/request_task_security_test.go +++ b/service/rpc/request_task_security_test.go @@ -11,6 +11,7 @@ import ( "gorm.io/gorm" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" pb "github.com/nezhahq/nezha/proto" "github.com/nezhahq/nezha/service/singleton" ) @@ -257,7 +258,10 @@ func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, cron originalDB := singleton.DB originalConf := singleton.Conf originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalNotification := singleton.NotificationShared originalServerShared := singleton.ServerShared + originalServiceSentinel := singleton.ServiceSentinelShared originalCronShared := singleton.CronShared originalUserInfoMap := singleton.UserInfoMap originalAgentSecretToUserID := singleton.AgentSecretToUserId @@ -275,6 +279,8 @@ func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, cron singleton.DB = db singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}} singleton.Loc = time.UTC + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest() if err := singleton.DB.AutoMigrate(model.Server{}, model.Cron{}); err != nil { t.Fatal(err) } @@ -297,13 +303,14 @@ func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, cron singleton.CronShared = singleton.NewCronClass() t.Cleanup(func() { - if singleton.CronShared != nil && singleton.CronShared.Cron != nil { - singleton.CronShared.Stop() - } - sqlDB.Close() + singleton.CronShared.Close() + _ = sqlDB.Close() singleton.DB = originalDB singleton.Conf = originalConf singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.NotificationShared = originalNotification + singleton.ServiceSentinelShared = originalServiceSentinel singleton.ServerShared = originalServerShared singleton.CronShared = originalCronShared singleton.UserLock.Lock() diff --git a/service/rpc/state_metrics_test.go b/service/rpc/state_metrics_test.go new file mode 100644 index 00000000..07e3165b --- /dev/null +++ b/service/rpc/state_metrics_test.go @@ -0,0 +1,42 @@ +package rpc + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/tsdb" +) + +func TestStateMetricsWriterRunsOnlyForCurrentGeneration(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + oldLease := server.AttachStateStream(stateGenerationStream{}) + newLease := server.AttachStateStream(stateGenerationStream{}) + oldCalls := 0 + newCalls := 0 + oldWriter := writeServerMetrics + writeServerMetrics = func(*tsdb.ServerMetrics) error { + newCalls++ + return nil + } + t.Cleanup(func() { writeServerMetrics = oldWriter }) + + // When + oldAccepted := server.UpdateStateIfCurrentWithSideEffect(oldLease, &model.HostState{Uptime: 11}, time.Unix(100, 0), func() error { + oldCalls++ + return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(100, 0)}) + }) + newAccepted := server.UpdateStateIfCurrentWithSideEffect(newLease, &model.HostState{Uptime: 22}, time.Unix(200, 0), func() error { + return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(200, 0)}) + }) + + // Then + require.False(t, oldAccepted) + require.True(t, newAccepted) + require.Zero(t, oldCalls) + require.Equal(t, 1, newCalls) +} diff --git a/service/rpc/state_stream_generation_test.go b/service/rpc/state_stream_generation_test.go new file mode 100644 index 00000000..04f62310 --- /dev/null +++ b/service/rpc/state_stream_generation_test.go @@ -0,0 +1,206 @@ +package rpc + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/tsdb" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestReportSystemState_HandlerWaitsForMetricsBeforeReceipt(t *testing.T) { + // Given + reporter := requestTaskSecurityServer(9, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{200: {Role: model.RoleMember}}, map[string]uint64{"reporter-secret": 200}) + stop := make(chan struct{}) + stream := &stateGenerationHandlerStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)), + states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: stop, + } + stream.states <- &pb.State{Uptime: 44} + metricsStarted := make(chan *tsdb.ServerMetrics, 1) + metricsRelease := make(chan struct{}) + oldWriter := writeServerMetrics + writeServerMetrics = func(metrics *tsdb.ServerMetrics) error { + metricsStarted <- metrics + <-metricsRelease + return nil + } + t.Cleanup(func() { writeServerMetrics = oldWriter }) + + // When + done := make(chan error, 1) + go func() { done <- NewNezhaHandler().ReportSystemState(stream) }() + metrics := <-metricsStarted + select { + case <-stream.receipts: + t.Fatal("receipt sent before metrics writer completed") + default: + } + close(metricsRelease) + + // Then + require.Equal(t, reporter.ID, metrics.ServerID) + require.Equal(t, uint64(44), metrics.Uptime) + require.NotNil(t, <-stream.receipts) + current, ok := singleton.ServerShared.Get(reporter.ID) + require.True(t, ok) + require.Equal(t, current.RuntimeSnapshot().LastActive, metrics.Timestamp) + close(stop) + require.ErrorIs(t, <-done, context.Canceled) +} + +type stateGenerationStream struct{} + +func (stateGenerationStream) Send(*pb.Receipt) error { return nil } +func (stateGenerationStream) Recv() (*pb.State, error) { return nil, nil } +func (stateGenerationStream) SetHeader(metadata.MD) error { return nil } +func (stateGenerationStream) SendHeader(metadata.MD) error { return nil } +func (stateGenerationStream) SetTrailer(metadata.MD) {} +func (stateGenerationStream) Context() context.Context { return context.Background() } +func (stateGenerationStream) SendMsg(any) error { return nil } +func (stateGenerationStream) RecvMsg(any) error { return nil } + +type stateGenerationHandlerStream struct { + ctx context.Context + states chan *pb.State + receipts chan *pb.Receipt + stop <-chan struct{} +} + +func (s *stateGenerationHandlerStream) Send(receipt *pb.Receipt) error { + s.receipts <- receipt + return nil +} + +func (s *stateGenerationHandlerStream) Recv() (*pb.State, error) { + select { + case state := <-s.states: + return state, nil + case <-s.stop: + return nil, context.Canceled + } +} + +func (s *stateGenerationHandlerStream) SetHeader(metadata.MD) error { return nil } +func (s *stateGenerationHandlerStream) SendHeader(metadata.MD) error { return nil } +func (s *stateGenerationHandlerStream) SetTrailer(metadata.MD) {} +func (s *stateGenerationHandlerStream) Context() context.Context { return s.ctx } +func (s *stateGenerationHandlerStream) SendMsg(any) error { return nil } +func (s *stateGenerationHandlerStream) RecvMsg(any) error { return nil } + +func TestReportSystemState_HandlerOldStreamCannotClearNewerState(t *testing.T) { + // Given + reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + oldStop := make(chan struct{}) + newStop := make(chan struct{}) + oldStream := &stateGenerationHandlerStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)), + states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: oldStop, + } + newStream := &stateGenerationHandlerStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)), + states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: newStop, + } + oldStream.states <- &pb.State{Uptime: 11} + newStream.states <- &pb.State{Uptime: 22} + handler := NewNezhaHandler() + oldDone := make(chan error, 1) + newDone := make(chan error, 1) + go func() { oldDone <- handler.ReportSystemState(oldStream) }() + <-oldStream.receipts + + // When + go func() { newDone <- handler.ReportSystemState(newStream) }() + <-newStream.receipts + close(newStop) + require.ErrorIs(t, <-newDone, context.Canceled) + close(oldStop) + require.ErrorIs(t, <-oldDone, context.Canceled) + + // Then + server, ok := singleton.ServerShared.Get(reporter.ID) + require.True(t, ok) + require.Equal(t, uint64(22), server.State.Uptime) + require.True(t, server.LastActive.IsZero()) +} + +func TestReportSystemState_OldStreamCannotUpdateNewerGeneration(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + oldStream := stateGenerationStream{} + newStream := stateGenerationStream{} + oldLease := server.AttachStateStream(oldStream) + updateGate := make(chan struct{}) + updateDone := make(chan bool, 1) + oldState := &model.HostState{Uptime: 11} + newState := &model.HostState{Uptime: 22} + oldTime := time.Unix(100, 0) + newTime := time.Unix(200, 0) + var waitGroup sync.WaitGroup + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + <-updateGate + updateDone <- server.UpdateStateIfCurrent(oldLease, oldState, oldTime) + }() + + // When + newLease := server.AttachStateStream(newStream) + close(updateGate) + oldUpdateAccepted := <-updateDone + newUpdateAccepted := server.UpdateStateIfCurrent(newLease, newState, newTime) + waitGroup.Wait() + + // Then + require.False(t, oldUpdateAccepted) + require.True(t, newUpdateAccepted) + require.Equal(t, newState, server.State) + require.Equal(t, newTime, server.LastActive) +} + +func TestReportSystemState_OldCleanupCannotClearNewerGeneration(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + oldLease := server.AttachStateStream(stateGenerationStream{}) + newLease := server.AttachStateStream(stateGenerationStream{}) + state := &model.HostState{Uptime: 22} + activeAt := time.Unix(200, 0) + require.True(t, server.UpdateStateIfCurrent(newLease, state, activeAt)) + + // When + oldCleanup := server.ClearStateStreamIfCurrent(oldLease) + + // Then + require.False(t, oldCleanup) + require.Equal(t, state, server.State) + require.Equal(t, activeAt, server.LastActive) +} + +func TestReportSystemState_CurrentCleanupClearsOnlineVisibility(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + lease := server.AttachStateStream(stateGenerationStream{}) + activeAt := time.Unix(300, 0) + require.True(t, server.UpdateStateIfCurrent(lease, &model.HostState{Uptime: 33}, activeAt)) + + // When + cleared := server.ClearStateStreamIfCurrent(lease) + + // Then + require.True(t, cleared) + require.True(t, server.LastActive.IsZero()) +} diff --git a/service/rpc/wait_for_agent_revoke_test.go b/service/rpc/wait_for_agent_revoke_test.go index e247b7d6..8d82d20a 100644 --- a/service/rpc/wait_for_agent_revoke_test.go +++ b/service/rpc/wait_for_agent_revoke_test.go @@ -16,7 +16,9 @@ import ( func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) { h := NewNezhaHandler() const streamID = "kill-switch-wait" - h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer) + if err := h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer); err != nil { + t.Fatalf("create waiter stream: %v", err) + } done := make(chan struct { io any @@ -35,9 +37,16 @@ func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) { dur time.Duration }{stream, ok, time.Since(start)} }() + waiter, err := h.GetStream(streamID) + if err != nil { + t.Fatalf("get waiter context: %v", err) + } + select { + case <-waiter.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not enter its blocking select") + } - // 让 WaitForAgent 真的进入 select 等待，再触发 kill switch。 - time.Sleep(50 * time.Millisecond) if revoked := h.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 { t.Fatalf("expected to revoke exactly 1 MCP stream, got %d", revoked) } @@ -54,3 +63,46 @@ func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) { t.Fatalf("WaitForAgent never returned after revoke; kill switch did not wake the waiter") } } + +func TestRevokeStreamsForServerWakesWaitForAgentAndPreservesNewGeneration(t *testing.T) { + h := NewNezhaHandler() + const streamID = "server-revoke-generation" + if err := h.CreateStream(streamID, 0, 7); err != nil { + t.Fatalf("create waiter stream: %v", err) + } + done := make(chan bool, 1) + go func() { + _, ok := h.WaitForAgent(context.Background(), streamID, time.Minute) + done <- ok + }() + waiter, err := h.GetStream(streamID) + if err != nil { + t.Fatalf("get waiter stream: %v", err) + } + select { + case <-waiter.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not reach its blocking select") + } + + h.RevokeStreamsForServer(7) + select { + case ok := <-done: + if ok { + t.Fatal("WaitForAgent must return false after server revocation") + } + case <-time.After(time.Second): + t.Fatal("server revocation did not wake WaitForAgent") + } + h.RevokeStreamsForServer(7) + if err := h.CreateStream(streamID, 0, 8); err != nil { + t.Fatalf("new generation must reuse released ID: %v", err) + } + h.RevokeStreamsForServer(7) + if h.StreamCount() != 1 { + t.Fatalf("new generation must remain tracked, got %d streams", h.StreamCount()) + } + if err := h.CloseStream(streamID); err != nil { + t.Fatalf("cleanup new generation: %v", err) + } +} diff --git a/service/singleton/crontask.go b/service/singleton/crontask.go index 0d846ca2..d8b7d6cb 100644 --- a/service/singleton/crontask.go +++ b/service/singleton/crontask.go @@ -24,6 +24,18 @@ type CronClass struct { *cron.Cron pendingAlertTriggerTasksMu sync.Mutex pendingAlertTriggerTasks map[uint64]map[uint64][]time.Time + closeOnce sync.Once +} + +// Close stops the scheduler and joins every job before a test restores globals. +// The embedded cron.Stop only exposes the completion context; callers must await it. +func (c *CronClass) Close() { + if c == nil || c.Cron == nil { + return + } + c.closeOnce.Do(func() { + <-c.Cron.Stop().Done() + }) } func NewCronClass() *CronClass { diff --git a/service/singleton/crontask_lifecycle_test.go b/service/singleton/crontask_lifecycle_test.go new file mode 100644 index 00000000..1566bbcc --- /dev/null +++ b/service/singleton/crontask_lifecycle_test.go @@ -0,0 +1,97 @@ +package singleton + +import ( + "context" + "testing" + "time" + + "github.com/robfig/cron/v3" + "github.com/stretchr/testify/require" +) + +const lifecycleTestTimeout = time.Second + +func TestCronClassClose_waitsForRunningJobs(t *testing.T) { + // Given + started := make(chan struct{}) + release := make(chan struct{}) + events := make(chan string, 9) + cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())} + _, err := cronClass.AddFunc("@every 1ns", func() { + defer func() { events <- "job" }() + close(started) + <-release + }) + require.NoError(t, err) + cronClass.Start() + <-started + + // When + closed := make(chan struct{}) + for range 8 { + go func() { + cronClass.Close() + events <- "close" + closed <- struct{}{} + }() + } + + // Then + close(release) + firstEvent := awaitCronLifecycleEvent(t, events, "cron lifecycle did not complete") + if firstEvent != "job" { + t.Fatalf("Close returned before the running cron job returned: first event=%q", firstEvent) + } + for range 8 { + awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return") + } + for range 8 { + if event := awaitCronLifecycleEvent(t, events, "concurrent Close call did not complete"); event != "close" { + t.Fatalf("unexpected cron lifecycle event: %q", event) + } + } +} + +func TestCronClassClose_isIdempotentAndNilSafe(t *testing.T) { + cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())} + cronClass.Start() + + closed := make(chan struct{}) + for range 8 { + go func() { + cronClass.Close() + closed <- struct{}{} + }() + } + for range 8 { + awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return") + } + cronClass.Close() + var nilCronClass *CronClass + nilCronClass.Close() + (&CronClass{}).Close() +} + +func awaitCronLifecycleSignal(t *testing.T, signal <-chan struct{}, message string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout) + defer cancel() + select { + case <-signal: + case <-ctx.Done(): + t.Fatal(message) + } +} + +func awaitCronLifecycleEvent(t *testing.T, events <-chan string, message string) string { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout) + defer cancel() + select { + case event := <-events: + return event + case <-ctx.Done(): + t.Fatal(message) + return "" + } +}