Buriburizaemon/nezha_domains
1
0
Fork 0
mirror of https://github.com/Buriburizaem0n/nezha_domains.git synced 2026-09-19 09:40:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rpc): make IO stream lifecycle race-safe

Browse Source 
Co-authored-by: naiba/CloudCode <hi+cloudcode@nai.ba>
This commit is contained in:
naiba
2026-07-20 04:26:19 +00:00
co-authored by naiba/CloudCode
parent 8b47ff141f
commit c756ef9385
23 changed files with 2114 additions and 690 deletions
Download Patch File Download Diff File Expand all files Collapse all files
service/rpc/geoip_rpc.go
+57
View File
@@ -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
}
service/rpc/io_stream.go
+58 -336
View File
@@ -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
}
service/rpc/io_stream_lifecycle.go
+128
View File
@@ -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
}
service/rpc/io_stream_lifecycle_test.go
+259
View File
@@ -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)
}
}
service/rpc/io_stream_quota.go
+62
View File
@@ -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)
}
service/rpc/io_stream_quota_test.go
+106
View File
@@ -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)
}
}
service/rpc/io_stream_registry.go
+85
View File
@@ -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
}
service/rpc/io_stream_rpc.go
+58
View File
@@ -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
}
service/rpc/io_stream_state.go
+95
View File
@@ -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{})
}
service/rpc/io_stream_state_expectation_test.go
+142
View File
@@ -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)
}
}
service/rpc/io_stream_state_test.go
+68
View File
@@ -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)
}
}
service/rpc/io_stream_state_wait_test.go
+220
View File
@@ -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")
}
}
service/rpc/io_stream_test.go
-126
View File
@@ -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 {
service/rpc/nezha.go
+90 -214
View File
@@ -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
}
service/rpc/request_task_security_test.go
+11 -4
View File
@@ -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()
service/rpc/state_metrics_test.go
+42
View File
@@ -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)
}
service/rpc/state_stream_generation_test.go
+206
View File
@@ -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())
}
service/rpc/wait_for_agent_revoke_test.go
+55 -3
View File
@@ -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)
}
}
service/singleton/crontask.go
+12
View File
@@ -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 {
service/singleton/crontask_lifecycle_test.go
+97
View File
@@ -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 ""
}
}