mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-22 19:20:12 +00:00
fix(rpc): cap concurrent IO streams per user and per server
GHSA-jg62-j5h6-8mpq: the terminal and file-manager endpoints created unbounded IO streams; an authenticated member could open thousands, each spawning goroutines, a 1MiB buffer and an agent-side PTY, exhausting dashboard and agent resources. CreateStream now enforces a per-user (20) and per-server (40) cap in the existing ioStreamMutex critical section, using the stream map as the single source of truth. Dashboard-internal streams (uid==0: NAT, server transfer, MCP transfer) skip the per-user cap but still count per-server. Adds caps, exemption, slot-release and no-leak regression tests.
This commit is contained in:
@@ -49,7 +49,9 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c), server.ID)
|
if err := rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c), server.ID); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
fmData, _ := json.Marshal(&model.TaskFM{
|
fmData, _ := json.Marshal(&model.TaskFM{
|
||||||
StreamID: streamId,
|
StreamID: streamId,
|
||||||
|
|||||||
@@ -962,7 +962,9 @@ func openFsTransferStream(ctx context.Context, serverID uint64, req *model.FsTra
|
|||||||
}
|
}
|
||||||
req.StreamID = streamId
|
req.StreamID = streamId
|
||||||
|
|
||||||
rpc.NezhaHandlerSingleton.CreateStreamWithPurpose(streamId, 0, serverID, rpc.PurposeMCPTransfer)
|
if err := rpc.NezhaHandlerSingleton.CreateStreamWithPurpose(streamId, 0, serverID, rpc.PurposeMCPTransfer); err != nil {
|
||||||
|
return nil, func() {}, err
|
||||||
|
}
|
||||||
cleanup := func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }
|
cleanup := func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }
|
||||||
|
|
||||||
body, err := json.Marshal(req)
|
body, err := json.Marshal(req)
|
||||||
|
|||||||
@@ -47,7 +47,9 @@ func createTerminal(c *gin.Context) (*model.CreateTerminalResponse, error) {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c), server.ID)
|
if err := rpc.NezhaHandlerSingleton.CreateStream(streamId, getUid(c), server.ID); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
terminalData, _ := json.Marshal(&model.TerminalTask{
|
terminalData, _ := json.Marshal(&model.TerminalTask{
|
||||||
StreamID: streamId,
|
StreamID: streamId,
|
||||||
|
|||||||
@@ -217,7 +217,11 @@ func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) {
|
|||||||
// IS required though — the receiving agent must prove it is the server the
|
// IS required though — the receiving agent must prove it is the server the
|
||||||
// NAT config addressed, otherwise any agent that snoops the streamId can
|
// NAT config addressed, otherwise any agent that snoops the streamId can
|
||||||
// answer NAT traffic on behalf of an unrelated host.
|
// answer NAT traffic on behalf of an unrelated host.
|
||||||
rpcService.NezhaHandlerSingleton.CreateStream(streamId, 0, server.ID)
|
if err := rpcService.NezhaHandlerSingleton.CreateStream(streamId, 0, server.ID); err != nil {
|
||||||
|
w.WriteHeader(http.StatusTooManyRequests)
|
||||||
|
w.Write(fmt.Appendf(nil, "stream limit: %v", err))
|
||||||
|
return
|
||||||
|
}
|
||||||
defer rpcService.NezhaHandlerSingleton.CloseStream(streamId)
|
defer rpcService.NezhaHandlerSingleton.CloseStream(streamId)
|
||||||
|
|
||||||
taskData, err := json.Marshal(model.TaskNAT{
|
taskData, err := json.Marshal(model.TaskNAT{
|
||||||
|
|||||||
@@ -48,14 +48,44 @@ var bufPool = sync.Pool{
|
|||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) {
|
const (
|
||||||
s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy)
|
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) {
|
func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error {
|
||||||
s.ioStreamMutex.Lock()
|
s.ioStreamMutex.Lock()
|
||||||
defer s.ioStreamMutex.Unlock()
|
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{
|
s.ioStreams[streamId] = &ioStreamContext{
|
||||||
creatorUserID: creatorUserID,
|
creatorUserID: creatorUserID,
|
||||||
targetServerID: targetServerID,
|
targetServerID: targetServerID,
|
||||||
@@ -64,6 +94,7 @@ func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID ui
|
|||||||
agentIoConnectCh: make(chan struct{}),
|
agentIoConnectCh: make(chan struct{}),
|
||||||
revokedCh: make(chan struct{}),
|
revokedCh: make(chan struct{}),
|
||||||
}
|
}
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// IsStreamAuthorizedForAgent reports whether the connecting agent is the
|
// IsStreamAuthorizedForAgent reports whether the connecting agent is the
|
||||||
@@ -270,8 +301,6 @@ func (s *NezhaHandler) CloseStream(streamId string) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
// UserConnected publishes the user-side IO under ioStreamMutex so concurrent
|
// UserConnected publishes the user-side IO under ioStreamMutex so concurrent
|
||||||
// Revoke* / WaitForAgent / StartStream see a consistent stream view.
|
// Revoke* / WaitForAgent / StartStream see a consistent stream view.
|
||||||
// Without the lock, the bare assignment to stream.userIo races with the
|
// Without the lock, the bare assignment to stream.userIo races with the
|
||||||
|
|||||||
@@ -0,0 +1,67 @@
|
|||||||
|
package rpc
|
||||||
|
|
||||||
|
import (
|
||||||
|
"runtime"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// settleGoroutines lets transient goroutines wind down so the count reflects
|
||||||
|
// only durable leaks, not in-flight teardown.
|
||||||
|
func settleGoroutines() int {
|
||||||
|
var n int
|
||||||
|
for i := 0; i < 50; i++ {
|
||||||
|
runtime.GC()
|
||||||
|
time.Sleep(20 * time.Millisecond)
|
||||||
|
n = runtime.NumGoroutine()
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestStartStream_NoGoroutineLeakAfterClose verifies the bidirectional relay in
|
||||||
|
// StartStream does not strand a goroutine. StartStream launches two
|
||||||
|
// io.CopyBuffer goroutines (user<-agent and agent<-user) but returns after the
|
||||||
|
// first one finishes. The second goroutine stays blocked in CopyBuffer until
|
||||||
|
// its endpoints are closed. CloseStream closes both endpoints, which must
|
||||||
|
// unblock and drain that second goroutine. If it doesn't, every terminal / fm /
|
||||||
|
// NAT session leaks one goroutine for the lifetime of the dashboard.
|
||||||
|
func TestStartStream_NoGoroutineLeakAfterClose(t *testing.T) {
|
||||||
|
base := settleGoroutines()
|
||||||
|
|
||||||
|
const n = 20
|
||||||
|
for i := 0; i < n; i++ {
|
||||||
|
h := NewNezhaHandler()
|
||||||
|
const id = "leak-stream"
|
||||||
|
|
||||||
|
if err := h.CreateStream(id, 1, 1); err != nil {
|
||||||
|
t.Fatalf("CreateStream: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
userIo, agentIo := newPipeReadWriter(), newPipeReadWriter()
|
||||||
|
h.AgentConnected(id, agentIo)
|
||||||
|
h.UserConnected(id, userIo)
|
||||||
|
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
_ = h.StartStream(id, time.Second*5)
|
||||||
|
close(done)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Close one endpoint so the first CopyBuffer returns and StartStream
|
||||||
|
// unblocks, mirroring a peer disconnect.
|
||||||
|
time.Sleep(10 * time.Millisecond)
|
||||||
|
userIo.Close()
|
||||||
|
<-done
|
||||||
|
|
||||||
|
// The caller's defer CloseStream closes both endpoints, which must
|
||||||
|
// drain the still-blocked second copy goroutine.
|
||||||
|
_ = h.CloseStream(id)
|
||||||
|
agentIo.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
after := settleGoroutines()
|
||||||
|
if grew := after - base; grew > 2 {
|
||||||
|
t.Fatalf("goroutine leak in StartStream relay: ran %d streams, goroutines grew by %d (base=%d after=%d)",
|
||||||
|
n, grew, base, after)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,6 +1,8 @@
|
|||||||
package rpc
|
package rpc
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"reflect"
|
"reflect"
|
||||||
"testing"
|
"testing"
|
||||||
@@ -98,6 +100,130 @@ 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 {
|
func newPipeReadWriter() io.ReadWriteCloser {
|
||||||
r, w := io.Pipe()
|
r, w := io.Pipe()
|
||||||
return struct {
|
return struct {
|
||||||
|
|||||||
Reference in New Issue
Block a user