diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index b390882ad..f282b4d6e 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -19,11 +19,16 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" cache: false + + - name: Show Go version + run: go version - - name: Go mod tidy - run: go mod tidy + - name: Verify Go module files + run: | + go mod tidy + git diff --exit-code -- go.mod go.sum - name: Install embedded build assets run: go run ./scripts/pre_install.go @@ -50,21 +55,34 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" cache: false + + - name: Show Go version + run: go version - - name: Go mod tidy - run: go mod tidy + - name: Verify Go module files + run: | + go mod tidy + git diff --exit-code -- go.mod go.sum - name: Install embedded build assets run: go run ./scripts/pre_install.go - - name: Race detection — core, parser, stream + - name: Race detection — client, core, parser, RPC, listener, stream run: >- - go test -race -count=1 -timeout 300s - ./server/internal/core - ./server/internal/parser/... - ./server/internal/stream + go test -race -count=1 -timeout 300s + ./client/core/... + ./server/internal/core + ./server/internal/parser/... + ./server/internal/stream + ./server/listener + ./server/rpc + + - name: Parser fuzz smoke tests + run: | + go test ./server/internal/parser/malefic -run=^$ -fuzz=^FuzzMaleficParserReadHeader$ -fuzztime=5s + go test ./server/internal/parser/malefic -run=^$ -fuzz=^FuzzMaleficParserParse$ -fuzztime=5s mock_implant: runs-on: ubuntu-22.04 @@ -77,11 +95,16 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" cache: false + + - name: Show Go version + run: go version - - name: Go mod tidy - run: go mod tidy + - name: Verify Go module files + run: | + go mod tidy + git diff --exit-code -- go.mod go.sum - name: Install embedded build assets run: go run ./scripts/pre_install.go @@ -107,11 +130,16 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" cache: false + + - name: Show Go version + run: go version - - name: Go mod tidy - run: go mod tidy + - name: Verify Go module files + run: | + go mod tidy + git diff --exit-code -- go.mod go.sum - name: Install embedded build assets run: go run ./scripts/pre_install.go diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index ec1d5109b..40020d15f 100644 --- a/.github/workflows/nightly.yaml +++ b/.github/workflows/nightly.yaml @@ -23,7 +23,10 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" + + - name: Show Go version + run: go version - name: Configure Git run: | diff --git a/.github/workflows/realimplant.yaml b/.github/workflows/realimplant.yaml index 001b0e35c..dc9ecf649 100644 --- a/.github/workflows/realimplant.yaml +++ b/.github/workflows/realimplant.yaml @@ -22,9 +22,12 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" cache: false + - name: Show Go version + run: go version + - name: Validate real implant environment shell: pwsh run: | @@ -38,6 +41,9 @@ jobs: - name: Go mod tidy run: go mod tidy + - name: Verify Go module files + run: git diff --exit-code -- go.mod go.sum + - name: Real implant E2E tests shell: pwsh run: | diff --git a/.github/workflows/releaser.yaml b/.github/workflows/releaser.yaml index 82c6098eb..d9fdc82cc 100644 --- a/.github/workflows/releaser.yaml +++ b/.github/workflows/releaser.yaml @@ -24,8 +24,10 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" - - run: go version + go-version: "1.25.0" + + - name: Show Go version + run: go version - name: Configure Git run: | @@ -75,4 +77,3 @@ jobs: dist/vsix/iom.vsix env: GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }} - diff --git a/.github/workflows/update-search-index.yaml b/.github/workflows/update-search-index.yaml index 61b62ff2f..d1165731e 100644 --- a/.github/workflows/update-search-index.yaml +++ b/.github/workflows/update-search-index.yaml @@ -53,9 +53,12 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: "1.24.13" + go-version: "1.25.0" cache: false + - name: Show Go version + run: go version + - name: Update community submodule to latest run: | git config --global url."https://${{ secrets.PAT_TOKEN }}@github.com/".insteadOf "https://github.com/" diff --git a/client/command/build/artifact_command_test.go b/client/command/build/artifact_command_test.go index e3cedba77..4535cde68 100644 --- a/client/command/build/artifact_command_test.go +++ b/client/command/build/artifact_command_test.go @@ -48,12 +48,16 @@ func TestArtifactPublishAddsWebsiteContent(t *testing.T) { } calls := h.Recorder.Calls() - if len(calls) != 1 || calls[0].Method != "AddWebsiteContent" { + if len(calls) != 2 || calls[0].Method != "DownloadArtifact" || calls[1].Method != "AddWebsiteContent" { t.Fatalf("calls = %#v", calls) } - add := calls[0].Request.(*clientpb.Website) + download := calls[0].Request.(*clientpb.Artifact) + if download.Name != "demo-artifact" { + t.Fatalf("download artifact request = %#v", download) + } + add := calls[1].Request.(*clientpb.Website) content := add.Contents["/payload.bin"] - if add.Name != "site-a" || content == nil || content.GetType() != consts.ArtifactWebcontent || content.GetFile() != "demo-artifact" { + if add.Name != "site-a" || content == nil || string(content.Content) != "artifact-bin" { t.Fatalf("add website content request = %#v", add) } } diff --git a/client/command/build/commands.go b/client/command/build/commands.go index 060d09437..93f4fdd0d 100644 --- a/client/command/build/commands.go +++ b/client/command/build/commands.go @@ -630,13 +630,13 @@ artifact prune --failed --older-than 720h func Register(con *core.Console) { // EventCallback requires initialized Server; skip in doc-gen mode (genlua) if con.Server != nil { - con.EventCallback[consts.CtrlArtifactDownload] = func(event *clientpb.Event) { + con.SetEventCallback(consts.CtrlArtifactDownload, func(event *clientpb.Event) { err := WriteOriginArtifact(con, event.Job.Name) if err != nil { con.Log.Errorf("write artifact %s error: %s", event.Job.Name, err) return } - } + }) } con.RegisterServerFunc("search_artifact", SearchArtifact, diff --git a/client/command/implant.go b/client/command/implant.go index 5feb9a8bb..96cec7a39 100644 --- a/client/command/implant.go +++ b/client/command/implant.go @@ -105,10 +105,11 @@ func makeRunners(implantCmd *cobra.Command, con *core.Console) (pre, post func(c return nil } sid, _ := cmd.Flags().GetString("use") - if sid == "" && con.ActiveTarget.Session == nil { + activeSession := con.ActiveTarget.Get() + if sid == "" && activeSession == nil { return fmt.Errorf("no implant to run command on") - } else if sid == "" && con.ActiveTarget.Session != nil { - sid = con.ActiveTarget.Session.SessionId + } else if sid == "" { + sid = activeSession.SessionId } var session *client.Session diff --git a/client/command/listener/pipeline_integration_test.go b/client/command/listener/pipeline_integration_test.go index 4d405d14d..0ee818438 100644 --- a/client/command/listener/pipeline_integration_test.go +++ b/client/command/listener/pipeline_integration_test.go @@ -87,9 +87,9 @@ func TestStopPipelineCmdUsesDatabaseResolution(t *testing.T) { } testsupport.WaitForCondition(t, 5*time.Second, func() bool { - _, ok := clientHarness.Console.Pipelines["tcp-stop"] - return !ok - }, "client pipeline cache to remove stopped pipeline") + pipeline, err := clientHarness.Console.FindCachedPipeline("tcp-stop", nil) + return err == nil && !pipeline.GetEnable() + }, "client pipeline cache to retain stopped pipeline as disabled") model, err := h.GetPipeline("tcp-stop", h.ListenerID()) if err != nil { diff --git a/client/command/website/webcontent.go b/client/command/website/webcontent.go index 5fae29431..bba2ef67b 100644 --- a/client/command/website/webcontent.go +++ b/client/command/website/webcontent.go @@ -2,14 +2,15 @@ package website import ( "errors" - "github.com/chainreactors/IoM-go/consts" + "github.com/chainreactors/IoM-go/proto/client/clientpb" "github.com/chainreactors/malice-network/client/core" + "github.com/chainreactors/malice-network/helper/utils/fileutils" + "github.com/chainreactors/malice-network/helper/utils/output" "github.com/chainreactors/malice-network/helper/utils/pe" "github.com/chainreactors/tui" "github.com/evertras/bubble-table/table" "github.com/spf13/cobra" - "net/url" "os" "path/filepath" "strconv" @@ -235,30 +236,24 @@ func AddArtifactContentCmd(cmd *cobra.Command, con *core.Console) error { func AddArtifactContent(con *core.Console, artifactName, websiteName, format, rdi, webPath, contentType, name, comment, auth string) (*clientpb.WebContent, error) { rpcFormat := normalizeArtifactContentFormat(format) + artifact, err := con.Rpc.DownloadArtifact(con.Context(), &clientpb.Artifact{ + Name: artifactName, + Format: rpcFormat, + Rdi: rdi, + }) + if err != nil { + return nil, err + } + if len(artifact.Bin) == 0 { + return nil, errors.New("artifact maybe not download in server") + } if name == "" { name = artifactName } - - resolvedWebsite, listenerID, _ := resolveWebsiteTarget(con, websiteName) - website := &clientpb.Website{ - Name: resolvedWebsite, - ListenerId: listenerID, - Contents: map[string]*clientpb.WebContent{ - webPath: { - WebsiteId: resolvedWebsite, - ListenerId: listenerID, - File: artifactName, - Path: webPath, - Type: consts.ArtifactWebcontent, - Url: artifactContentOptions(rpcFormat, rdi), - ContentType: contentType, - Name: name, - Comment: comment, - Auth: auth, - }, - }, + if webPath == "" { + webPath = defaultArtifactWebPath(artifact, artifactName, rpcFormat) } - return con.Rpc.AddWebsiteContent(con.Context(), website) + return AddWebContentData(con, websiteName, artifact.Bin, webPath, contentType, name, comment, auth) } func normalizeArtifactContentFormat(format string) string { @@ -268,15 +263,25 @@ func normalizeArtifactContentFormat(format string) string { return format } -func artifactContentOptions(format, rdi string) string { - values := url.Values{} - if format != "" { - values.Set("format", format) +func defaultArtifactWebPath(artifact *clientpb.Artifact, fallbackName, format string) string { + name := artifact.GetName() + if name == "" { + name = fallbackName + } + ext := artifact.GetFormat() + if f, ok := output.SupportedFormats[strings.ToLower(format)]; ok && f.Extension != "" { + ext = f.Extension + } + if ext == "" { + ext, _ = fileutils.GetExtensionByBytes(artifact.GetBin()) + } + if ext != "" && !strings.HasPrefix(ext, ".") { + ext = "." + ext } - if rdi != "" { - values.Set("rdi", rdi) + if ext != "" && strings.HasSuffix(strings.ToLower(name), strings.ToLower(ext)) { + ext = "" } - return values.Encode() + return "/" + name + ext } // RemoveWebContentCmd - 删除网站内容 diff --git a/client/command/website/webcontent_cmd_test.go b/client/command/website/webcontent_cmd_test.go index a5b97a7bf..87768bb36 100644 --- a/client/command/website/webcontent_cmd_test.go +++ b/client/command/website/webcontent_cmd_test.go @@ -208,37 +208,60 @@ func TestAddWebContentCmdSupportsArtifactFlag(t *testing.T) { if err := AddWebContentCmd(cmd, con); err != nil { t.Fatalf("AddWebContentCmd failed: %v", err) } - if len(rpc.calls) != 1 { - t.Fatalf("call count = %d, want 1", len(rpc.calls)) + if len(rpc.calls) != 2 { + t.Fatalf("call count = %d, want 2", len(rpc.calls)) } - add := rpc.calls[0].request.(*clientpb.Website) - content := add.GetContents()["/payload.bin"] - if rpc.calls[0].method != "AddWebsiteContent" || add.Name != "site-a" || add.ListenerId != "listener-a" || content == nil { - t.Fatalf("add call = %s %#v, content %#v, want scoped website payload", rpc.calls[0].method, add, content) + download := rpc.calls[0].request.(*clientpb.Artifact) + if download.Name != "beacon" || download.Format != "raw" { + t.Fatalf("download request = %#v, want beacon/raw", download) } - if len(content.Content) != 0 || content.Type != consts.ArtifactWebcontent || content.File != "beacon" || content.Url != "format=raw" || content.ContentType != "application/octet-stream" || content.Name != "payload" { - t.Fatalf("content = %#v, want server-side artifact reference", content) + add := rpc.calls[1].request.(*clientpb.Website) + content := add.GetContents()["/payload.bin"] + if content == nil || content.ContentType != "application/octet-stream" || content.Name != "payload" { + t.Fatalf("content = %#v, want artifact website content", content) } } -func TestAddArtifactContentAddsServerSideArtifactReference(t *testing.T) { +func TestAddArtifactContentDownloadsAndAddsWebsiteContent(t *testing.T) { rpc := &websiteTestRPC{} con := newWebsiteTestConsole(rpc) con.Pipelines["listener-a:site-a"] = scopedWebsitePipeline("site-a", "listener-a") - if _, err := AddArtifactContent(con, "beacon", "listener-a:site-a", "shellcode", "", "/payload.bin", "application/octet-stream", "payload", "from artifact", "none"); err != nil { + if _, err := AddArtifactContent(con, "beacon", "listener-a:site-a", "shellcode", "rdi-test", "/payload.bin", "application/octet-stream", "payload", "from artifact", "none"); err != nil { t.Fatalf("AddArtifactContent failed: %v", err) } - if len(rpc.calls) != 1 { - t.Fatalf("call count = %d, want 1", len(rpc.calls)) + if len(rpc.calls) != 2 { + t.Fatalf("call count = %d, want 2", len(rpc.calls)) } - add := rpc.calls[0].request.(*clientpb.Website) + download := rpc.calls[0].request.(*clientpb.Artifact) + if rpc.calls[0].method != "DownloadArtifact" || download.Name != "beacon" || download.Format != "raw" || download.Rdi != "rdi-test" { + t.Fatalf("download call = %s %#v, want beacon/raw with RDI", rpc.calls[0].method, download) + } + add := rpc.calls[1].request.(*clientpb.Website) content := add.GetContents()["/payload.bin"] - if rpc.calls[0].method != "AddWebsiteContent" || add.Name != "site-a" || add.ListenerId != "listener-a" || content == nil { - t.Fatalf("add call = %s %#v, content %#v, want scoped website payload", rpc.calls[0].method, add, content) + if rpc.calls[1].method != "AddWebsiteContent" || add.Name != "site-a" || add.ListenerId != "listener-a" || content == nil { + t.Fatalf("add call = %s %#v, content %#v, want scoped website payload", rpc.calls[1].method, add, content) + } + if string(content.Content) != "artifact-binary" || content.Name != "payload" || content.Comment != "from artifact" || content.Auth != "none" { + t.Fatalf("content = %#v, want artifact bytes with metadata", content) + } +} + +func TestDefaultArtifactWebPathDoesNotDuplicateExtension(t *testing.T) { + tests := []struct { + name string + artifact *clientpb.Artifact + want string + }{ + {name: "existing extension", artifact: &clientpb.Artifact{Name: "payload.exe", Format: ".exe"}, want: "/payload.exe"}, + {name: "append extension", artifact: &clientpb.Artifact{Name: "payload", Format: ".exe"}, want: "/payload.exe"}, } - if len(content.Content) != 0 || content.Type != consts.ArtifactWebcontent || content.File != "beacon" || content.Name != "payload" || content.Comment != "from artifact" || content.Auth != "none" { - t.Fatalf("content = %#v, want server-side artifact reference with metadata", content) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := defaultArtifactWebPath(tt.artifact, "fallback", ""); got != tt.want { + t.Fatalf("defaultArtifactWebPath() = %q, want %q", got, tt.want) + } + }) } } diff --git a/client/core/server.go b/client/core/server.go index 4c8114ef5..43acb1a8d 100644 --- a/client/core/server.go +++ b/client/core/server.go @@ -6,7 +6,9 @@ import ( "fmt" "io" "reflect" + "strconv" "sync" + "time" "github.com/chainreactors/IoM-go/client" "github.com/chainreactors/IoM-go/consts" @@ -17,10 +19,13 @@ import ( "github.com/chainreactors/malice-network/helper/intermediate" "github.com/chainreactors/malice-network/helper/utils/output" "google.golang.org/grpc" + "google.golang.org/grpc/metadata" ) var ErrLuaVMDead = fmt.Errorf("lua vm is dead") +const eventStreamHeaderTimeout = 10 * time.Second + func wrapToTaskContext(event *clientpb.Event) *clientpb.TaskContext { return &clientpb.TaskContext{ Task: event.Task, @@ -31,13 +36,15 @@ func wrapToTaskContext(event *clientpb.Event) *clientpb.TaskContext { type Server struct { *client.ServerState - taskMessageMu sync.Mutex - taskMessages map[string]string - eventHookMu sync.RWMutex - eventHookIDs map[client.EventCondition][]uint64 - nextEventHookID uint64 - eventHandlerMu sync.Mutex - eventHandlerActive bool + taskMessageMu sync.Mutex + taskMessages map[string]string + eventHookMu sync.RWMutex + eventHookIDs map[client.EventCondition][]uint64 + nextEventHookID uint64 + eventHandlerMu sync.Mutex + eventHandlerActive bool + suppressStartupOutput bool + startupHistoryPending bool // Quiet suppresses console event output while still updating internal state. Quiet bool @@ -76,6 +83,18 @@ func (s *Server) EventHandlerRunning() bool { return s.eventHandlerActive } +func (s *Server) needsStartupHistory() bool { + s.eventHandlerMu.Lock() + defer s.eventHandlerMu.Unlock() + return s.startupHistoryPending +} + +func (s *Server) markStartupHistoryComplete() { + s.eventHandlerMu.Lock() + s.startupHistoryPending = false + s.eventHandlerMu.Unlock() +} + func taskMessageKey(sessionID string, taskID uint32) string { return fmt.Sprintf("%s-%d", sessionID, taskID) } @@ -123,12 +142,12 @@ func NewServerWithOptions(conn *grpc.ClientConn, config *mtls.ClientConfig, supp if err != nil { return nil, err } - ser := &Server{ServerState: s, taskMessages: make(map[string]string)} - events, err := ser.GetEvent(context.Background(), &clientpb.Int{}) - if err != nil { - return nil, err + ser := &Server{ + ServerState: s, + taskMessages: make(map[string]string), + suppressStartupOutput: suppressStartupOutput, + startupHistoryPending: true, } - ser.replayCachedEvents(events.GetEvents(), suppressStartupOutput) return ser, nil } @@ -437,18 +456,82 @@ func (s *Server) EventHandler() { } defer s.releaseEventHandler() - eventStream, err := s.Rpc.Events(context.Background(), &clientpb.Empty{}) + replayStartupHistory := s.needsStartupHistory() + eventBaseContext, cancelEventStream := context.WithCancel(context.Background()) + defer cancelEventStream() + eventContext := eventBaseContext + if replayStartupHistory { + eventContext = metadata.NewOutgoingContext(eventContext, metadata.Pairs( + consts.EventStreamHistoryReplayHeader, "true", + )) + } + eventStream, err := s.Rpc.Events(eventContext, &clientpb.Empty{}) if err != nil { return } - // Wait until the server has installed the event subscription before taking - // a new snapshot. Events produced during Update then remain queued in order. - if _, err := eventStream.Header(); err != nil { + type headerResult struct { + header metadata.MD + err error + } + headerCh := make(chan headerResult, 1) + go func() { + header, headerErr := eventStream.Header() + headerCh <- headerResult{header: header, err: headerErr} + }() + headerTimer := time.NewTimer(eventStreamHeaderTimeout) + defer headerTimer.Stop() + var header metadata.MD + select { + case result := <-headerCh: + if result.err != nil { + return + } + header = result.header + case <-headerTimer.C: + client.Log.Warnf("event stream header timed out after %s\n", eventStreamHeaderTimeout) return } - s.Update() - if s.Session != nil { - s.UpdateSession(s.GetInteractive().SessionId) + if replayStartupHistory { + historyCount := 0 + if values := header.Get(consts.EventStreamHistoryCountHeader); len(values) > 0 { + parsed, parseErr := strconv.Atoi(values[0]) + if parseErr != nil || parsed < 0 { + client.Log.Warnf("invalid event history count %q\n", values[0]) + return + } + historyCount = parsed + } else { + // Older servers cannot provide an atomic history/live cut. Fetching + // history after subscription avoids loss but may replay a concurrent + // important event once; this is an intentional compatibility fallback. + historyContext, cancelHistory := context.WithTimeout(eventBaseContext, eventStreamHeaderTimeout) + history, historyErr := s.GetEvent(historyContext, &clientpb.Int{}) + cancelHistory() + if historyErr != nil { + client.Log.Warnf("failed to load event history: %v\n", historyErr) + return + } + s.replayCachedEvents(history.GetEvents(), s.suppressStartupOutput) + } + for i := 0; i < historyCount; i++ { + event, recvErr := eventStream.Recv() + if recvErr != nil { + client.Log.Warnf("event history receive failed: %v\n", recvErr) + return + } + if !s.suppressStartupOutput && event != nil && event.Type != consts.EventTask { + s.handleEventOutput(event) + } + } + s.markStartupHistoryComplete() + } + if err := s.Update(); err != nil { + client.Log.Warnf("failed to refresh server state: %v\n", err) + } + if session := s.Get(); session != nil { + if _, err := s.UpdateSession(session.SessionId); err != nil { + client.Log.Warnf("failed to refresh active session %s: %v\n", session.SessionId, err) + } } client.Log.Info("starting event loop\n") defer func() { @@ -456,12 +539,18 @@ func (s *Server) EventHandler() { }() for { event, err := eventStream.Recv() - if err == io.EOF || event == nil { + if err != nil { + if err != io.EOF { + client.Log.Warnf("event stream receive failed: %v\n", err) + } + return + } + if event == nil { return } s.dispatchEventHooks(event) - if fn, ok := s.EventCallback[event.Op]; ok { + if fn, ok := s.GetEventCallback(event.Op); ok { fn(event) } s.HandlerEvent(event) @@ -509,7 +598,6 @@ func (s *Server) HandlerEvent(event *clientpb.Event) { s.ReconcileEvent(event) s.handleEventOutput(event) } - func (s *Server) handleEventOutput(event *clientpb.Event) { if s == nil || event == nil { return diff --git a/client/core/server_state_test.go b/client/core/server_state_test.go index e478f0f5b..081d9cb42 100644 --- a/client/core/server_state_test.go +++ b/client/core/server_state_test.go @@ -86,6 +86,17 @@ func TestEventHandlerRunningTracksClaimLifecycle(t *testing.T) { } } +func TestStartupHistoryIsRequestedOnlyUntilCompleted(t *testing.T) { + s := &Server{startupHistoryPending: true} + if !s.needsStartupHistory() { + t.Fatal("new server should request startup history") + } + s.markStartupHistoryComplete() + if s.needsStartupHistory() { + t.Fatal("completed startup history should not be requested again") + } +} + func TestRenderEventAppliesColoringForSessionRegister(t *testing.T) { event := &clientpb.Event{ Type: consts.EventSession, diff --git a/helper/utils/compress/compress.go b/helper/utils/compress/compress.go index 149641ce9..f0d60f147 100644 --- a/helper/utils/compress/compress.go +++ b/helper/utils/compress/compress.go @@ -2,6 +2,11 @@ package compress import "github.com/golang/snappy" +// DecodedLen returns the uncompressed length declared by a Snappy block. +func DecodedLen(compressedData []byte) (int, error) { + return snappy.DecodedLen(compressedData) +} + // Compress compresses the input data using Snappy and returns the compressed data or an error. func Compress(data []byte) ([]byte, error) { // Use snappy compressor to compress the input data diff --git a/helper/utils/compress/compress_test.go b/helper/utils/compress/compress_test.go index 0159552e9..c2ca4dde6 100644 --- a/helper/utils/compress/compress_test.go +++ b/helper/utils/compress/compress_test.go @@ -1,8 +1,13 @@ package compress import ( + "bytes" + "encoding/binary" + "errors" "fmt" "testing" + + "github.com/golang/snappy" ) // TestSnappyCompression tests the compress and decompress functionality using Snappy. @@ -32,3 +37,52 @@ func TestSnappyCompression(t *testing.T) { t.Errorf("Decompressed data doesn't match original data.\nExpected: %s\nGot: %s", originalData, decompressedData) } } + +func TestValidateSnappyBlock(t *testing.T) { + for _, original := range [][]byte{ + nil, + []byte("small payload"), + bytes.Repeat([]byte("compressible"), 1024), + } { + compressed := snappy.Encode(nil, original) + decodedLength, err := snappy.DecodedLen(compressed) + if err != nil { + t.Fatalf("decoded length: %v", err) + } + if err := ValidateSnappyBlock(compressed, decodedLength); err != nil { + t.Fatalf("ValidateSnappyBlock rejected valid block: %v", err) + } + } +} + +func TestValidateSnappyBlockRejectsHeaderOnlyLargeBlock(t *testing.T) { + var header [binary.MaxVarintLen64]byte + n := binary.PutUvarint(header[:], 2<<30) + if err := ValidateSnappyBlock(header[:n], 2<<30); !errors.Is(err, snappy.ErrCorrupt) { + t.Fatalf("ValidateSnappyBlock error = %v, want snappy.ErrCorrupt", err) + } +} + +func FuzzValidateSnappyBlockMatchesDecoder(f *testing.F) { + for _, seed := range [][]byte{ + nil, + snappy.Encode(nil, []byte("valid")), + snappy.Encode(nil, bytes.Repeat([]byte("copy"), 1024)), + {0x05}, + {0x01, 0x01, 0x00}, + } { + f.Add(seed) + } + + f.Fuzz(func(t *testing.T, compressed []byte) { + decodedLength, err := snappy.DecodedLen(compressed) + if err != nil || decodedLength > 1<<20 { + return + } + validateErr := ValidateSnappyBlock(compressed, decodedLength) + _, decodeErr := snappy.Decode(nil, compressed) + if (validateErr == nil) != (decodeErr == nil) { + t.Fatalf("validator/decode mismatch: validate=%v decode=%v data=%x", validateErr, decodeErr, compressed) + } + }) +} diff --git a/helper/utils/compress/snappy_validate.go b/helper/utils/compress/snappy_validate.go new file mode 100644 index 000000000..5505cca68 --- /dev/null +++ b/helper/utils/compress/snappy_validate.go @@ -0,0 +1,83 @@ +package compress + +import ( + "encoding/binary" + + "github.com/golang/snappy" +) + +// ValidateSnappyBlock verifies the complete block structure without allocating +// the declared decoded buffer. decodedLength must come from snappy.DecodedLen. +func ValidateSnappyBlock(compressedData []byte, decodedLength int) error { + declaredLength, headerLength := binary.Uvarint(compressedData) + if headerLength <= 0 || decodedLength < 0 || declaredLength != uint64(decodedLength) { + return snappy.ErrCorrupt + } + + src := compressedData[headerLength:] + limit := uint64(decodedLength) + var produced uint64 + for cursor := 0; cursor < len(src); { + tag := src[cursor] + var length, offset uint64 + + switch tag & 0x03 { + case 0: // Literal. + literalLength := uint64(tag >> 2) + headerBytes := 1 + if literalLength >= 60 { + extraBytes := int(literalLength - 59) + if len(src)-cursor < 1+extraBytes { + return snappy.ErrCorrupt + } + literalLength = 0 + for i := 0; i < extraBytes; i++ { + literalLength |= uint64(src[cursor+1+i]) << (8 * i) + } + headerBytes += extraBytes + } + length = literalLength + 1 + cursor += headerBytes + if length > limit-produced || length > uint64(len(src)-cursor) { + return snappy.ErrCorrupt + } + cursor += int(length) + produced += length + continue + + case 1: // Copy with an 11-bit offset. + if len(src)-cursor < 2 { + return snappy.ErrCorrupt + } + length = 4 + uint64(tag>>2&0x07) + offset = uint64(tag&0xe0)<<3 | uint64(src[cursor+1]) + cursor += 2 + + case 2: // Copy with a 16-bit offset. + if len(src)-cursor < 3 { + return snappy.ErrCorrupt + } + length = 1 + uint64(tag>>2) + offset = uint64(binary.LittleEndian.Uint16(src[cursor+1 : cursor+3])) + cursor += 3 + + case 3: // Copy with a 32-bit offset. + if len(src)-cursor < 5 { + return snappy.ErrCorrupt + } + length = 1 + uint64(tag>>2) + offset = uint64(binary.LittleEndian.Uint32(src[cursor+1 : cursor+5])) + cursor += 5 + } + + if offset == 0 || offset > produced || length > limit-produced { + return snappy.ErrCorrupt + } + produced += length + } + + if produced != limit { + return snappy.ErrCorrupt + } + return nil +} diff --git a/server/client_server_integration_test.go b/server/client_server_integration_test.go index b019daf8f..9b4e51910 100644 --- a/server/client_server_integration_test.go +++ b/server/client_server_integration_test.go @@ -100,7 +100,7 @@ func TestClientServerControlPlaneIntegration(t *testing.T) { testsupport.WaitForEvent(t, stopEvents, "stop pipeline event") testsupport.WaitForCondition(t, 5*time.Second, func() bool { - _, ok := server.Pipelines[pipeline.Name] - return !ok - }, "event reconciliation to remove stopped pipeline") + cached, err := server.FindCachedPipeline(pipeline.Name, nil) + return err == nil && !cached.GetEnable() + }, "event reconciliation to retain stopped pipeline as disabled") } diff --git a/server/internal/core/connection.go b/server/internal/core/connection.go index b88fbb680..bb7c4d787 100644 --- a/server/internal/core/connection.go +++ b/server/internal/core/connection.go @@ -2,6 +2,7 @@ package core import ( "context" + "errors" "fmt" "strings" "sync" @@ -30,6 +31,11 @@ var ( ErrConnectionRemoved = fmt.Errorf("connection removed") ) +const ( + payloadReadBaseTimeout = 30 * time.Second + payloadReadMinBytesPerSecond = 512 << 10 +) + // listenerSessions 管理 listener 端的 session 信息 type listenerSessions struct { sessions *sync.Map // map[uint32]*clientpb.Session @@ -285,11 +291,22 @@ func (c *Connection) Send(ctx context.Context, conn *cryptostream.Conn) error { func (c *Connection) buildResponse(conn *cryptostream.Conn, length uint32) error { var msg *implantpb.Spites if length >= 2 { - var err error - msg, err = c.Parser.ReadMessage(conn, length) - if err != nil { - return fmt.Errorf("error reading message:%s %w", conn.RemoteAddr(), err) + if err := conn.SetReadDeadline(time.Now().Add(payloadReadTimeout(length))); err != nil { + return fmt.Errorf("set payload read deadline for %s: %w", conn.RemoteAddr(), err) + } + readMsg, readErr := c.Parser.ReadMessage(conn, length) + clearErr := conn.SetReadDeadline(time.Time{}) + if readErr != nil { + readErr = fmt.Errorf("error reading message:%s %w", conn.RemoteAddr(), readErr) + if clearErr != nil { + return errors.Join(readErr, fmt.Errorf("clear payload read deadline for %s: %w", conn.RemoteAddr(), clearErr)) + } + return readErr + } + if clearErr != nil { + return fmt.Errorf("clear payload read deadline for %s: %w", conn.RemoteAddr(), clearErr) } + msg = readMsg if msg.Spites == nil { msg = types.BuildPingSpites() } @@ -306,6 +323,11 @@ func (c *Connection) buildResponse(conn *cryptostream.Conn, length uint32) error return nil } +func payloadReadTimeout(length uint32) time.Duration { + seconds := (uint64(length) + payloadReadMinBytesPerSecond - 1) / payloadReadMinBytesPerSecond + return payloadReadBaseTimeout + time.Duration(seconds)*time.Second +} + func (c *Connection) Handler(ctx context.Context, conn *cryptostream.Conn) error { var err error _, length, err := c.Parser.ReadHeader(conn) diff --git a/server/internal/core/connection_runtime_test.go b/server/internal/core/connection_runtime_test.go index 662c9ab54..52c2660b8 100644 --- a/server/internal/core/connection_runtime_test.go +++ b/server/internal/core/connection_runtime_test.go @@ -44,6 +44,14 @@ func (w errorWriter) Write([]byte) (int, error) { return 0, w.err } +type errorReader struct { + err error +} + +func (r errorReader) Read([]byte) (int, error) { + return 0, r.err +} + type testAddr string func (a testAddr) Network() string { return "tcp" } @@ -55,6 +63,32 @@ type testConnRWC struct { func (c testConnRWC) RemoteAddr() net.Addr { return testAddr("127.0.0.1:0") } +type deadlineRecordingRWC struct { + io.ReadWriteCloser + mu sync.Mutex + deadlines []time.Time + deadlineErrors []error +} + +func (c *deadlineRecordingRWC) RemoteAddr() net.Addr { return testAddr("127.0.0.1:0") } + +func (c *deadlineRecordingRWC) SetReadDeadline(deadline time.Time) error { + c.mu.Lock() + call := len(c.deadlines) + c.deadlines = append(c.deadlines, deadline) + defer c.mu.Unlock() + if call < len(c.deadlineErrors) { + return c.deadlineErrors[call] + } + return nil +} + +func (c *deadlineRecordingRWC) readDeadlines() []time.Time { + c.mu.Lock() + defer c.mu.Unlock() + return append([]time.Time(nil), c.deadlines...) +} + func TestConnectionReceiveLoopFailureMarksConnectionDead(t *testing.T) { conn := &Connection{ SessionID: "session-a", @@ -109,6 +143,108 @@ func TestConnectionSendReturnsWriteError(t *testing.T) { } } +func TestConnectionBuildResponseBoundsPayloadRead(t *testing.T) { + rwc := &deadlineRecordingRWC{ + ReadWriteCloser: cryptostream.WrapReadWriteCloser(bytes.NewReader([]byte{1, 2}), io.Discard, nil), + } + packetParser := testPacketParser{} + streamConn := &cryptostream.Conn{ + ReadWriteCloser: rwc, + Parser: &parser.MessageParser{ + Implant: "test", + PacketParser: packetParser, + }, + } + connection := &Connection{ + SessionID: "payload-deadline", + PipelineID: "missing-forward", + Parser: streamConn.Parser, + } + + before := time.Now() + if err := connection.buildResponse(streamConn, 2); err != nil { + t.Fatalf("buildResponse failed: %v", err) + } + deadlines := rwc.readDeadlines() + if len(deadlines) != 2 { + t.Fatalf("read deadline calls = %d, want 2", len(deadlines)) + } + if deadlines[0].Before(before.Add(payloadReadBaseTimeout)) { + t.Fatalf("payload read deadline = %v, want at least base timeout", deadlines[0]) + } + if !deadlines[1].IsZero() { + t.Fatalf("cleared read deadline = %v, want zero", deadlines[1]) + } +} + +func TestPayloadReadTimeoutScalesWithDeclaredLength(t *testing.T) { + if got := payloadReadTimeout(1); got != payloadReadBaseTimeout+time.Second { + t.Fatalf("timeout for 1 byte = %v", got) + } + if got := payloadReadTimeout(payloadReadMinBytesPerSecond * 2); got != payloadReadBaseTimeout+2*time.Second { + t.Fatalf("timeout for two transfer units = %v", got) + } +} + +func TestConnectionBuildResponseReportsDeadlineFailures(t *testing.T) { + setErr := errors.New("set deadline failed") + readErr := errors.New("read failed") + clearErr := errors.New("clear deadline failed") + tests := []struct { + name string + reader io.Reader + deadlineErrors []error + wantErrors []error + }{ + { + name: "set deadline", + reader: bytes.NewReader([]byte{1, 2}), + deadlineErrors: []error{setErr}, + wantErrors: []error{setErr}, + }, + { + name: "clear deadline", + reader: bytes.NewReader([]byte{1, 2}), + deadlineErrors: []error{nil, clearErr}, + wantErrors: []error{clearErr}, + }, + { + name: "read and clear deadline", + reader: errorReader{err: readErr}, + deadlineErrors: []error{nil, clearErr}, + wantErrors: []error{readErr, clearErr}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rwc := &deadlineRecordingRWC{ + ReadWriteCloser: cryptostream.WrapReadWriteCloser(tt.reader, io.Discard, nil), + deadlineErrors: tt.deadlineErrors, + } + streamConn := &cryptostream.Conn{ + ReadWriteCloser: rwc, + Parser: &parser.MessageParser{ + Implant: "test", + PacketParser: testPacketParser{}, + }, + } + connection := &Connection{ + SessionID: "payload-deadline-error", + PipelineID: "missing-forward", + Parser: streamConn.Parser, + } + + err := connection.buildResponse(streamConn, 2) + for _, wantErr := range tt.wantErrors { + if !errors.Is(err, wantErr) { + t.Fatalf("buildResponse error = %v, want error %v", err, wantErr) + } + } + }) + } +} + func TestConnectionsRemoveDeletesConnection(t *testing.T) { pool := &connections{connections: &sync.Map{}} conn := &Connection{SessionID: "session-remove"} diff --git a/server/internal/core/event.go b/server/internal/core/event.go index 659e1018d..b05317f35 100644 --- a/server/internal/core/event.go +++ b/server/internal/core/event.go @@ -29,6 +29,7 @@ const ( var ( ErrEventBrokerUnavailable = errors.New("event broker unavailable") ErrEventBrokerQueueFull = errors.New("event broker queue full") + ErrEventSubscriberSlow = errors.New("event subscriber queue full") eventBrokerRestartBackoff = 200 * time.Millisecond ) @@ -279,17 +280,21 @@ func (event *Event) ToProtobuf() *clientpb.Event { } } +type eventSubscription struct { + events chan Event + ready chan struct{} + history chan []Event +} + type eventBroker struct { stop chan struct{} publish chan Event - subscribe chan chan Event + subscribe chan eventSubscription unsubscribe chan chan Event send chan Event notifier inotify.Notifier - lock *sync.Mutex cache *RingCache - ready sync.Map subscribeV2 chan eventV2SubscribeRequest unsubscribeV2 chan chan SequencedEvent @@ -303,6 +308,9 @@ type eventBroker struct { alive atomic.Bool managed atomic.Bool startOnce sync.Once + stopOnce sync.Once + publishMu sync.Mutex + stopped bool } func (broker *eventBroker) run() error { @@ -313,12 +321,7 @@ func (broker *eventBroker) run() error { defer func() { broker.alive.Store(false) for sub := range subscribers { - func(ch chan Event) { - defer func() { - _ = recover() - }() - close(ch) - }(sub) + closeEventSubscriber(sub) } for sub := range v2Subscribers { func(ch chan SequencedEvent) { @@ -332,12 +335,20 @@ func (broker *eventBroker) run() error { for { select { case <-broker.stop: + broker.recordAcceptedEvents() return nil - case sub := <-broker.subscribe: - subscribers[sub] = struct{}{} - if ready, ok := broker.ready.LoadAndDelete(sub); ok { - close(ready.(chan struct{})) + case subscription := <-broker.subscribe: + var history []Event + if subscription.history != nil { + for _, event := range broker.GetAll() { + history = append(history, *event) + } + } + subscribers[subscription.events] = struct{}{} + if subscription.history != nil { + subscription.history <- history } + close(subscription.ready) case sub := <-broker.unsubscribe: delete(subscribers, sub) case request := <-broker.subscribeV2: @@ -352,21 +363,19 @@ func (broker *eventBroker) run() error { case sub := <-broker.unsubscribeV2: delete(v2Subscribers, sub) case event := <-broker.publish: + sequenced := broker.recordAcceptedEvent(event) if event.Important { logs.Log.Infof("event.%s - %s", event.EventType, event.String()) } else if event.EventType != consts.EventHeartbeat { logs.Log.Debugf("event.%s - %s", event.EventType, event.String()) } - broker.lock.Lock() for sub := range subscribers { if err := broker.dispatch(sub, event); err != nil { delete(subscribers, sub) - logs.Log.Warnf("drop broken event subscriber: %s", ErrorText(err)) + closeEventSubscriber(sub) + logs.Log.Warnf("disconnect event subscriber: %s", ErrorText(err)) } } - broker.lock.Unlock() - - sequenced := broker.recordV2Event(event) for channel, subscriber := range v2Subscribers { if err := broker.dispatchV2(subscriber, sequenced); err != nil { delete(v2Subscribers, channel) @@ -377,6 +386,32 @@ func (broker *eventBroker) run() error { } } +func (broker *eventBroker) recordAcceptedEvents() { + for { + select { + case event := <-broker.publish: + broker.recordAcceptedEvent(event) + default: + return + } + } +} + +func closeEventSubscriber(sub chan Event) { + defer func() { + _ = recover() + }() + close(sub) +} + +func (broker *eventBroker) recordAcceptedEvent(event Event) SequencedEvent { + if event.Important { + eventCopy := event + broker.cache.Add(&eventCopy) + } + return broker.recordV2Event(event) +} + func (broker *eventBroker) ensureV2() { broker.v2Once.Do(func() { if broker.subscribeV2 == nil { @@ -547,10 +582,12 @@ func (broker *eventBroker) dispatch(sub chan Event, event Event) (err error) { }() select { case sub <- event: + return nil default: - // channel full, drop event; subscriber may catch up later + // Preserve legacy V1 behavior by dropping the current event without + // disconnecting the stream. V2 consumers receive reset/replay signals. + return nil } - return nil } func (broker *eventBroker) Start() { @@ -575,30 +612,65 @@ func (broker *eventBroker) Start() { // Stop - Close the broker channel func (broker *eventBroker) Stop() { - close(broker.stop) + if broker == nil { + return + } + broker.stopOnce.Do(func() { + broker.publishMu.Lock() + broker.stopped = true + close(broker.stop) + broker.publishMu.Unlock() + }) } // Subscribe - Generate a new subscription channel -func (broker *eventBroker) Subscribe() chan Event { +func (broker *eventBroker) Subscribe() (chan Event, error) { + events, _, err := broker.subscribeEvents(false) + return events, err +} + +func (broker *eventBroker) SubscribeWithHistory() (chan Event, []Event, error) { + return broker.subscribeEvents(true) +} + +func (broker *eventBroker) subscribeEvents(includeHistory bool) (chan Event, []Event, error) { events := make(chan Event, eventBufSize) ready := make(chan struct{}) - broker.ready.Store(events, ready) - broker.subscribe <- events - // Do not return until run has installed the subscriber. Callers may publish - // immediately after Subscribe and must not lose that first event. - <-ready - return events + var history chan []Event + if includeHistory { + history = make(chan []Event, 1) + } + select { + case broker.subscribe <- eventSubscription{events: events, ready: ready, history: history}: + case <-broker.stop: + return nil, nil, ErrEventBrokerUnavailable + } + select { + case <-ready: + if history != nil { + return events, <-history, nil + } + return events, nil, nil + case <-broker.stop: + return nil, nil, ErrEventBrokerUnavailable + } } // Unsubscribe - Remove a subscription channel func (broker *eventBroker) Unsubscribe(events chan Event) { - broker.unsubscribe <- events + select { + case broker.unsubscribe <- events: + case <-broker.stop: + } //close(events) } // SubscribeV2 registers a resumable event subscription. The first value in // the returned channel is always a Ready envelope containing stream bounds. -func (broker *eventBroker) SubscribeV2(options EventSubscription) chan SequencedEvent { +func (broker *eventBroker) SubscribeV2(options EventSubscription) (chan SequencedEvent, error) { + if broker == nil { + return nil, ErrEventBrokerUnavailable + } broker.ensureV2() capacity := broker.historyCapacity + 1 if capacity < eventBufSize { @@ -606,18 +678,33 @@ func (broker *eventBroker) SubscribeV2(options EventSubscription) chan Sequenced } events := make(chan SequencedEvent, capacity) ready := make(chan struct{}) - broker.subscribeV2 <- eventV2SubscribeRequest{ + request := eventV2SubscribeRequest{ options: options, channel: events, ready: ready, } - <-ready - return events + select { + case broker.subscribeV2 <- request: + case <-broker.stop: + return nil, ErrEventBrokerUnavailable + } + select { + case <-ready: + return events, nil + case <-broker.stop: + return nil, ErrEventBrokerUnavailable + } } func (broker *eventBroker) UnsubscribeV2(events chan SequencedEvent) { + if broker == nil { + return + } broker.ensureV2() - broker.unsubscribeV2 <- events + select { + case broker.unsubscribeV2 <- events: + case <-broker.stop: + } } // Publish - Push a message to all subscribers @@ -634,17 +721,26 @@ func (broker *eventBroker) TryPublish(event Event) error { if broker == nil { return ErrEventBrokerUnavailable } - if event.Important { - broker.cache.Add(&event) + broker.publishMu.Lock() + if broker.stopped { + broker.publishMu.Unlock() + return ErrEventBrokerUnavailable } if broker.managed.Load() && !broker.alive.Load() { + broker.publishMu.Unlock() return ErrEventBrokerUnavailable } select { case broker.publish <- event: default: + // History caching now happens on dequeue (recordAcceptedEvent), so an + // Important event dropped here is also absent from the replay history. + // This is an accepted trade-off under queue saturation; the drop is + // surfaced by Publish's error log rather than silently swallowed. + broker.publishMu.Unlock() return ErrEventBrokerQueueFull } + broker.publishMu.Unlock() if event.IsNotify { broker.Notify(event) } @@ -675,12 +771,11 @@ func NewBroker() *eventBroker { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } broker.Start() ticker := GlobalTicker diff --git a/server/internal/core/event_runtime_test.go b/server/internal/core/event_runtime_test.go index 4bf4488c3..3b2c06ae6 100644 --- a/server/internal/core/event_runtime_test.go +++ b/server/internal/core/event_runtime_test.go @@ -54,22 +54,12 @@ func TestEventBrokerRunStopClosesSubscribers(t *testing.T) { resultCh <- broker.run() }() - broker.subscribe <- sub - deadline := time.After(2 * time.Second) -subscribed: - for { - broker.publish <- Event{ - EventType: consts.EventBroadcast, - Op: "ready", - Message: "ready", - } - select { - case <-sub: - break subscribed - case <-time.After(20 * time.Millisecond): - case <-deadline: - t.Fatal("subscriber did not receive initial event") - } + ready := make(chan struct{}) + broker.subscribe <- eventSubscription{events: sub, ready: ready} + select { + case <-ready: + case <-time.After(2 * time.Second): + t.Fatal("subscriber was not registered") } close(broker.stop) @@ -88,6 +78,61 @@ subscribed: } } +func TestEventBrokerStopIsIdempotent(t *testing.T) { + broker := newTestBroker() + broker.Stop() + broker.Stop() +} + +func TestEventBrokerStopCachesAcceptedImportantEvents(t *testing.T) { + for i := 0; i < 100; i++ { + broker := newTestBroker() + result := make(chan error, 1) + go func() { result <- broker.run() }() + if err := broker.TryPublish(Event{ + EventType: consts.EventBroadcast, + Op: "accepted-before-stop", + Important: true, + }); err != nil { + t.Fatalf("iteration %d: TryPublish failed: %v", i, err) + } + broker.Stop() + select { + case err := <-result: + if err != nil { + t.Fatalf("iteration %d: broker.run error: %v", i, err) + } + case <-time.After(2 * time.Second): + t.Fatalf("iteration %d: broker.run did not stop", i) + } + if history := broker.GetAll(); len(history) != 1 || history[0].Op != "accepted-before-stop" { + t.Fatalf("iteration %d: cached history = %#v", i, history) + } + } +} + +func TestEventBrokerKeepsSlowLegacySubscriberConnected(t *testing.T) { + broker := newTestBroker() + sub := make(chan Event, eventBufSize) + for i := 0; i < eventBufSize; i++ { + sub <- Event{EventType: consts.EventBroadcast, Op: "queued"} + } + err := broker.dispatch(sub, Event{EventType: consts.EventBroadcast, Op: "overflow"}) + if err != nil { + t.Fatalf("dispatch overflow error = %v, want nil", err) + } + for i := 0; i < eventBufSize; i++ { + <-sub + } + want := Event{EventType: consts.EventBroadcast, Op: "recovered"} + if err := broker.dispatch(sub, want); err != nil { + t.Fatalf("dispatch after recovery error = %v", err) + } + if got := <-sub; got.Op != want.Op { + t.Fatalf("event after recovery = %q, want %q", got.Op, want.Op) + } +} + func TestEventBrokerTryPublishReturnsUnavailableWhenStopped(t *testing.T) { broker := newTestBroker() broker.managed.Store(true) @@ -119,7 +164,9 @@ func TestEventBrokerStartSurvivesBrokenSubscriber(t *testing.T) { closedSub := make(chan Event, 1) close(closedSub) - broker.subscribe <- closedSub + closedReady := make(chan struct{}) + broker.subscribe <- eventSubscription{events: closedSub, ready: closedReady} + <-closedReady if err := broker.TryPublish(Event{ EventType: consts.EventBroadcast, Op: "panic", @@ -128,7 +175,10 @@ func TestEventBrokerStartSurvivesBrokenSubscriber(t *testing.T) { t.Fatalf("TryPublish panic trigger error = %v", err) } - sub := broker.Subscribe() + sub, err := broker.Subscribe() + if err != nil { + t.Fatalf("Subscribe error = %v", err) + } defer broker.Unsubscribe(sub) deadline = time.After(2 * time.Second) @@ -177,13 +227,138 @@ subscribed: } } +func TestEventBrokerSubscribeWaitsForRegistration(t *testing.T) { + broker := newTestBroker() + returned := make(chan chan Event, 1) + go func() { + sub, _ := broker.Subscribe() + returned <- sub + }() + + select { + case <-returned: + t.Fatal("Subscribe returned before the broker registered the subscriber") + case <-time.After(50 * time.Millisecond): + } + + go func() { + _ = broker.run() + }() + defer broker.Stop() + + var sub chan Event + select { + case sub = <-returned: + case <-time.After(2 * time.Second): + t.Fatal("Subscribe did not return after registration") + } + + want := Event{EventType: consts.EventBroadcast, Op: "after-ready"} + if err := broker.TryPublish(want); err != nil { + t.Fatalf("TryPublish error = %v", err) + } + select { + case got := <-sub: + if got.Op != want.Op { + t.Fatalf("event op = %q, want %q", got.Op, want.Op) + } + case <-time.After(2 * time.Second): + t.Fatal("registered subscriber missed the first event") + } +} + +func TestEventBrokerSubscribeWithHistoryHasNoPublishBoundaryGap(t *testing.T) { + for i := 0; i < 100; i++ { + broker := newTestBroker() + go func() { _ = broker.run() }() + + boundary := Event{ + EventType: consts.EventBroadcast, + Op: "boundary", + Important: true, + } + if err := broker.TryPublish(boundary); err != nil { + broker.Stop() + t.Fatalf("iteration %d: publish boundary event: %v", i, err) + } + events, history, err := broker.SubscribeWithHistory() + if err != nil { + broker.Stop() + t.Fatalf("iteration %d: subscribe with history: %v", i, err) + } + marker := Event{EventType: consts.EventBroadcast, Op: "marker"} + if err := broker.TryPublish(marker); err != nil { + broker.Stop() + t.Fatalf("iteration %d: publish marker: %v", i, err) + } + + boundaryCount := 0 + for _, event := range history { + if event.Op == boundary.Op { + boundaryCount++ + } + } + for { + select { + case event, ok := <-events: + if !ok { + broker.Stop() + t.Fatalf("iteration %d: subscription closed before marker", i) + } + if event.Op == boundary.Op { + boundaryCount++ + } + if event.Op == marker.Op { + if boundaryCount != 1 { + broker.Stop() + t.Fatalf("iteration %d: boundary event count = %d, want 1", i, boundaryCount) + } + broker.Unsubscribe(events) + broker.Stop() + goto nextIteration + } + case <-time.After(2 * time.Second): + broker.Stop() + t.Fatalf("iteration %d: marker event not received", i) + } + } + nextIteration: + } +} + +func TestEventBrokerSubscribeAfterStopReturnsUnavailable(t *testing.T) { + broker := newTestBroker() + broker.Stop() + if _, err := broker.Subscribe(); !errors.Is(err, ErrEventBrokerUnavailable) { + t.Fatalf("Subscribe error = %v, want %v", err, ErrEventBrokerUnavailable) + } +} + +func TestEventBrokerSubscribeRacingStopDoesNotHang(t *testing.T) { + for i := 0; i < 100; i++ { + broker := newTestBroker() + go func() { _ = broker.run() }() + result := make(chan error, 1) + go func() { + _, err := broker.Subscribe() + result <- err + }() + broker.Stop() + select { + case <-result: + case <-time.After(2 * time.Second): + t.Fatalf("Subscribe hung while racing Stop at iteration %d", i) + } + } +} + func TestEventBrokerSubscribeIsReadyBeforeReturn(t *testing.T) { broker := newTestBroker() broker.Start() defer broker.Stop() for i := 0; i < 100; i++ { - sub := broker.Subscribe() + sub := mustSubscribeEvent(t, broker) event := Event{ EventType: consts.EventBroadcast, Op: fmt.Sprintf("ready-%d", i), @@ -208,7 +383,7 @@ func TestEventBrokerV2DeliversReadyThenOrderedEvents(t *testing.T) { broker.Start() defer broker.Stop() - sub := broker.SubscribeV2(EventSubscription{}) + sub := mustSubscribeEventV2(t, broker, EventSubscription{}) defer broker.UnsubscribeV2(sub) ready := waitSequencedEvent(t, sub) @@ -243,7 +418,7 @@ func TestEventBrokerV2ReplaysEventsAfterCursor(t *testing.T) { broker.Start() defer broker.Stop() - initial := broker.SubscribeV2(EventSubscription{}) + initial := mustSubscribeEventV2(t, broker, EventSubscription{}) ready := waitSequencedEvent(t, initial) if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: "one", Important: true}); err != nil { t.Fatalf("TryPublish(one) failed: %v", err) @@ -255,14 +430,14 @@ func TestEventBrokerV2ReplaysEventsAfterCursor(t *testing.T) { _ = waitSequencedEvent(t, initial) broker.UnsubscribeV2(initial) - barrier := broker.Subscribe() + barrier := mustSubscribeEvent(t, broker) if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: "three", Important: true}); err != nil { t.Fatalf("TryPublish(three) failed: %v", err) } _ = waitEvent(t, barrier) broker.Unsubscribe(barrier) - reconnected := broker.SubscribeV2(EventSubscription{ + reconnected := mustSubscribeEventV2(t, broker, EventSubscription{ StreamID: ready.StreamID, AfterSequence: one.Sequence, Replay: true, @@ -289,7 +464,7 @@ func TestEventBrokerV2SignalsResetWhenCursorFallsBehindHistory(t *testing.T) { broker.Start() defer broker.Stop() - initial := broker.SubscribeV2(EventSubscription{}) + initial := mustSubscribeEventV2(t, broker, EventSubscription{}) ready := waitSequencedEvent(t, initial) if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: "one", Important: true}); err != nil { t.Fatalf("TryPublish(one) failed: %v", err) @@ -297,7 +472,7 @@ func TestEventBrokerV2SignalsResetWhenCursorFallsBehindHistory(t *testing.T) { one := waitSequencedEvent(t, initial) broker.UnsubscribeV2(initial) - barrier := broker.Subscribe() + barrier := mustSubscribeEvent(t, broker) for _, op := range []string{"two", "three", "four"} { if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: op, Important: true}); err != nil { t.Fatalf("TryPublish(%s) failed: %v", op, err) @@ -306,7 +481,7 @@ func TestEventBrokerV2SignalsResetWhenCursorFallsBehindHistory(t *testing.T) { } broker.Unsubscribe(barrier) - reconnected := broker.SubscribeV2(EventSubscription{ + reconnected := mustSubscribeEventV2(t, broker, EventSubscription{ StreamID: ready.StreamID, AfterSequence: one.Sequence, Replay: true, @@ -328,7 +503,7 @@ func TestEventBrokerV2HeartbeatsDoNotEvictReplayableEvents(t *testing.T) { broker.Start() defer broker.Stop() - initial := broker.SubscribeV2(EventSubscription{}) + initial := mustSubscribeEventV2(t, broker, EventSubscription{}) ready := waitSequencedEvent(t, initial) if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: "one", Important: true}); err != nil { t.Fatalf("TryPublish(one) failed: %v", err) @@ -336,7 +511,7 @@ func TestEventBrokerV2HeartbeatsDoNotEvictReplayableEvents(t *testing.T) { one := waitSequencedEvent(t, initial) broker.UnsubscribeV2(initial) - barrier := broker.Subscribe() + barrier := mustSubscribeEvent(t, broker) for index := 0; index < 3; index++ { if err := broker.TryPublish(Event{EventType: consts.EventHeartbeat, Op: consts.CtrlHeartbeat1s}); err != nil { t.Fatalf("TryPublish(heartbeat %d) failed: %v", index, err) @@ -349,7 +524,7 @@ func TestEventBrokerV2HeartbeatsDoNotEvictReplayableEvents(t *testing.T) { _ = waitEvent(t, barrier) broker.Unsubscribe(barrier) - reconnected := broker.SubscribeV2(EventSubscription{ + reconnected := mustSubscribeEventV2(t, broker, EventSubscription{ StreamID: ready.StreamID, AfterSequence: one.Sequence, Replay: true, @@ -372,7 +547,7 @@ func TestEventBrokerV2TransientEventsDoNotEvictReplayableEvents(t *testing.T) { broker.Start() defer broker.Stop() - initial := broker.SubscribeV2(EventSubscription{}) + initial := mustSubscribeEventV2(t, broker, EventSubscription{}) ready := waitSequencedEvent(t, initial) if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: "one", Important: true}); err != nil { t.Fatalf("TryPublish(one) failed: %v", err) @@ -380,7 +555,7 @@ func TestEventBrokerV2TransientEventsDoNotEvictReplayableEvents(t *testing.T) { one := waitSequencedEvent(t, initial) broker.UnsubscribeV2(initial) - barrier := broker.Subscribe() + barrier := mustSubscribeEvent(t, broker) for index := 0; index < 5; index++ { if err := broker.TryPublish(Event{EventType: consts.EventSession, Op: consts.CtrlSessionCheckin}); err != nil { t.Fatalf("TryPublish(checkin %d) failed: %v", index, err) @@ -393,7 +568,7 @@ func TestEventBrokerV2TransientEventsDoNotEvictReplayableEvents(t *testing.T) { _ = waitEvent(t, barrier) broker.Unsubscribe(barrier) - reconnected := broker.SubscribeV2(EventSubscription{ + reconnected := mustSubscribeEventV2(t, broker, EventSubscription{ StreamID: ready.StreamID, AfterSequence: one.Sequence, Replay: true, @@ -418,11 +593,11 @@ func TestEventBrokerV2ReplaysPersistedSessionUpdate(t *testing.T) { EventBroker = broker defer func() { EventBroker = oldBroker }() - initial := broker.SubscribeV2(EventSubscription{}) + initial := mustSubscribeEventV2(t, broker, EventSubscription{}) ready := waitSequencedEvent(t, initial) broker.UnsubscribeV2(initial) - barrier := broker.Subscribe() + barrier := mustSubscribeEvent(t, broker) session := newTestSession("session-update-replay") session.PushUpdate("note updated") if event := waitEvent(t, barrier); !event.Important { @@ -430,7 +605,7 @@ func TestEventBrokerV2ReplaysPersistedSessionUpdate(t *testing.T) { } broker.Unsubscribe(barrier) - reconnected := broker.SubscribeV2(EventSubscription{ + reconnected := mustSubscribeEventV2(t, broker, EventSubscription{ StreamID: ready.StreamID, AfterSequence: ready.Sequence, Replay: true, @@ -471,7 +646,7 @@ func TestEventBrokerV2FiltersTopicsAndHeartbeats(t *testing.T) { broker.Start() defer broker.Stop() - sub := broker.SubscribeV2(EventSubscription{Topics: []string{consts.EventSession}}) + sub := mustSubscribeEventV2(t, broker, EventSubscription{Topics: []string{consts.EventSession}}) defer broker.UnsubscribeV2(sub) _ = waitSequencedEvent(t, sub) @@ -491,6 +666,24 @@ func TestEventBrokerV2FiltersTopicsAndHeartbeats(t *testing.T) { } } +func mustSubscribeEvent(t *testing.T, broker *eventBroker) chan Event { + t.Helper() + events, err := broker.Subscribe() + if err != nil { + t.Fatalf("Subscribe failed: %v", err) + } + return events +} + +func mustSubscribeEventV2(t *testing.T, broker *eventBroker, options EventSubscription) chan SequencedEvent { + t.Helper() + events, err := broker.SubscribeV2(options) + if err != nil { + t.Fatalf("SubscribeV2 failed: %v", err) + } + return events +} + func waitSequencedEvent(t *testing.T, ch <-chan SequencedEvent) SequencedEvent { t.Helper() select { diff --git a/server/internal/core/forward.go b/server/internal/core/forward.go index 3c98cd46c..1d5868031 100644 --- a/server/internal/core/forward.go +++ b/server/internal/core/forward.go @@ -3,6 +3,7 @@ package core import ( "context" "encoding/binary" + "errors" "fmt" "strconv" "sync" @@ -20,6 +21,10 @@ import ( "google.golang.org/protobuf/proto" ) +const forwardShutdownTimeout = 5 * time.Second + +var ErrForwardShutdownTimeout = errors.New("forward shutdown timed out") + func forwardRawIDBytes(rawID uint32) []byte { data := make([]byte, 4) binary.LittleEndian.PutUint32(data, rawID) @@ -109,13 +114,30 @@ func (f *forwarders) Remove(id string) error { if fw == nil { return nil } - f.forwarders.Delete(id) - fw.shutdown() - err := fw.Close() - if err != nil { - return err + return f.removeIfSame(id, fw) +} + +// RemoveIfSame atomically detaches fw when it is still the registered +// generation for id, then aborts only that generation. It intentionally does +// not call Pipeline.Close so callers can finish an in-flight start rollback +// before releasing their lifecycle gate. A replacement is always preserved. +func (f *forwarders) RemoveIfSame(id string, fw *Forward) error { + if fw == nil { + return nil } - return nil + f.forwarders.CompareAndDelete(id, fw) + return fw.Abort() +} + +func (f *forwarders) removeIfSame(id string, fw *Forward) error { + if fw == nil { + return nil + } + if !f.forwarders.CompareAndDelete(id, fw) { + return fw.Abort() + } + abortErr := fw.Abort() + return errors.Join(abortErr, fw.Close()) } func (f *forwarders) Send(id string, msg *Message) { @@ -129,6 +151,7 @@ func (f *forwarders) Send(id string, msg *Message) { func NewForward(rpc ForwardClient, pipeline Pipeline) (*Forward, error) { var err error + ctx, cancel := context.WithCancel(context.Background()) listenerID := "" if pb := pipeline.ToProtobuf(); pb != nil { listenerID = pb.ListenerId @@ -138,25 +161,31 @@ func NewForward(rpc ForwardClient, pipeline Pipeline) (*Forward, error) { ListenerRpc: rpc, Pipeline: pipeline, ListenerId: listenerID, - ctx: context.Background(), + ctx: ctx, + cancel: cancel, done: make(chan struct{}), + handlerDone: make(chan struct{}), } forward.alive.Store(true) - forward.Stream, err = rpc.OpenForwardStream(context.Background(), pipeline) + forward.Stream, err = rpc.OpenForwardStream(ctx, pipeline) if err != nil { + cancel() return nil, err } - GoGuarded("forward:"+pipeline.ID(), forward.Handler, forward.handleRuntimeError(), forward.shutdown) + GoGuarded("forward:"+pipeline.ID(), func() error { + defer close(forward.handlerDone) + return forward.Handler() + }, forward.handleRuntimeError(), forward.shutdown) return forward, nil } // Forward is a struct that handles messages from listener and server type Forward struct { - ctx context.Context - count int + ctx context.Context + cancel context.CancelFunc Pipeline ListenerId string Stream ForwardStream @@ -164,9 +193,12 @@ type Forward struct { ListenerRpc forwardRPCClient - alive atomic.Bool - done chan struct{} - closeOnce sync.Once + alive atomic.Bool + done chan struct{} + closeOnce sync.Once + handlerDone chan struct{} + abortOnce sync.Once + abortErr error } func (f *Forward) RuntimeKey() string { @@ -183,7 +215,6 @@ func (f *Forward) Add(msg *Message) { } select { case f.implantC <- msg: - f.count++ case <-f.done: logs.Log.Warnf("forward %s closed, dropping message from %s", f.ID(), msg.SessionID) } @@ -192,12 +223,53 @@ func (f *Forward) Add(msg *Message) { func (f *Forward) shutdown() { f.alive.Store(false) f.closeOnce.Do(func() { + if f.cancel != nil { + f.cancel() + } close(f.done) }) } -func (f *Forward) Count() int { - return f.count +// Abort closes an uncommitted forward without touching the global registry or +// closing its pipeline. It is safe to call more than once. +func (f *Forward) Abort() error { + ctx, cancel := context.WithTimeout(context.Background(), forwardShutdownTimeout) + defer cancel() + return f.AbortContext(ctx) +} + +func (f *Forward) AbortContext(ctx context.Context) error { + if f == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + f.abortOnce.Do(func() { + f.shutdown() + streamClosed := make(chan error, 1) + if stream, ok := f.Stream.(interface{ CloseSend() error }); ok { + go func() { + streamClosed <- stream.CloseSend() + }() + } else { + streamClosed <- nil + } + select { + case f.abortErr = <-streamClosed: + case <-ctx.Done(): + f.abortErr = fmt.Errorf("%w while closing stream: %v", ErrForwardShutdownTimeout, ctx.Err()) + return + } + if f.handlerDone != nil { + select { + case <-f.handlerDone: + case <-ctx.Done(): + f.abortErr = errors.Join(f.abortErr, fmt.Errorf("%w: %v", ErrForwardShutdownTimeout, ctx.Err())) + } + } + }) + return f.abortErr } func (f *Forward) Context(sid string) context.Context { @@ -211,7 +283,17 @@ func (f *Forward) Context(sid string) context.Context { // Handler is a loop that handles messages from implant func (f *Forward) Handler() error { - for msg := range f.implantC { + for { + var msg *Message + select { + case <-f.done: + return nil + case received, ok := <-f.implantC: + if !ok { + return nil + } + msg = received + } _, err := f.ListenerRpc.Checkin(f.Context(msg.SessionID), &implantpb.Ping{}) if err != nil { logs.Log.Warnf("forward %s checkin failed for session %s: %v", f.ID(), msg.SessionID, err) @@ -260,7 +342,6 @@ func (f *Forward) Handler() error { } } } - return nil } func (f *Forward) handleRuntimeError() GoErrorHandler { diff --git a/server/internal/core/forward_runtime_test.go b/server/internal/core/forward_runtime_test.go index f69cc2b48..591fd80ee 100644 --- a/server/internal/core/forward_runtime_test.go +++ b/server/internal/core/forward_runtime_test.go @@ -6,6 +6,7 @@ import ( "sync" "sync/atomic" "testing" + "time" "github.com/chainreactors/IoM-go/proto/client/clientpb" implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" @@ -37,6 +38,47 @@ func (testForwardStream) Recv() (*clientpb.SpiteRequest, error) { return nil, errors.New("not used") } +type abortForwardStream struct { + closeCalls atomic.Int32 +} + +type blockingCloseForwardStream struct { + started chan struct{} + release chan struct{} + done chan struct{} +} + +func (*blockingCloseForwardStream) Send(*clientpb.SpiteResponse) error { return nil } +func (*blockingCloseForwardStream) Recv() (*clientpb.SpiteRequest, error) { + return nil, errors.New("not used") +} +func (s *blockingCloseForwardStream) CloseSend() error { + defer close(s.done) + close(s.started) + <-s.release + return nil +} + +func (*abortForwardStream) Send(*clientpb.SpiteResponse) error { return nil } +func (*abortForwardStream) Recv() (*clientpb.SpiteRequest, error) { + return nil, errors.New("not used") +} +func (s *abortForwardStream) CloseSend() error { + s.closeCalls.Add(1) + return nil +} + +type abortForwardClient struct { + *testForwardRPC + stream *abortForwardStream + ctx context.Context +} + +func (c *abortForwardClient) OpenForwardStream(ctx context.Context, _ Pipeline) (ForwardStream, error) { + c.ctx = ctx + return c.stream, nil +} + type testPipeline struct { id string closeErr error @@ -52,6 +94,86 @@ func (p testPipeline) ToProtobuf() *clientpb.Pipeline { return &clientpb.Pipeline{Name: p.id} } +func TestForwardAbortClosesStreamCancelsContextAndStopsHandler(t *testing.T) { + stream := &abortForwardStream{} + client := &abortForwardClient{testForwardRPC: &testForwardRPC{}, stream: stream} + forward, err := NewForward(client, testPipeline{id: "pipe-abort"}) + if err != nil { + t.Fatalf("NewForward failed: %v", err) + } + + if err := forward.Abort(); err != nil { + t.Fatalf("Abort failed: %v", err) + } + if err := forward.Abort(); err != nil { + t.Fatalf("second Abort failed: %v", err) + } + if got := stream.closeCalls.Load(); got != 1 { + t.Fatalf("CloseSend calls = %d, want 1", got) + } + select { + case <-client.ctx.Done(): + default: + t.Fatal("OpenForwardStream context was not canceled") + } + select { + case <-forward.handlerDone: + default: + t.Fatal("forward handler was still running after Abort returned") + } + if got := Forwarders.Get(forward.RuntimeKey()); got != nil { + t.Fatalf("aborted forward remained in registry: %#v", got) + } +} + +func TestForwardAbortContextReturnsWhenHandlerDoesNotStop(t *testing.T) { + forward := &Forward{ + Pipeline: testPipeline{id: "pipe-timeout"}, + Stream: testForwardStream{}, + done: make(chan struct{}), + handlerDone: make(chan struct{}), + } + forward.alive.Store(true) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := forward.AbortContext(ctx); !errors.Is(err, ErrForwardShutdownTimeout) { + t.Fatalf("AbortContext error = %v, want ErrForwardShutdownTimeout", err) + } +} + +func TestForwardAbortContextReturnsWhenCloseSendDoesNotStop(t *testing.T) { + stream := &blockingCloseForwardStream{ + started: make(chan struct{}), + release: make(chan struct{}), + done: make(chan struct{}), + } + forward := &Forward{ + Pipeline: testPipeline{id: "pipe-close-timeout"}, + Stream: stream, + done: make(chan struct{}), + } + forward.alive.Store(true) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + err := forward.AbortContext(ctx) + if !errors.Is(err, ErrForwardShutdownTimeout) { + t.Fatalf("AbortContext error = %v, want ErrForwardShutdownTimeout", err) + } + select { + case <-stream.started: + default: + t.Fatal("CloseSend was not called") + } + close(stream.release) + select { + case <-stream.done: + case <-time.After(time.Second): + t.Fatal("CloseSend did not finish after release") + } +} + func TestForwardHandlerReturnsStreamSendError(t *testing.T) { want := errors.New("stream send failed") forward := &Forward{ @@ -97,6 +219,41 @@ func TestForwardersRemoveDeletesOnCloseError(t *testing.T) { } } +func TestForwardersRemoveIfSamePreservesReplacement(t *testing.T) { + store := &forwarders{forwarders: &sync.Map{}} + oldForward := &Forward{ + Pipeline: testPipeline{id: "shared"}, + Stream: testForwardStream{}, + done: make(chan struct{}), + } + newForward := &Forward{ + Pipeline: testPipeline{id: "shared"}, + Stream: testForwardStream{}, + done: make(chan struct{}), + } + oldForward.alive.Store(true) + newForward.alive.Store(true) + store.Add(oldForward) + store.Add(newForward) + + if err := store.removeIfSame("shared", oldForward); err != nil { + t.Fatalf("removeIfSame returned error: %v", err) + } + if got := store.Get("shared"); got != newForward { + t.Fatalf("replacement forward = %#v, want %#v", got, newForward) + } + select { + case <-oldForward.done: + default: + t.Fatal("stale forward was not aborted") + } + select { + case <-newForward.done: + t.Fatal("replacement forward was aborted") + default: + } +} + func TestForwardHandlerCheckinCalledOncePerMessage(t *testing.T) { rpc := &testForwardRPC{} stream := &capturingForwardStream{} diff --git a/server/internal/core/safe_test.go b/server/internal/core/safe_test.go index aed714365..177d2f544 100644 --- a/server/internal/core/safe_test.go +++ b/server/internal/core/safe_test.go @@ -18,12 +18,11 @@ func newTestBroker() *eventBroker { return &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } } diff --git a/server/internal/core/secure_test.go b/server/internal/core/secure_test.go index 0ca8594de..d670ce5a1 100644 --- a/server/internal/core/secure_test.go +++ b/server/internal/core/secure_test.go @@ -1,6 +1,7 @@ package core import ( + "fmt" "sync" "testing" @@ -159,6 +160,48 @@ func TestSecureManager_UpdateKeyPair(t *testing.T) { sm.mu.Unlock() } +func TestSessionUpdateKeyPairConcurrentFieldsStaySynchronized(t *testing.T) { + sess := newSecureTestSession() + sess.SecureManager = NewSecureSpiteManager(sess) + + const iterations = 500 + for i := 0; i < iterations; i++ { + publicKey := fmt.Sprintf("public-%d", i) + privateKey := fmt.Sprintf("private-%d", i) + start := make(chan struct{}) + var updates sync.WaitGroup + updates.Add(2) + + go func() { + defer updates.Done() + <-start + sess.UpdatePublicKey(publicKey) + }() + go func() { + defer updates.Done() + <-start + sess.UpdatePrivateKey(privateKey) + }() + + close(start) + updates.Wait() + + sess.stateMu.RLock() + stored := sess.KeyPair + sess.stateMu.RUnlock() + sess.SecureManager.mu.Lock() + managed := sess.SecureManager.keyPair + sess.SecureManager.mu.Unlock() + + if stored != managed { + t.Fatalf("iteration %d: session and manager key pairs reference different values", i) + } + if stored.GetPublicKey() != publicKey || stored.GetPrivateKey() != privateKey { + t.Fatalf("iteration %d: key pair = %q/%q, want %q/%q", i, stored.GetPublicKey(), stored.GetPrivateKey(), publicKey, privateKey) + } + } +} + // TestSecureManager_UpdateKeyPairNil verifies updating with nil key pair does not panic. func TestSecureManager_UpdateKeyPairNil(t *testing.T) { t.Parallel() diff --git a/server/internal/core/session.go b/server/internal/core/session.go index 8f0f63892..6804d1194 100644 --- a/server/internal/core/session.go +++ b/server/internal/core/session.go @@ -37,9 +37,7 @@ var ( // example when the command is not supported on the platform ErrUnknownMessageType = errors.New("unknown message type") - // ErrImplantSendTimeout - The implant did not respond prior to timeout deadline - ErrImplantSendTimeout = errors.New("implant timeout") - ErrSpiteStreamClosed = errors.New("spite stream writer closed") + ErrSpiteStreamClosed = errors.New("spite stream writer closed") // DB function variables — swappable in tests for mocking sessionDBSave = func(s *models.Session) error { return db.SaveSessionModel(s) } @@ -273,6 +271,7 @@ type Session struct { keepaliveMu sync.Mutex keepaliveEnabled bool anyMu sync.RWMutex + stateMu sync.RWMutex Ctx context.Context Cancel context.CancelFunc @@ -355,6 +354,8 @@ func (w *SpiteStreamWriter) finish(err error) { } func (s *Session) Abstract() string { + s.stateMu.RLock() + defer s.stateMu.RUnlock() if s.Os == nil { return fmt.Sprintf("%s(%s)", s.Name, s.ID) } else { @@ -495,6 +496,12 @@ func (s *Session) expectedCheckinDeadline() (time.Time, bool) { if s == nil { return time.Time{}, false } + s.stateMu.RLock() + defer s.stateMu.RUnlock() + return s.expectedCheckinDeadlineLocked() +} + +func (s *Session) expectedCheckinDeadlineLocked() (time.Time, bool) { lastCheckin := s.LastCheckinUnix() if lastCheckin <= 0 { return time.Time{}, false @@ -538,10 +545,16 @@ func (s *Session) isAliveAt(now time.Time) bool { if s == nil { return false } + s.stateMu.RLock() + defer s.stateMu.RUnlock() + return s.isAliveAtLocked(now) +} + +func (s *Session) isAliveAtLocked(now time.Time) bool { if s.Type == consts.BindPipeline { return true } - deadline, ok := s.expectedCheckinDeadline() + deadline, ok := s.expectedCheckinDeadlineLocked() if !ok { return true } @@ -553,66 +566,65 @@ func (s *Session) isAlived() bool { } func (s *Session) ToProtobuf() *clientpb.Session { - session := &clientpb.Session{ - Type: s.Type, - SessionId: s.ID, - RawId: s.RawID, - Note: s.Note, - Name: s.Name, - GroupName: s.Group, - IsAlive: s.isAlived(), - IsPrivilege: s.IsPrivilege, - Target: s.Target, - PipelineId: s.PipelineID, - ListenerId: s.ListenerID, - Os: s.Os, - Process: s.Process, - LastCheckin: s.LastCheckinUnix(), - Filepath: s.SessionContext.Filepath, - Workdir: s.SessionContext.WorkDir, - Locate: s.SessionContext.Locale, - Proxy: s.SessionContext.ProxyURL, - Timer: &implantpb.Timer{Expression: s.Expression, Jitter: s.Jitter}, - CreatedAt: s.CreatedAt.Unix(), - Tasks: s.Tasks.ToProtobuf(), - Modules: s.Modules, - Addons: s.Addons, - KeyPair: s.KeyPair, // 添加密钥对 - Data: s.Marshal(), - } - session.IsInitialized = s.Initialized - return session + pb := s.ToProtobufLite() + pb.Tasks = s.Tasks.ToProtobuf() + return pb } func (s *Session) ToProtobufLite() *clientpb.Session { - session := &clientpb.Session{ - Type: s.Type, - SessionId: s.ID, - RawId: s.RawID, - Note: s.Note, - Name: s.Name, - GroupName: s.Group, - IsAlive: s.isAlived(), - IsPrivilege: s.IsPrivilege, - Target: s.Target, - PipelineId: s.PipelineID, - ListenerId: s.ListenerID, - Os: s.Os, - Process: s.Process, - LastCheckin: s.LastCheckinUnix(), - Filepath: s.SessionContext.Filepath, - Workdir: s.SessionContext.WorkDir, - Locate: s.SessionContext.Locale, - Proxy: s.SessionContext.ProxyURL, - Timer: &implantpb.Timer{Expression: s.Expression, Jitter: s.Jitter}, - CreatedAt: s.CreatedAt.Unix(), - Modules: s.Modules, - Addons: s.Addons, - KeyPair: s.KeyPair, - Data: s.Marshal(), - } - session.IsInitialized = s.Initialized - return session + s.stateMu.RLock() + defer s.stateMu.RUnlock() + return s.toProtobufLiteLocked(time.Now()) +} + +func (s *Session) toProtobufLiteLocked(now time.Time) *clientpb.Session { + modules := append([]string(nil), s.Modules...) + addons := cloneSessionAddons(s.Addons) + var keyPair *clientpb.KeyPair + if s.KeyPair != nil { + keyPair = proto.Clone(s.KeyPair).(*clientpb.KeyPair) + } + + return &clientpb.Session{ + Type: s.Type, + SessionId: s.ID, + RawId: s.RawID, + Note: s.Note, + Name: s.Name, + GroupName: s.Group, + IsAlive: s.isAliveAtLocked(now), + IsPrivilege: s.IsPrivilege, + Target: s.Target, + PipelineId: s.PipelineID, + ListenerId: s.ListenerID, + Os: cloneSysInfoOS(s.Os), + Process: cloneSysInfoProcess(s.Process), + LastCheckin: s.LastCheckinUnix(), + Filepath: s.SessionContext.Filepath, + Workdir: s.SessionContext.WorkDir, + Locate: s.SessionContext.Locale, + Proxy: s.SessionContext.ProxyURL, + Timer: &implantpb.Timer{Expression: s.Expression, Jitter: s.Jitter}, + CreatedAt: s.CreatedAt.Unix(), + Modules: modules, + Addons: addons, + KeyPair: keyPair, + Data: s.Marshal(), + IsInitialized: s.Initialized, + } +} + +func cloneSessionAddons(addons []*implantpb.Addon) []*implantpb.Addon { + if addons == nil { + return nil + } + cloned := make([]*implantpb.Addon, len(addons)) + for i, addon := range addons { + if addon != nil { + cloned[i] = proto.Clone(addon).(*implantpb.Addon) + } + } + return cloned } func (s *Session) Save() error { @@ -620,6 +632,8 @@ func (s *Session) Save() error { } func (s *Session) ToModel() *models.Session { + s.stateMu.RLock() + name := s.Name sessModel := &models.Session{ SessionID: s.ID, RawID: s.RawID, @@ -630,11 +644,12 @@ func (s *Session) ToModel() *models.Session { Type: s.Type, PipelineID: s.PipelineID, ListenerID: s.ListenerID, - IsAlive: s.isAlived(), + IsAlive: s.isAliveAtLocked(time.Now()), LastCheckin: s.LastCheckinUnix(), DataString: s.Marshal(), } - artifact, err := sessionDBGetArtifact(s.Name) + s.stateMu.RUnlock() + artifact, err := sessionDBGetArtifact(name) if err == nil && artifact.ProfileName != "" { if _, profileErr := sessionDBGetProfile(artifact.ProfileName); profileErr == nil { sessModel.ProfileName = artifact.ProfileName @@ -664,6 +679,169 @@ func (s *Session) SaveAndNotify(msg string) error { return nil } +func (s *Session) SetNote(note string) { + s.stateMu.Lock() + s.Note = note + s.stateMu.Unlock() +} + +func (s *Session) SetGroup(group string) { + s.stateMu.Lock() + s.Group = group + s.stateMu.Unlock() +} + +func (s *Session) SetTimer(expression string, jitter float64) { + if s == nil { + return + } + s.stateMu.Lock() + s.ensureSessionContextLocked() + s.Expression = expression + s.Jitter = jitter + s.stateMu.Unlock() +} + +func (s *Session) TimerSnapshot() (string, float64) { + if s == nil { + return "", 0 + } + s.stateMu.RLock() + defer s.stateMu.RUnlock() + if s.SessionContext == nil || s.SessionInfo == nil { + return "", 0 + } + return s.Expression, s.Jitter +} + +func (s *Session) ensureSessionContextLocked() { + if s.SessionContext == nil { + s.SessionContext = &client.SessionContext{} + } + if s.SessionInfo == nil { + s.SessionInfo = &client.SessionInfo{} + } +} + +func (s *Session) RoutingSnapshot() (string, string) { + _, listenerID, pipelineID := s.ConnectionSnapshot() + return listenerID, pipelineID +} + +func (s *Session) ConnectionSnapshot() (string, string, string) { + if s == nil { + return "", "", "" + } + s.stateMu.RLock() + defer s.stateMu.RUnlock() + return s.Target, s.ListenerID, s.PipelineID +} + +func (s *Session) ReplaceAddons(addons []*implantpb.Addon) { + if s == nil { + return + } + s.stateMu.Lock() + s.ensureSessionContextLocked() + s.Addons = mergeSessionAddons(nil, addons) + s.stateMu.Unlock() +} + +func (s *Session) MergeAddons(addons []*implantpb.Addon) { + if s == nil || len(addons) == 0 { + return + } + s.stateMu.Lock() + s.ensureSessionContextLocked() + s.Addons = mergeSessionAddons(s.Addons, addons) + s.stateMu.Unlock() +} + +func (s *Session) AddonsSnapshot() []*implantpb.Addon { + if s == nil { + return nil + } + s.stateMu.RLock() + defer s.stateMu.RUnlock() + if s.SessionContext == nil { + return nil + } + return cloneSessionAddons(s.Addons) +} + +func (s *Session) HasAddon(name string) bool { + if s == nil || name == "" { + return false + } + s.stateMu.RLock() + defer s.stateMu.RUnlock() + if s.SessionContext == nil { + return false + } + for _, addon := range s.Addons { + if addon != nil && addon.GetName() == name { + return true + } + } + return false +} + +func mergeSessionAddons(existing, incoming []*implantpb.Addon) []*implantpb.Addon { + merged := make([]*implantpb.Addon, 0, len(existing)+len(incoming)) + indexes := make(map[string]int, len(existing)+len(incoming)) + appendOrReplace := func(addon *implantpb.Addon) { + if addon == nil || addon.GetName() == "" { + return + } + cloned := proto.Clone(addon).(*implantpb.Addon) + if index, ok := indexes[addon.GetName()]; ok { + merged[index] = cloned + return + } + indexes[addon.GetName()] = len(merged) + merged = append(merged, cloned) + } + for _, addon := range existing { + appendOrReplace(addon) + } + for _, addon := range incoming { + appendOrReplace(addon) + } + return merged +} + +func (s *Session) ApplyModules(modules *implantpb.Modules, appendOnly bool) { + if modules == nil { + return + } + s.stateMu.Lock() + defer s.stateMu.Unlock() + + if appendOnly { + s.Modules = append(s.Modules, modules.GetModules()...) + if bundleMap := modules.GetBundleMap(); len(bundleMap) > 0 { + if s.BundleMap == nil { + s.BundleMap = make(map[string]string) + } + for key, value := range bundleMap { + s.BundleMap[key] = value + } + } + return + } + + s.Modules = append([]string(nil), modules.GetModules()...) + bundleMap := modules.GetBundleMap() + if bundleMap == nil { + s.BundleMap = nil + return + } + s.BundleMap = make(map[string]string, len(bundleMap)) + for key, value := range bundleMap { + s.BundleMap[key] = value + } +} + // GetPipelineEncryptionKey returns the transport encryption key from the // session's pipeline config. Returns "" if not found. func (s *Session) GetPipelineEncryptionKey() string { @@ -688,16 +866,27 @@ func (s *Session) GetPacketLength() int { } func (s *Session) findPipeline() (*clientpb.Pipeline, bool) { - if s == nil || s.PipelineID == "" { + if s == nil { + return nil, false + } + s.stateMu.RLock() + pipelineID := s.PipelineID + listenerID := s.ListenerID + s.stateMu.RUnlock() + if pipelineID == "" { return nil, false } - if s.ListenerID != "" { - return Listeners.FindByListener(s.ListenerID, s.PipelineID) + if listenerID != "" { + return Listeners.FindByListener(listenerID, pipelineID) } - return Listeners.Find(s.PipelineID) + return Listeners.Find(pipelineID) } func (s *Session) Update(req *clientpb.RegisterSession) { + if req == nil || req.RegisterData == nil { + return + } + s.stateMu.Lock() s.Name = req.RegisterData.Name s.PipelineID = req.PipelineId s.ListenerID = req.ListenerId @@ -706,15 +895,18 @@ func (s *Session) Update(req *clientpb.RegisterSession) { s.Expression = req.RegisterData.Timer.Expression s.Jitter = req.RegisterData.Timer.Jitter } - s.SessionContext.Update(req) - + s.Modules = append([]string(nil), req.RegisterData.Module...) + s.Addons = cloneSessionAddons(req.RegisterData.Addons) // SecureManager现在使用固定的100次交互计数，不需要更新间隔 + shouldPublishInit := false if req.RegisterData.Sysinfo != nil { - if !s.Initialized { - s.Publish(consts.CtrlSessionInit, fmt.Sprintf("session %s init", s.ID), true, true) - } - s.UpdateSysInfo(req.RegisterData.Sysinfo) + shouldPublishInit = !s.Initialized + s.updateSysInfoLocked(req.RegisterData.Sysinfo) + } + s.stateMu.Unlock() + if shouldPublishInit { + s.Publish(consts.CtrlSessionInit, fmt.Sprintf("session %s init", s.ID), true, true) } err := s.Save() @@ -727,6 +919,18 @@ func (s *Session) UpdateSysInfo(info *implantpb.SysInfo) { if info == nil { return } + s.stateMu.Lock() + defer s.stateMu.Unlock() + s.updateSysInfoLocked(info) +} + +func (s *Session) SetWorkDir(workDir string) { + s.stateMu.Lock() + s.WorkDir = workDir + s.stateMu.Unlock() +} + +func (s *Session) updateSysInfoLocked(info *implantpb.SysInfo) { s.Initialized = true osInfo := mergeSysInfoOS(s.Os, info.GetOs()) processInfo := mergeSysInfoProcess(s.Process, info.GetProcess()) @@ -887,15 +1091,20 @@ func hasSignatureMetadata(value *implantpb.Process) bool { } func (s *Session) FillSysInfo() { - artifact, err := sessionDBGetArtifact(s.Name) + s.stateMu.RLock() + name := s.Name + s.stateMu.RUnlock() + artifact, err := sessionDBGetArtifact(name) if err != nil { - logs.Log.Errorf("failed to find atrtifact %s: %s", s.Name, err) + logs.Log.Errorf("failed to find atrtifact %s: %s", name, err) return } + s.stateMu.Lock() s.Os = &implantpb.Os{ Name: artifact.Os, Arch: artifact.Arch, } + s.stateMu.Unlock() } func (s *Session) Publish(Op string, msg string, notify bool, important bool) { @@ -932,17 +1141,6 @@ func (s *Session) Request(msg *clientpb.SpiteRequest, stream grpc.ServerStream) return stream.SendMsg(msg) } -func (s *Session) RequestAndWait(msg *clientpb.SpiteRequest, stream grpc.ServerStream, timeout time.Duration) (*implantpb.Spite, error) { - ch := make(chan *implantpb.Spite, 16) - s.StoreResp(msg.Task.TaskId, ch) - err := s.Request(msg, stream) - if err != nil { - return nil, err - } - resp := <-ch - return resp, nil -} - // RequestWithStream - 'async' means that the response is not returned immediately, but is returned through the channel 'ch func (s *Session) RequestWithStream(msg *clientpb.SpiteRequest, stream grpc.ServerStream, timeout time.Duration) (*SpiteStreamWriter, chan *implantpb.Spite, error) { respCh := make(chan *implantpb.Spite, 16) @@ -1045,15 +1243,6 @@ func (s *Session) DeleteResp(taskId uint32) { } } -// UpdateKeyPair 更新KeyPair并同步到SecureManager -func (s *Session) UpdateKeyPair(keyPair *clientpb.KeyPair) { - s.SessionContext.KeyPair = keyPair - // 更新SecureManager中的KeyPair引用 - if s.SecureManager != nil { - s.SecureManager.UpdateKeyPair(keyPair) - } -} - // SetKeepalive updates the keepalive state. Returns the previous state. func (s *Session) SetKeepalive(enabled bool) bool { s.keepaliveMu.Lock() @@ -1223,14 +1412,21 @@ func (s *Session) initializeSecureManager(req *clientpb.RegisterSession) error { } func (s *Session) UpdatePublicKey(key string) { - s.UpdateKeyPairFieldsAndPushCtrl(key, "") + s.UpdateKeyPair(key, "") } func (s *Session) UpdatePrivateKey(key string) { - s.UpdateKeyPairFieldsAndPushCtrl("", key) + s.UpdateKeyPair("", key) } -func (s *Session) UpdateKeyPairFieldsAndPushCtrl(publicKey string, privateKey string) { +// UpdateKeyPair merges the given non-empty public/private key fields into the +// session's current key pair, syncs the SecureManager, and pushes the update to +// the listener; an empty field leaves the existing value unchanged. The merge +// and store run under a single write lock so concurrent field updates (e.g. a +// re-register updating the public key while a key rotation updates both) cannot +// read the same base snapshot and clobber each other. +func (s *Session) UpdateKeyPair(publicKey string, privateKey string) { + s.stateMu.Lock() next := &clientpb.KeyPair{} if s.KeyPair != nil { next.PublicKey = s.KeyPair.PublicKey @@ -1242,7 +1438,16 @@ func (s *Session) UpdateKeyPairFieldsAndPushCtrl(publicKey string, privateKey st if privateKey != "" { next.PrivateKey = privateKey } - s.UpdateKeyPair(next) + // next is freshly allocated here and shared with no caller, so no defensive + // clone is needed before storing it. + s.SessionContext.KeyPair = next + if s.SecureManager != nil { + // Keep the session and manager updates in the same critical section so + // concurrent callers cannot publish an older manager value after a newer + // session value has already been stored. + s.SecureManager.UpdateKeyPair(next) + } + s.stateMu.Unlock() s.PushCtrl() } diff --git a/server/internal/core/session_lifecycle_test.go b/server/internal/core/session_lifecycle_test.go index 940b38b9a..71c84382c 100644 --- a/server/internal/core/session_lifecycle_test.go +++ b/server/internal/core/session_lifecycle_test.go @@ -25,12 +25,11 @@ func TestLifecycle_TickerMarksDead(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -88,12 +87,11 @@ func TestLifecycle_TickerKeepsAlive(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -167,12 +165,11 @@ func TestLifecycle_RegisterCheckinDeadReborn(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*4), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -279,12 +276,11 @@ func TestLifecycle_ConcurrentCheckinAndTicker(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*10), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -363,12 +359,11 @@ func TestLifecycle_MixedAliveDeadSessions(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*4), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -466,12 +461,11 @@ func TestEdge_RegisterLongSilenceThenReborn(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*4), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -594,12 +588,11 @@ func TestEdge_ConcurrentRemove(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*10), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -786,12 +779,11 @@ func TestEdge_RapidFlapping(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*20), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -936,12 +928,11 @@ func TestEdge_MultipleDeathsInSameTicker(t *testing.T) { broker := &eventBroker{ stop: make(chan struct{}), publish: make(chan Event, eventBufSize*10), - subscribe: make(chan chan Event, eventBufSize), + subscribe: make(chan eventSubscription), unsubscribe: make(chan chan Event, eventBufSize), send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker diff --git a/server/internal/core/session_state_accessors_test.go b/server/internal/core/session_state_accessors_test.go new file mode 100644 index 000000000..283f96543 --- /dev/null +++ b/server/internal/core/session_state_accessors_test.go @@ -0,0 +1,80 @@ +package core + +import ( + "fmt" + "sync" + "testing" + + "github.com/chainreactors/IoM-go/client" + implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" +) + +func TestSessionAddonAccessorsCloneDeduplicateAndRefresh(t *testing.T) { + sess := &Session{} + seatbelt := &implantpb.Addon{Name: "seatbelt", Type: "bof", Depend: "exec"} + sess.ReplaceAddons([]*implantpb.Addon{nil, seatbelt, {Name: "seatbelt", Type: "assembly"}, {Name: ""}}) + seatbelt.Type = "mutated-after-replace" + + sess.MergeAddons([]*implantpb.Addon{ + {Name: "seatbelt", Type: "assembly", Depend: "execute_dll"}, + {Name: "sharpview", Type: "assembly", Depend: "exec"}, + }) + addons := sess.AddonsSnapshot() + if len(addons) != 2 { + t.Fatalf("addon count = %d, want 2: %#v", len(addons), addons) + } + if addons[0].GetName() != "seatbelt" || addons[0].GetType() != "assembly" || addons[0].GetDepend() != "execute_dll" { + t.Fatalf("seatbelt addon = %#v, want refreshed clone", addons[0]) + } + if !sess.HasAddon("sharpview") || sess.HasAddon("missing") { + t.Fatalf("HasAddon results unexpected: addons=%#v", addons) + } + addons[0].Type = "mutated-snapshot" + if got := sess.AddonsSnapshot()[0].GetType(); got != "assembly" { + t.Fatalf("stored addon type = %q after snapshot mutation, want assembly", got) + } +} + +func TestSessionTimerAndRoutingSnapshotsDoNotTear(t *testing.T) { + sess := &Session{} + const iterations = 2000 + var writers sync.WaitGroup + writers.Add(1) + go func() { + defer writers.Done() + for i := 0; i < iterations; i++ { + suffix := i % 2 + sess.SetTimer(fmt.Sprintf("expr-%d", suffix), float64(suffix)) + sess.stateMu.Lock() + sess.Target = fmt.Sprintf("target-%d", suffix) + sess.ListenerID = fmt.Sprintf("listener-%d", suffix) + sess.PipelineID = fmt.Sprintf("pipeline-%d", suffix) + sess.stateMu.Unlock() + } + }() + + for i := 0; i < iterations; i++ { + expression, jitter := sess.TimerSnapshot() + if expression != "" && expression != fmt.Sprintf("expr-%d", int(jitter)) { + t.Fatalf("torn timer snapshot = %q/%v", expression, jitter) + } + target, listenerID, pipelineID := sess.ConnectionSnapshot() + if listenerID != "" && (target[len(target)-1:] != listenerID[len(listenerID)-1:] || listenerID[len(listenerID)-1:] != pipelineID[len(pipelineID)-1:]) { + t.Fatalf("torn connection snapshot = %q/%q/%q", target, listenerID, pipelineID) + } + } + writers.Wait() +} + +func TestSessionTimerAccessorsInitializeMissingSessionInfo(t *testing.T) { + for _, sess := range []*Session{ + {}, + {SessionContext: &client.SessionContext{}}, + } { + sess.SetTimer("*/5 * * * *", 0.25) + expression, jitter := sess.TimerSnapshot() + if expression != "*/5 * * * *" || jitter != 0.25 { + t.Fatalf("timer snapshot = %q/%v, want initialized values", expression, jitter) + } + } +} diff --git a/server/internal/core/session_test.go b/server/internal/core/session_test.go index 96c6c5737..d0a42836f 100644 --- a/server/internal/core/session_test.go +++ b/server/internal/core/session_test.go @@ -3,6 +3,7 @@ package core import ( "context" "errors" + "strings" "sync" "sync/atomic" "testing" @@ -466,6 +467,128 @@ func TestSessionRequestWithStreamWriterReturnsErrorAfterSendFailure(t *testing.T } } +func TestSessionSysInfoSnapshotConcurrent(t *testing.T) { + sess := newTestSession("sysinfo-snapshot") + sess.UpdateSysInfo(&implantpb.SysInfo{ + Os: &implantpb.Os{Name: "linux", Arch: "amd64", Username: "user-a"}, + Process: &implantpb.Process{Name: "agent-a", Arch: "amd64"}, + IsPrivilege: true, + Filepath: "/tmp/agent-a", + Workdir: "/tmp/a", + }) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + for i := 0; i < 500; i++ { + suffix := "a" + isPrivilege := true + if i%2 == 0 { + suffix = "b" + isPrivilege = false + } + sess.UpdateSysInfo(&implantpb.SysInfo{ + Os: &implantpb.Os{Name: "linux", Arch: "amd64", Username: "user-" + suffix}, + Process: &implantpb.Process{Name: "agent-" + suffix, Arch: "amd64"}, + IsPrivilege: isPrivilege, + Filepath: "/tmp/agent-" + suffix, + Workdir: "/tmp/" + suffix, + }) + } + }() + go func() { + defer wg.Done() + for i := 0; i < 500; i++ { + pb := sess.ToProtobufLite() + if pb.Os == nil || pb.Process == nil { + t.Errorf("ToProtobufLite returned incomplete sysinfo: %#v", pb) + return + } + suffix := "a" + wantPrivilege := true + if pb.Filepath == "/tmp/agent-b" { + suffix = "b" + wantPrivilege = false + } + if pb.Filepath != "/tmp/agent-"+suffix || pb.Workdir != "/tmp/"+suffix || + pb.Os.Username != "user-"+suffix || pb.Process.Name != "agent-"+suffix || + pb.IsPrivilege != wantPrivilege { + t.Errorf("ToProtobufLite mixed sysinfo generations: %#v", pb) + return + } + } + }() + wg.Wait() +} + +func TestSessionStateSnapshotConcurrentWithUpdate(t *testing.T) { + cleanup := installTestDBMocks() + defer cleanup() + + sess := newTestSession("state-snapshot") + sess.UpdateSysInfo(&implantpb.SysInfo{}) + register := func(suffix string) *clientpb.RegisterSession { + return &clientpb.RegisterSession{ + PipelineId: "pipeline-" + suffix, + ListenerId: "listener-" + suffix, + RegisterData: &implantpb.Register{ + Name: "agent-" + suffix, + Proxy: "proxy-" + suffix, + Module: []string{"module-" + suffix}, + Addons: []*implantpb.Addon{{Name: "addon-" + suffix}}, + Timer: &implantpb.Timer{Expression: "*/5 * * * *", Jitter: 0.1}, + Sysinfo: &implantpb.SysInfo{ + Os: &implantpb.Os{Name: "linux", Arch: "amd64", Username: "user-" + suffix}, + Process: &implantpb.Process{Name: "process-" + suffix, Arch: "amd64"}, + Filepath: "/tmp/agent-" + suffix, + Workdir: "/tmp/" + suffix, + }, + }, + } + } + + sess.Update(register("a")) + const iterations = 500 + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + for i := 0; i < iterations; i++ { + suffix := "a" + if i%2 == 0 { + suffix = "b" + } + sess.Update(register(suffix)) + } + }() + go func() { + defer wg.Done() + <-start + for i := 0; i < iterations; i++ { + pb := sess.ToProtobufLite() + suffix := strings.TrimPrefix(pb.Name, "agent-") + if suffix != "a" && suffix != "b" { + t.Errorf("unexpected session generation: %#v", pb) + return + } + if pb.PipelineId != "pipeline-"+suffix || pb.ListenerId != "listener-"+suffix || + pb.Proxy != "proxy-"+suffix || len(pb.Modules) != 1 || pb.Modules[0] != "module-"+suffix || + pb.Os == nil || pb.Os.Username != "user-"+suffix || pb.Process == nil || pb.Process.Name != "process-"+suffix { + t.Errorf("session snapshot mixed generations: %#v", pb) + return + } + _ = sess.ToProtobuf() + _ = sess.ToModel() + } + }() + + close(start) + wg.Wait() +} + // ---------- Task management ---------- func TestSession_NewTask_IncrementsSeq(t *testing.T) { diff --git a/server/internal/parser/malefic/parser.go b/server/internal/parser/malefic/parser.go index 79c8b2732..c9c756390 100644 --- a/server/internal/parser/malefic/parser.go +++ b/server/internal/parser/malefic/parser.go @@ -3,6 +3,7 @@ package malefic import ( "bytes" "encoding/binary" + "errors" "fmt" "github.com/chainreactors/IoM-go/consts" "github.com/chainreactors/IoM-go/proto/client/clientpb" @@ -24,6 +25,12 @@ const ( HeaderLength = 9 DefaultStartDelimiter = 0xd1 DefaultEndDelimiter = 0xd2 + maxDecodedPayloadSize = uint64(2 << 30) +) + +var ( + ErrDecodedPayloadTooLarge = errors.New("decoded payload exceeds 2 GiB limit") + ErrInvalidCompressedPayload = errors.New("invalid Snappy payload") ) func NewMaleficParser() *MaleficParser { @@ -34,11 +41,11 @@ func NewMaleficParser() *MaleficParser { } type MaleficParser struct { - StartDelimiter byte - EndDelimiter byte - MaxPacketLength uint32 - keyPair *clientpb.KeyPair // Age 密钥对，用于加解密 - privateKeys []string + StartDelimiter byte + EndDelimiter byte + MaxPacketLength uint32 + keyPair *clientpb.KeyPair // Age 密钥对，用于加解密 + privateKeys []string } // maxPacketLen returns the per-pipeline limit or falls back to global config. @@ -104,13 +111,21 @@ func (parser *MaleficParser) readHeader(conn io.ReadWriteCloser) (uint32, uint32 } sessionId := ParseSid(header) length := binary.LittleEndian.Uint32(header[MsgSessionEnd:]) + framedLength := uint64(length) + 1 + if framedLength > uint64(^uint32(0)) { + return 0, 0, fmt.Errorf("packet length %d for session %x overflows framed uint32 length", length, sessionId) + } + maxLen := parser.maxPacketLen() - if maxLen > 0 && length > maxLen+consts.KB*16 { - logs.Log.Warnf("[parser] large packet from session %x: %d bytes (limit %d), accepting anyway", - sessionId, length, maxLen) + if maxLen > 0 { + warningThreshold := uint64(maxLen) + uint64(consts.KB)*16 + if uint64(length) > warningThreshold { + logs.Log.Warnf("[parser] large packet from session %x: %d bytes (chunk size %d), accepting anyway", + sessionId, length, maxLen) + } } - return sessionId, length + 1, nil + return sessionId, uint32(framedLength), nil } func (parser *MaleficParser) ReadHeader(conn io.ReadWriteCloser) (uint32, uint32, error) { @@ -150,9 +165,20 @@ func (parser *MaleficParser) Parse(buf []byte) (*implantpb.Spites, error) { } } - buf, err := compress.Decompress(buf) + decodedLength, err := compress.DecodedLen(buf) + if err != nil { + return nil, fmt.Errorf("%w: read decoded length: %v", ErrInvalidCompressedPayload, err) + } + if uint64(decodedLength) > maxDecodedPayloadSize { + return nil, fmt.Errorf("%w: %d bytes", ErrDecodedPayloadTooLarge, decodedLength) + } + if err := compress.ValidateSnappyBlock(buf, decodedLength); err != nil { + return nil, fmt.Errorf("%w: validate block: %v", ErrInvalidCompressedPayload, err) + } + + buf, err = compress.Decompress(buf) if err != nil { - logs.Log.Debugf("trying plaintext: %v", err) + return nil, fmt.Errorf("%w: decompress: %v", ErrInvalidCompressedPayload, err) } spites := &implantpb.Spites{} diff --git a/server/internal/parser/malefic/parser_test.go b/server/internal/parser/malefic/parser_test.go index c98e24151..5daba1390 100644 --- a/server/internal/parser/malefic/parser_test.go +++ b/server/internal/parser/malefic/parser_test.go @@ -79,15 +79,30 @@ func TestMaleficParser_ReadHeader_PacketTooLarge(t *testing.T) { buf := &rwcBuf{bytes.NewBuffer(header)} sid, length, err := p.ReadHeader(buf) - // Large packets are now accepted with a warning, not rejected if err != nil { - t.Fatalf("expected large packet to be accepted, got error: %v", err) + t.Fatalf("expected packet beyond the chunking threshold to remain compatible: %v", err) } if sid != 1 { - t.Fatalf("expected session_id=1, got %d", sid) + t.Fatalf("ReadHeader SID = %d, want 1", sid) } - if length == 0 { - t.Fatal("expected non-zero length") + if length != hugeLen+1 { + t.Fatalf("ReadHeader length = %d, want %d", length, hugeLen+1) + } +} + +func TestMaleficParser_ReadHeader_CompatibilityMarginBoundary(t *testing.T) { + const maxPacketLength = uint32(1024) + p := NewMaleficParser() + p.MaxPacketLength = maxPacketLength + declaredLength := maxPacketLength + uint32(consts.KB)*16 + header := buildHeader(DefaultStartDelimiter, 1, declaredLength) + + _, length, err := p.ReadHeader(&rwcBuf{bytes.NewBuffer(header)}) + if err != nil { + t.Fatalf("expected packet at the compatibility boundary to be accepted: %v", err) + } + if length != declaredLength+1 { + t.Fatalf("ReadHeader length = %d, want %d", length, declaredLength+1) } } @@ -149,28 +164,48 @@ func TestMaleficParser_Parse_InvalidEndDelimiter(t *testing.T) { func TestMaleficParser_Parse_EmptyPayload(t *testing.T) { p := NewMaleficParser() - // Single byte: just the end delimiter. After stripping it, protobuf - // unmarshal receives an empty buffer. payload := []byte{DefaultEndDelimiter} - spites, err := p.Parse(payload) - // Empty protobuf unmarshal on Spites{} might succeed (empty message is valid) - // or fail depending on implementation. Either way, no panic. - if err != nil { - t.Logf("Parse empty payload returned error (expected): %v", err) - } else if spites != nil && len(spites.Spites) != 0 { - t.Logf("Parse empty payload returned spites with %d items", len(spites.Spites)) + if _, err := p.Parse(payload); !errors.Is(err, ErrInvalidCompressedPayload) { + t.Fatalf("Parse error = %v, want ErrInvalidCompressedPayload", err) + } +} + +func TestMaleficParser_Parse_RejectsDecodedPayloadAboveLimit(t *testing.T) { + p := NewMaleficParser() + var header [binary.MaxVarintLen64]byte + n := binary.PutUvarint(header[:], maxDecodedPayloadSize+1) + payload := append(append([]byte(nil), header[:n]...), DefaultEndDelimiter) + + _, err := p.Parse(payload) + if !errors.Is(err, ErrDecodedPayloadTooLarge) { + t.Fatalf("Parse error = %v, want ErrDecodedPayloadTooLarge", err) + } +} + +func TestMaleficParser_Parse_RejectsHeaderOnlyPayloadAtLimit(t *testing.T) { + p := NewMaleficParser() + var header [binary.MaxVarintLen64]byte + n := binary.PutUvarint(header[:], maxDecodedPayloadSize) + payload := append(append([]byte(nil), header[:n]...), DefaultEndDelimiter) + + if _, err := p.Parse(payload); !errors.Is(err, ErrInvalidCompressedPayload) { + t.Fatalf("Parse error = %v, want ErrInvalidCompressedPayload", err) } } func TestMaleficParser_Parse_ZeroLengthSlice(t *testing.T) { p := NewMaleficParser() - defer func() { - if r := recover(); r != nil { - t.Logf("BUG: Parse panics on zero-length slice: %v", r) - } - }() - // This will index buf[len(buf)-1] which is buf[-1] -> panic - _, _ = p.Parse([]byte{}) + if _, err := p.Parse(nil); err == nil { + t.Fatal("expected zero-length payload to be rejected") + } +} + +func TestMaleficParser_Parse_RejectsMalformedSnappyPayload(t *testing.T) { + p := NewMaleficParser() + payload := []byte{0xff, 0xfe, 0xfd, DefaultEndDelimiter} + if _, err := p.Parse(payload); !errors.Is(err, ErrInvalidCompressedPayload) { + t.Fatalf("Parse error = %v, want ErrInvalidCompressedPayload", err) + } } func TestMaleficParser_MarshalParse_RoundTrip(t *testing.T) { @@ -244,8 +279,6 @@ func TestParseSid(t *testing.T) { } } -// TestParseSid_ShortData documents a real bug: ParseSid does not bounds-check -// the input slice. Data shorter than 5 bytes causes a panic. func TestParseSid_ShortData(t *testing.T) { shortInputs := []struct { name string @@ -259,15 +292,9 @@ func TestParseSid_ShortData(t *testing.T) { for _, tc := range shortInputs { t.Run(tc.name, func(t *testing.T) { - defer func() { - if r := recover(); r != nil { - t.Logf("BUG CONFIRMED: ParseSid panics on %s input (len=%d): %v", - tc.name, len(tc.data), r) - } - }() - _ = ParseSid(tc.data) - // If we reach here without panic, the bug may have been fixed - // with a bounds check. That would be the correct behavior. + if got := ParseSid(tc.data); got != 0 { + t.Fatalf("ParseSid() = %d, want 0", got) + } }) } } diff --git a/server/internal/parser/malefic/untrusted_input_fuzz_test.go b/server/internal/parser/malefic/untrusted_input_fuzz_test.go new file mode 100644 index 000000000..1b62f6d46 --- /dev/null +++ b/server/internal/parser/malefic/untrusted_input_fuzz_test.go @@ -0,0 +1,104 @@ +package malefic + +import ( + "bytes" + "encoding/binary" + "errors" + "testing" + + "github.com/chainreactors/IoM-go/types" + "github.com/golang/snappy" +) + +const fuzzMaxPayloadLength = 64 << 10 + +func FuzzMaleficParserReadHeader(f *testing.F) { + seeds := [][]byte{ + nil, + {DefaultStartDelimiter}, + buildHeader(DefaultStartDelimiter, 1, 0), + buildHeader(DefaultStartDelimiter, 0x10203040, 1024), + buildHeader(DefaultStartDelimiter, 0x10203040, ^uint32(0)-1), + buildHeader(DefaultStartDelimiter, 0x10203040, ^uint32(0)), + buildHeader(0xff, 1, 16), + } + for _, seed := range seeds { + f.Add(seed) + } + + f.Fuzz(func(t *testing.T, data []byte) { + // ReadHeader consumes exactly nine bytes. Capping the corpus input avoids + // spending fuzz resources on irrelevant trailing data. + if len(data) > HeaderLength+16 { + return + } + + p := NewMaleficParser() + p.MaxPacketLength = 1024 + sid, length, err := p.ReadHeader(&rwcBuf{Buffer: bytes.NewBuffer(data)}) + + if len(data) < HeaderLength { + if err == nil { + t.Fatal("truncated header was accepted") + } + return + } + if data[MsgStart] != DefaultStartDelimiter { + if !errors.Is(err, types.ErrInvalidStart) { + t.Fatalf("invalid delimiter error = %v, want ErrInvalidStart", err) + } + return + } + declaredLength := binary.LittleEndian.Uint32(data[MsgSessionEnd:HeaderLength]) + if declaredLength == ^uint32(0) { + if err == nil { + t.Fatal("overflowing framed length was accepted") + } + return + } + if err != nil { + t.Fatalf("valid header returned error: %v", err) + } + + wantSID := binary.LittleEndian.Uint32(data[MsgSessionStart:MsgSessionEnd]) + wantLength := declaredLength + 1 + if sid != wantSID || length != wantLength { + t.Fatalf("ReadHeader() = (%d, %d), want (%d, %d)", sid, length, wantSID, wantLength) + } + }) +} + +func FuzzMaleficParserParse(f *testing.F) { + p := NewMaleficParser() + f.Add([]byte(nil)) + f.Add([]byte{DefaultEndDelimiter}) + f.Add([]byte{0x00, DefaultEndDelimiter}) + f.Add([]byte{0xff, 0xff, DefaultEndDelimiter}) + + f.Fuzz(func(t *testing.T, payload []byte) { + if len(payload) > fuzzMaxPayloadLength { + return + } + + // A tiny Snappy frame can declare a very large decoded size. Check the + // metadata before calling production Parse so this fuzz target cannot + // exhaust the host while still exercising malformed and bounded frames. + if len(payload) > 0 && payload[len(payload)-1] == DefaultEndDelimiter { + decodedLength, err := snappy.DecodedLen(payload[:len(payload)-1]) + if err == nil { + if uint64(decodedLength) > maxDecodedPayloadSize { + _, parseErr := p.Parse(payload) + if !errors.Is(parseErr, ErrDecodedPayloadTooLarge) { + t.Fatalf("oversized decoded payload error = %v, want ErrDecodedPayloadTooLarge", parseErr) + } + return + } + if decodedLength > fuzzMaxPayloadLength { + return + } + } + } + + _, _ = p.Parse(payload) + }) +} diff --git a/server/internal/parser/malefic/untrusted_input_test.go b/server/internal/parser/malefic/untrusted_input_test.go new file mode 100644 index 000000000..011dd10f0 --- /dev/null +++ b/server/internal/parser/malefic/untrusted_input_test.go @@ -0,0 +1,21 @@ +package malefic + +import ( + "bytes" + "strings" + "testing" +) + +func TestMaleficParserReadHeader_RejectsMaxDeclaredLength(t *testing.T) { + p := NewMaleficParser() + p.MaxPacketLength = 1024 + header := buildHeader(DefaultStartDelimiter, 0x10203040, ^uint32(0)) + + _, _, err := p.ReadHeader(&rwcBuf{Buffer: bytes.NewBuffer(header)}) + if err == nil { + t.Fatal("ReadHeader accepted a declared body length of math.MaxUint32; want a boundary error before uint32 wraparound") + } + if !strings.Contains(err.Error(), "overflows framed uint32 length") { + t.Fatalf("ReadHeader error = %q, want framed length overflow context", err) + } +} diff --git a/server/internal/parser/parser.go b/server/internal/parser/parser.go index 91268166b..3261090c3 100644 --- a/server/internal/parser/parser.go +++ b/server/internal/parser/parser.go @@ -57,6 +57,10 @@ type MessageParser struct { PacketParser } +const maxPayloadReadSize = uint32(512<<20) + 1 + +var ErrPayloadTooLarge = errors.New("encoded payload exceeds 512 MiB limit") + // WithSecure 为 MessageParser 添加安全支持 func (mp *MessageParser) WithSecure(keyPair *clientpb.KeyPair) { switch mp.Implant { @@ -80,8 +84,7 @@ func (mp *MessageParser) WithMaxPacketLength(n uint32) { } func (parser *MessageParser) ReadMessage(conn io.ReadWriteCloser, length uint32) (*implantpb.Spites, error) { - buf := make([]byte, length) - _, err := io.ReadFull(conn, buf) + buf, err := readPayload(conn, length) if err != nil { return nil, err } @@ -94,8 +97,7 @@ func (parser *MessageParser) ReadPacket(conn io.ReadWriteCloser) (uint32, *impla return 0, nil, err } - buf := make([]byte, length) - _, err = io.ReadFull(conn, buf) + buf, err := readPayload(conn, length) if err != nil { return 0, nil, err } @@ -104,6 +106,20 @@ func (parser *MessageParser) ReadPacket(conn io.ReadWriteCloser) (uint32, *impla return sessionId, msg, err } +func readPayload(conn io.Reader, length uint32) ([]byte, error) { + if length > maxPayloadReadSize { + return nil, fmt.Errorf("%w: %d bytes", ErrPayloadTooLarge, length) + } + if length == 0 { + return []byte{}, nil + } + payload := make([]byte, int(length)) + if _, err := io.ReadFull(conn, payload); err != nil { + return nil, err + } + return payload, nil +} + func (parser *MessageParser) WritePacket(conn net.Conn, msg *implantpb.Spites, sid uint32) error { bs, err := parser.Marshal(msg, sid) if err != nil { diff --git a/server/internal/parser/parser_test.go b/server/internal/parser/parser_test.go index 3bc8549f5..79c44d81f 100644 --- a/server/internal/parser/parser_test.go +++ b/server/internal/parser/parser_test.go @@ -160,10 +160,7 @@ func TestReadPacket_RoundTrip(t *testing.T) { } } -// TestReadPacket_DiscardParseError documents the bug on line 104 of parser.go: -// ReadPacket returns nil error even when Parse fails because the return -// statement is `return sessionId, msg, nil` instead of `return sessionId, msg, err`. -func TestReadPacket_DiscardParseError(t *testing.T) { +func TestReadPacket_ReturnsParseError(t *testing.T) { parser, err := NewParser("malefic") if err != nil { t.Fatalf("failed to create parser: %v", err) @@ -176,7 +173,7 @@ func TestReadPacket_DiscardParseError(t *testing.T) { // but we make the body garbage so protobuf unmarshal fails. sid := uint32(1) payloadBody := []byte{0xFF, 0xFE, 0xFD} // garbage protobuf data - payload := append(payloadBody, 0xd2) // valid end delimiter + payload := append(payloadBody, 0xd2) // valid end delimiter header := make([]byte, 9) header[0] = 0xd1 // start delimiter @@ -204,21 +201,15 @@ func TestReadPacket_DiscardParseError(t *testing.T) { gotSid, msg, err := parser.ReadPacket(serverConn) - // BUG: err is nil even though Parse should have failed on garbage protobuf data. - // The msg will be nil (unmarshal failure) but the error is swallowed. - if err != nil { - // If this branch is reached, the bug has been fixed. - t.Logf("Bug appears fixed: ReadPacket now returns parse error: %v", err) - return + if err == nil { + t.Fatal("ReadPacket discarded the parse error") } - - // Document the bug: error is nil but msg is also nil. - if msg == nil { - t.Log("BUG CONFIRMED: ReadPacket returned nil message AND nil error. " + - "Line 104 discards the parse error (returns nil instead of err).") + if msg != nil { + t.Fatalf("ReadPacket message = %#v, want nil after parse failure", msg) + } + if gotSid != sid { + t.Fatalf("ReadPacket session ID = %d, want %d", gotSid, sid) } - - _ = gotSid } // --- WritePacket tests --- @@ -333,3 +324,15 @@ func TestReadMessage_ValidPayload(t *testing.T) { t.Fatalf("unexpected result: %+v", result) } } + +func TestReadMessage_RejectsPayloadAboveWireLimit(t *testing.T) { + parser, err := NewParser("malefic") + if err != nil { + t.Fatalf("failed to create parser: %v", err) + } + + _, err = parser.ReadMessage(&rwcBuf{bytes.NewBuffer(nil)}, maxPayloadReadSize+1) + if !errors.Is(err, ErrPayloadTooLarge) { + t.Fatalf("ReadMessage error = %v, want ErrPayloadTooLarge", err) + } +} diff --git a/server/internal/parser/untrusted_length_test.go b/server/internal/parser/untrusted_length_test.go new file mode 100644 index 000000000..1a912be9b --- /dev/null +++ b/server/internal/parser/untrusted_length_test.go @@ -0,0 +1,122 @@ +package parser + +import ( + "bytes" + "encoding/binary" + "errors" + "io" + "testing" + + "github.com/chainreactors/IoM-go/consts" + "github.com/chainreactors/IoM-go/proto/implant/implantpb" + "github.com/chainreactors/malice-network/server/internal/parser/malefic" +) + +type headerThenEOFConn struct { + header []byte + offset int + maxReadRequest int +} + +func (c *headerThenEOFConn) Read(p []byte) (int, error) { + if len(p) > c.maxReadRequest { + c.maxReadRequest = len(p) + } + if c.offset == len(c.header) { + return 0, io.EOF + } + n := copy(p, c.header[c.offset:]) + c.offset += n + return n, nil +} + +func (c *headerThenEOFConn) Write(p []byte) (int, error) { return len(p), nil } +func (c *headerThenEOFConn) Close() error { return nil } + +func TestReadPacket_DeclaredLengthBelowLimitReadsDeclaredPayload(t *testing.T) { + const configuredLimit = uint32(1024) + declaredLength := uint32(256 * 1024) + + header := make([]byte, malefic.HeaderLength) + header[malefic.MsgStart] = malefic.DefaultStartDelimiter + binary.LittleEndian.PutUint32(header[malefic.MsgSessionStart:malefic.MsgSessionEnd], 1) + binary.LittleEndian.PutUint32(header[malefic.MsgSessionEnd:], declaredLength) + + p, err := NewParser(consts.ImplantMalefic) + if err != nil { + t.Fatalf("NewParser returned an unexpected error: %v", err) + } + p.WithMaxPacketLength(configuredLimit) + conn := &headerThenEOFConn{header: header} + + _, _, err = p.ReadPacket(conn) + if !errors.Is(err, io.EOF) { + t.Fatalf("ReadPacket error = %v, want EOF after the header", err) + } + if conn.maxReadRequest != int(declaredLength+1) { + t.Fatalf("ReadPacket requested a %d-byte read buffer, want %d", conn.maxReadRequest, declaredLength+1) + } +} + +func TestReadPacket_RejectsDeclaredLengthAboveWireLimit(t *testing.T) { + declaredLength := uint32(512<<20) + 1 + header := make([]byte, malefic.HeaderLength) + header[malefic.MsgStart] = malefic.DefaultStartDelimiter + binary.LittleEndian.PutUint32(header[malefic.MsgSessionStart:malefic.MsgSessionEnd], 1) + binary.LittleEndian.PutUint32(header[malefic.MsgSessionEnd:], declaredLength) + + p, err := NewParser(consts.ImplantMalefic) + if err != nil { + t.Fatalf("NewParser returned an unexpected error: %v", err) + } + conn := &headerThenEOFConn{header: header} + + _, _, err = p.ReadPacket(conn) + if !errors.Is(err, ErrPayloadTooLarge) { + t.Fatalf("ReadPacket error = %v, want ErrPayloadTooLarge", err) + } + if conn.maxReadRequest != malefic.HeaderLength { + t.Fatalf("ReadPacket attempted a payload read of %d bytes after rejecting the header", conn.maxReadRequest) + } +} + +func TestReadPacket_AcceptsValidFrameLargerThanPacketLength(t *testing.T) { + const configuredPacketLength = uint32(1024) + p, err := NewParser(consts.ImplantMalefic) + if err != nil { + t.Fatalf("NewParser returned an unexpected error: %v", err) + } + p.WithMaxPacketLength(configuredPacketLength) + + stdout := make([]byte, 128*1024) + x := uint32(1) + for i := range stdout { + x ^= x << 13 + x ^= x >> 17 + x ^= x << 5 + stdout[i] = byte(x) + } + spites := &implantpb.Spites{Spites: []*implantpb.Spite{{ + TaskId: 1, + Body: &implantpb.Spite_ExecResponse{ExecResponse: &implantpb.ExecResponse{ + Stdout: stdout, + }}, + }}} + raw, err := p.Marshal(spites, 1) + if err != nil { + t.Fatalf("Marshal returned an unexpected error: %v", err) + } + declaredLength := binary.LittleEndian.Uint32(raw[malefic.MsgSessionEnd:malefic.HeaderLength]) + warningThreshold := configuredPacketLength + uint32(consts.KB)*16 + if declaredLength <= warningThreshold { + t.Fatalf("test frame length = %d, want larger than chunking threshold %d", declaredLength, warningThreshold) + } + + _, got, err := p.ReadPacket(&rwcBuf{Buffer: bytes.NewBuffer(raw)}) + if err != nil { + t.Fatalf("ReadPacket rejected a valid large frame: %v", err) + } + if !bytes.Equal(got.Spites[0].GetExecResponse().GetStdout(), stdout) { + t.Fatal("ReadPacket changed the large exec response payload") + } +} diff --git a/server/internal/stream/conn.go b/server/internal/stream/conn.go index 8310e78f9..e306c5f10 100644 --- a/server/internal/stream/conn.go +++ b/server/internal/stream/conn.go @@ -2,9 +2,12 @@ package cryptostream import ( "bytes" + "errors" + "fmt" "io" "net" "sync" + "time" ) func NewCryptoConn(conn net.Conn, cryptor Cryptor) *CryptoConn { @@ -78,6 +81,36 @@ func (sc *CryptoConn) Close() error { return sc.ReadWriteCloser.Close() } +func (sc *CryptoConn) SetDeadline(deadline time.Time) error { + if sc.Conn != nil { + return sc.Conn.SetDeadline(deadline) + } + if conn, ok := sc.ReadWriteCloser.(interface{ SetDeadline(time.Time) error }); ok { + return conn.SetDeadline(deadline) + } + return fmt.Errorf("set deadline: %w", errors.ErrUnsupported) +} + +func (sc *CryptoConn) SetReadDeadline(deadline time.Time) error { + if sc.Conn != nil { + return sc.Conn.SetReadDeadline(deadline) + } + if conn, ok := sc.ReadWriteCloser.(interface{ SetReadDeadline(time.Time) error }); ok { + return conn.SetReadDeadline(deadline) + } + return fmt.Errorf("set read deadline: %w", errors.ErrUnsupported) +} + +func (sc *CryptoConn) SetWriteDeadline(deadline time.Time) error { + if sc.Conn != nil { + return sc.Conn.SetWriteDeadline(deadline) + } + if conn, ok := sc.ReadWriteCloser.(interface{ SetWriteDeadline(time.Time) error }); ok { + return conn.SetWriteDeadline(deadline) + } + return fmt.Errorf("set write deadline: %w", errors.ErrUnsupported) +} + func (sc *CryptoConn) encrypt(data []byte) ([]byte, error) { reader := bytes.NewReader(data) writer := &bytes.Buffer{} diff --git a/server/internal/stream/conn_test.go b/server/internal/stream/conn_test.go index e430b593a..577fe8813 100644 --- a/server/internal/stream/conn_test.go +++ b/server/internal/stream/conn_test.go @@ -8,6 +8,7 @@ import ( "net" "sync" "testing" + "time" ) func TestCryptoConnConcurrentReadsSerializeOverflowBuffer(t *testing.T) { @@ -52,6 +53,17 @@ type simpleRWC struct { closed bool } +type deadlineRWC struct { + *simpleRWC + readDeadline time.Time + readErr error +} + +func (d *deadlineRWC) SetReadDeadline(deadline time.Time) error { + d.readDeadline = deadline + return d.readErr +} + func (s *simpleRWC) Read(p []byte) (int, error) { if s.closed { return 0, errors.New("closed") @@ -408,6 +420,41 @@ func TestCryptoRWC_WriteRead(t *testing.T) { } } +func TestCryptoRWCDeadlineMethodsReportUnsupported(t *testing.T) { + t.Parallel() + cryptor := NewXorEncryptor([]byte{0}, []byte{0}) + conn := NewCryptoRWC(&simpleRWC{buf: &bytes.Buffer{}}, cryptor) + deadline := time.Now().Add(time.Second) + if err := conn.SetDeadline(deadline); !errors.Is(err, errors.ErrUnsupported) { + t.Fatalf("SetDeadline error = %v, want errors.ErrUnsupported", err) + } + if err := conn.SetReadDeadline(deadline); !errors.Is(err, errors.ErrUnsupported) { + t.Fatalf("SetReadDeadline error = %v, want errors.ErrUnsupported", err) + } + if err := conn.SetWriteDeadline(deadline); !errors.Is(err, errors.ErrUnsupported) { + t.Fatalf("SetWriteDeadline error = %v, want errors.ErrUnsupported", err) + } +} + +func TestCryptoRWCReadDeadlineDelegatesAndReturnsError(t *testing.T) { + t.Parallel() + wantErr := errors.New("deadline failed") + rwc := &deadlineRWC{ + simpleRWC: &simpleRWC{buf: &bytes.Buffer{}}, + readErr: wantErr, + } + conn := NewCryptoRWC(rwc, NewXorEncryptor([]byte{0}, []byte{0})) + deadline := time.Now().Add(time.Second) + + err := conn.SetReadDeadline(deadline) + if !errors.Is(err, wantErr) { + t.Fatalf("SetReadDeadline error = %v, want %v", err, wantErr) + } + if !rwc.readDeadline.Equal(deadline) { + t.Fatalf("delegated deadline = %v, want %v", rwc.readDeadline, deadline) + } +} + // TestCryptoConn_RemoteAddr_WithConn verifies RemoteAddr returns the conn's remote addr. func TestCryptoConn_RemoteAddr_WithConn(t *testing.T) { t.Parallel() diff --git a/server/internal/stream/peekconn.go b/server/internal/stream/peekconn.go index d58905546..5c27426ac 100644 --- a/server/internal/stream/peekconn.go +++ b/server/internal/stream/peekconn.go @@ -1,15 +1,24 @@ package cryptostream import ( - "github.com/chainreactors/logs" - "github.com/chainreactors/malice-network/server/internal/parser" - "github.com/chainreactors/malice-network/server/internal/parser/malefic" + "errors" + "fmt" "io" "net" "sync" "time" + + "github.com/chainreactors/logs" + "github.com/chainreactors/malice-network/server/internal/parser" + "github.com/chainreactors/malice-network/server/internal/parser/malefic" ) +const initialPacketReadTimeout = 10 * time.Second + +type readDeadlineSetter interface { + SetReadDeadline(time.Time) error +} + type ReadWriteCloser struct { r io.Reader w io.Writer @@ -58,10 +67,26 @@ func PeekSid(conn *Conn) (uint32, error) { return malefic.ParseSid(data), nil } +// WrapPeekConn bounds the initial header read when the transport supports read +// deadlines. Callers using other stream types must provide their own timeout. func WrapPeekConn(conn io.ReadWriteCloser, cryptos []Cryptor, parserName string, maxPacketLength uint32) (*Conn, error) { + return wrapPeekConn(conn, cryptos, parserName, maxPacketLength, initialPacketReadTimeout) +} + +func wrapPeekConn(conn io.ReadWriteCloser, cryptos []Cryptor, parserName string, maxPacketLength uint32, readTimeout time.Duration) (*Conn, error) { bs := make([]byte, 9) - _, err := io.ReadFull(conn, bs) - if err != nil { + deadlineConn, hasDeadline := conn.(readDeadlineSetter) + if hasDeadline { + if err := deadlineConn.SetReadDeadline(time.Now().Add(readTimeout)); err != nil { + return nil, fmt.Errorf("set initial packet read deadline: %w", err) + } + } + _, readErr := io.ReadFull(conn, bs) + var clearErr error + if hasDeadline { + clearErr = deadlineConn.SetReadDeadline(time.Time{}) + } + if err := errors.Join(readErr, clearErr); err != nil { return nil, err } @@ -71,6 +96,7 @@ func WrapPeekConn(conn io.ReadWriteCloser, cryptos []Cryptor, parserName string, var c Cryptor var p *parser.MessageParser var i int + var err error for i, c = range cryptos { var de []byte de, err = Decrypt(c, bs) diff --git a/server/internal/stream/peekconn_deadline_test.go b/server/internal/stream/peekconn_deadline_test.go new file mode 100644 index 000000000..5eaad3383 --- /dev/null +++ b/server/internal/stream/peekconn_deadline_test.go @@ -0,0 +1,80 @@ +package cryptostream + +import ( + "encoding/binary" + "errors" + "net" + "testing" + "time" + + "github.com/chainreactors/IoM-go/consts" + "github.com/chainreactors/malice-network/server/internal/parser/malefic" +) + +func TestWrapPeekConnTimesOutIncompleteInitialPacket(t *testing.T) { + serverConn, clientConn := net.Pipe() + t.Cleanup(func() { + _ = serverConn.Close() + _ = clientConn.Close() + }) + cryptor, err := NewCryptor(consts.CryptorRAW, nil, nil) + if err != nil { + t.Fatalf("new raw cryptor: %v", err) + } + + started := time.Now() + _, err = wrapPeekConn(serverConn, []Cryptor{cryptor}, consts.ImplantMalefic, 0, 20*time.Millisecond) + var netErr net.Error + if !errors.As(err, &netErr) || !netErr.Timeout() { + t.Fatalf("WrapPeekConn error = %v, want network timeout", err) + } + if elapsed := time.Since(started); elapsed > time.Second { + t.Fatalf("initial packet timeout took %s, want under 1s", elapsed) + } +} + +func TestWrapPeekConnClearsInitialPacketDeadline(t *testing.T) { + serverConn, clientConn := net.Pipe() + t.Cleanup(func() { + _ = serverConn.Close() + _ = clientConn.Close() + }) + header := make([]byte, malefic.HeaderLength) + header[0] = malefic.DefaultStartDelimiter + binary.LittleEndian.PutUint32(header[1:5], 7) + binary.LittleEndian.PutUint32(header[5:9], 0) + go func() { _, _ = clientConn.Write(header) }() + + cryptor, err := NewCryptor(consts.CryptorRAW, nil, nil) + if err != nil { + t.Fatalf("new raw cryptor: %v", err) + } + conn, err := wrapPeekConn(serverConn, []Cryptor{cryptor}, consts.ImplantMalefic, 0, 20*time.Millisecond) + if err != nil { + t.Fatalf("WrapPeekConn failed: %v", err) + } + initial := make([]byte, len(header)) + if _, err := conn.Read(initial); err != nil { + t.Fatalf("drain initial packet: %v", err) + } + + time.Sleep(40 * time.Millisecond) + readResult := make(chan error, 1) + go func() { + buf := make([]byte, 1) + _, readErr := conn.Read(buf) + if readErr == nil && buf[0] != 0x42 { + readErr = errors.New("unexpected byte after initial packet") + } + readResult <- readErr + }() + go func() { _, _ = clientConn.Write([]byte{0x42}) }() + select { + case err := <-readResult: + if err != nil { + t.Fatalf("read after initial deadline: %v", err) + } + case <-time.After(time.Second): + t.Fatal("read stayed blocked after initial deadline should have been cleared") + } +} diff --git a/server/listener/cmux.go b/server/listener/cmux.go index d355e01e7..269ee36e2 100644 --- a/server/listener/cmux.go +++ b/server/listener/cmux.go @@ -6,11 +6,17 @@ import ( "fmt" "net" "net/http" + "time" "github.com/chainreactors/malice-network/server/internal/core" "github.com/soheilhy/cmux" ) +const ( + cmuxReadTimeout = 10 * time.Second + httpReadHeaderTimeout = 10 * time.Second +) + var serveCMux = func(m cmux.CMux) error { return m.Serve() } @@ -21,12 +27,18 @@ var serveHTTP = func(server *http.Server, ln net.Listener) error { func NewHTTPServer(handler http.Handler) *http.Server { return &http.Server{ - Handler: handler, + Handler: handler, + ReadHeaderTimeout: httpReadHeaderTimeout, } } func SpiltTLSListener(ln net.Listener) (cmux.CMux, net.Listener, net.Listener) { + return splitTLSListener(ln, cmuxReadTimeout) +} + +func splitTLSListener(ln net.Listener, readTimeout time.Duration) (cmux.CMux, net.Listener, net.Listener) { m := cmux.New(ln) + m.SetReadTimeout(readTimeout) tlsL := m.Match(cmux.TLS()) plainL := m.Match(cmux.Any()) return m, tlsL, plainL diff --git a/server/listener/cmux_server_test.go b/server/listener/cmux_server_test.go new file mode 100644 index 000000000..c9c0ae46e --- /dev/null +++ b/server/listener/cmux_server_test.go @@ -0,0 +1,73 @@ +package listener + +import ( + "errors" + "net" + "net/http" + "testing" + "time" +) + +func TestNewHTTPServerLimitsHeaderReadTime(t *testing.T) { + server := NewHTTPServer(http.NotFoundHandler()) + if server.ReadHeaderTimeout != httpReadHeaderTimeout { + t.Fatalf("ReadHeaderTimeout = %s, want %s", server.ReadHeaderTimeout, httpReadHeaderTimeout) + } + if server.ReadHeaderTimeout <= 0 { + t.Fatal("ReadHeaderTimeout must be positive") + } +} + +func TestSplitTLSListenerLimitsClassifierReadTime(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + mux, _, plain := splitTLSListener(listener, 20*time.Millisecond) + serveDone := make(chan error, 1) + go func() { serveDone <- mux.Serve() }() + t.Cleanup(func() { + _ = listener.Close() + mux.Close() + select { + case err := <-serveDone: + if err != nil && !errors.Is(err, net.ErrClosed) { + t.Errorf("cmux serve: %v", err) + } + case <-time.After(time.Second): + t.Error("cmux Serve did not stop") + } + }) + + conn, err := net.Dial("tcp", listener.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.Close() + if err := conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set client deadline: %v", err) + } + + accepted := make(chan net.Conn, 1) + acceptErr := make(chan error, 1) + go func() { + serverConn, err := plain.Accept() + if err != nil { + acceptErr <- err + return + } + accepted <- serverConn + }() + started := time.Now() + select { + case serverConn := <-accepted: + _ = serverConn.Close() + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("classifier timeout took %s", elapsed) + } + case err := <-acceptErr: + t.Fatalf("plain listener accept: %v", err) + case <-time.After(500 * time.Millisecond): + t.Fatal("idle connection remained stuck in cmux classification") + } +} diff --git a/server/listener/forward.go b/server/listener/forward.go index 0bf66c63c..e1bcc4bdd 100644 --- a/server/listener/forward.go +++ b/server/listener/forward.go @@ -7,6 +7,7 @@ import ( "io" "net" "sync" + "sync/atomic" "time" "github.com/chainreactors/IoM-go/proto/client/clientpb" @@ -22,6 +23,15 @@ import ( "google.golang.org/grpc/status" ) +const forwardStreamCloseTimeout = 5 * time.Second + +var ( + ErrForwardStreamClosing = errors.New("forward stream is closing") + ErrForwardStreamCloseTimeout = errors.New("forward stream close timed out") + ErrForwardStreamPoisoned = errors.New("forward stream has not drained") + ErrForwardStreamAlreadyAttached = errors.New("forward stream already has an active task stream") +) + type ForwardListener struct { lns *listener server *grpc.Server @@ -73,7 +83,7 @@ func NewForwardListener(cfg *configs.ListenerConfig) (*ForwardListener, error) { } return nil }, core.LogGuardedError("forward-listener:"+cfg.Name)) - Listener = lns + setCurrentListener(lns) logs.Log.Importantf("listener.forward - start name=%s address=%s", cfg.Name, ln.Addr().String()) return runtime, nil } @@ -128,8 +138,18 @@ func (s *forwardListenerService) TaskStream(stream forwardrpc.ForwardListener_Ta if err != nil { return err } - local := s.registry.get(listenerID, pipelineID) - local.attach(stream) + local, err := s.registry.open(listenerID, pipelineID) + if err != nil { + return status.Error(codes.Unavailable, err.Error()) + } + release, err := local.attach() + if err != nil { + if errors.Is(err, ErrForwardStreamAlreadyAttached) { + return status.Error(codes.AlreadyExists, err.Error()) + } + return status.Error(codes.Unavailable, err.Error()) + } + defer release() return local.serve(stream) } @@ -139,19 +159,28 @@ type forwardPipelineRPC struct { } func (c *forwardPipelineRPC) OpenForwardStream(ctx context.Context, pipeline core.Pipeline) (core.ForwardStream, error) { + if err := ctx.Err(); err != nil { + return nil, err + } listenerID := c.listenerID if pb := pipeline.ToProtobuf(); pb != nil && pb.ListenerId != "" { listenerID = pb.ListenerId } - return c.registry.get(listenerID, pipeline.ID()), nil + return c.registry.open(listenerID, pipeline.ID()) } func (c *forwardPipelineRPC) Register(ctx context.Context, in *clientpb.RegisterSession, _ ...grpc.CallOption) (*clientpb.Empty, error) { if in == nil { return nil, fmt.Errorf("register session is nil") } - stream := c.registry.get(in.ListenerId, in.PipelineId) - return &clientpb.Empty{}, stream.sendEvent(&clientpb.SpiteRequest{ + if err := ctx.Err(); err != nil { + return nil, err + } + stream, err := c.registry.open(in.ListenerId, in.PipelineId) + if err != nil { + return nil, err + } + return &clientpb.Empty{}, stream.sendEventContext(ctx, &clientpb.SpiteRequest{ ListenerId: in.ListenerId, Session: &clientpb.Session{ SessionId: in.SessionId, @@ -169,14 +198,20 @@ func (c *forwardPipelineRPC) Register(ctx context.Context, in *clientpb.Register } func (c *forwardPipelineRPC) Checkin(ctx context.Context, in *implantpb.Ping, _ ...grpc.CallOption) (*clientpb.Empty, error) { + if err := ctx.Err(); err != nil { + return nil, err + } sessionID, _ := metadataValue(ctx, "session_id") listenerID, _ := metadataValue(ctx, "listener_id") pipelineID, _ := metadataValue(ctx, "pipeline_id") if listenerID == "" { listenerID = c.listenerID } - stream := c.registry.get(listenerID, pipelineID) - return &clientpb.Empty{}, stream.sendEvent(&clientpb.SpiteRequest{ + stream, err := c.registry.open(listenerID, pipelineID) + if err != nil { + return nil, err + } + return &clientpb.Empty{}, stream.sendEventContext(ctx, &clientpb.SpiteRequest{ ListenerId: listenerID, Session: &clientpb.Session{ SessionId: sessionID, @@ -194,55 +229,153 @@ func (c *forwardPipelineRPC) GetArtifact(context.Context, *clientpb.Artifact, .. return nil, status.Error(codes.Unimplemented, "artifact fetch is not supported by forward listener transport") } -func (c *forwardPipelineRPC) removeStream(listenerID, pipelineID string) { - c.registry.remove(listenerID, pipelineID) +func (c *forwardPipelineRPC) retireStream(listenerID, pipelineID string) (func(), error) { + return c.registry.retire(listenerID, pipelineID) } type forwardStreamRegistry struct { - mu sync.Mutex - streams map[string]*forwardLocalStream + mu sync.Mutex + streams map[string]*forwardLocalStream + retiring map[string]int + poisoned map[string]*forwardLocalStream } func newForwardStreamRegistry() *forwardStreamRegistry { - return &forwardStreamRegistry{streams: make(map[string]*forwardLocalStream)} + return &forwardStreamRegistry{ + streams: make(map[string]*forwardLocalStream), + retiring: make(map[string]int), + poisoned: make(map[string]*forwardLocalStream), + } } -func (r *forwardStreamRegistry) get(listenerID, pipelineID string) *forwardLocalStream { +func (r *forwardStreamRegistry) open(listenerID, pipelineID string) (*forwardLocalStream, error) { key := core.PipelineRuntimeKey(listenerID, pipelineID) r.mu.Lock() defer r.mu.Unlock() if stream := r.streams[key]; stream != nil { - return stream + return stream, nil + } + if r.retiring[key] > 0 { + return nil, fmt.Errorf("%w: %s", ErrForwardStreamClosing, key) } - stream := &forwardLocalStream{ + if r.poisoned[key] != nil { + return nil, fmt.Errorf("%w: %s", ErrForwardStreamPoisoned, key) + } + stream := newForwardLocalStream(r, listenerID, pipelineID) + r.streams[key] = stream + return stream, nil +} + +func newForwardLocalStream(registry *forwardStreamRegistry, listenerID, pipelineID string) *forwardLocalStream { + return &forwardLocalStream{ + registry: registry, listenerID: listenerID, pipelineID: pipelineID, requests: make(chan *clientpb.SpiteRequest, 255), events: make(chan *clientpb.SpiteRequest, 255), done: make(chan struct{}), } - r.streams[key] = stream - return stream } -func (r *forwardStreamRegistry) remove(listenerID, pipelineID string) { +func (r *forwardStreamRegistry) retire(listenerID, pipelineID string) (func(), error) { + ctx, cancel := context.WithTimeout(context.Background(), forwardStreamCloseTimeout) + defer cancel() + return r.retireContext(ctx, listenerID, pipelineID) +} + +func (r *forwardStreamRegistry) retireContext(ctx context.Context, listenerID, pipelineID string) (func(), error) { key := core.PipelineRuntimeKey(listenerID, pipelineID) r.mu.Lock() + r.retiring[key]++ stream := r.streams[key] delete(r.streams, key) r.mu.Unlock() + + var once sync.Once + release := func() { + once.Do(func() { + r.mu.Lock() + defer r.mu.Unlock() + r.retiring[key]-- + if r.retiring[key] == 0 { + delete(r.retiring, key) + } + }) + } if stream != nil { - stream.close() + err := stream.closeContext(ctx) + r.poisonIfUndrained(key, stream, err) + return release, err + } + return release, nil +} + +func (r *forwardStreamRegistry) discard(stream *forwardLocalStream) error { + if r == nil || stream == nil { + return nil + } + key := core.PipelineRuntimeKey(stream.listenerID, stream.pipelineID) + r.mu.Lock() + if r.streams[key] != stream { + r.mu.Unlock() + return stream.close() + } + r.retiring[key]++ + delete(r.streams, key) + r.mu.Unlock() + + err := stream.close() + r.poisonIfUndrained(key, stream, err) + r.mu.Lock() + r.retiring[key]-- + if r.retiring[key] == 0 { + delete(r.retiring, key) + } + r.mu.Unlock() + return err +} + +func (r *forwardStreamRegistry) poisonIfUndrained(key string, stream *forwardLocalStream, closeErr error) { + if !errors.Is(closeErr, ErrForwardStreamCloseTimeout) || stream.drained == nil { + return } + r.mu.Lock() + r.poisoned[key] = stream + r.mu.Unlock() + go func() { + <-stream.drained + r.mu.Lock() + if r.poisoned[key] == stream { + delete(r.poisoned, key) + } + r.mu.Unlock() + }() } type forwardLocalStream struct { + registry *forwardStreamRegistry listenerID string pipelineID string requests chan *clientpb.SpiteRequest events chan *clientpb.SpiteRequest done chan struct{} closeOnce sync.Once + closed atomic.Bool + eventMu sync.Mutex + slotsOnce sync.Once + // eventSlots is a semaphore holding exactly one token per free slot in the + // events channel. Invariant: every entry sent into events must first acquire + // a slot (sendEventContext) and release it exactly once on consumption or on + // the closed-path early return. This lets senders wait for capacity outside + // eventMu (cancellable) and then commit the send while holding eventMu + // without ever blocking on a full channel. Any new events send path MUST + // keep this one-slot-per-in-flight-entry accounting or it will deadlock/leak. + eventSlots chan struct{} + remoteSend sync.WaitGroup + closeErr error + drained chan struct{} + attachMu sync.Mutex + attached bool } func (s *forwardLocalStream) Send(resp *clientpb.SpiteResponse) error { @@ -261,6 +394,13 @@ func (s *forwardLocalStream) Send(resp *clientpb.SpiteResponse) error { }) } +func (s *forwardLocalStream) CloseSend() error { + if s.registry != nil { + return s.registry.discard(s) + } + return s.close() +} + func (s *forwardLocalStream) Recv() (*clientpb.SpiteRequest, error) { select { case req, ok := <-s.requests: @@ -273,7 +413,26 @@ func (s *forwardLocalStream) Recv() (*clientpb.SpiteRequest, error) { } } -func (s *forwardLocalStream) attach(_ forwardrpc.ForwardListener_TaskStreamServer) {} +func (s *forwardLocalStream) attach() (func(), error) { + s.attachMu.Lock() + defer s.attachMu.Unlock() + if s.closed.Load() { + return nil, ErrForwardStreamClosing + } + if s.attached { + return nil, ErrForwardStreamAlreadyAttached + } + s.attached = true + + var once sync.Once + return func() { + once.Do(func() { + s.attachMu.Lock() + s.attached = false + s.attachMu.Unlock() + }) + }, nil +} func (s *forwardLocalStream) serve(stream forwardrpc.ForwardListener_TaskStreamServer) error { ctx, cancel := context.WithCancel(stream.Context()) @@ -297,13 +456,18 @@ func (s *forwardLocalStream) serve(stream forwardrpc.ForwardListener_TaskStreamS } }() go func() { + s.eventSlotPool() for { select { case event, ok := <-s.events: if !ok { return } - if err := stream.Send(event); err != nil { + s.releaseEventSlot() + if err := s.sendRemoteEvent(stream, event); err != nil { + if errors.Is(err, io.ErrClosedPipe) { + return + } errCh <- err return } @@ -326,28 +490,107 @@ func (s *forwardLocalStream) serve(stream forwardrpc.ForwardListener_TaskStreamS return err } -func (s *forwardLocalStream) close() { +func (s *forwardLocalStream) close() error { + ctx, cancel := context.WithTimeout(context.Background(), forwardStreamCloseTimeout) + defer cancel() + return s.closeContext(ctx) +} + +func (s *forwardLocalStream) closeContext(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } s.closeOnce.Do(func() { + s.eventMu.Lock() + s.closed.Store(true) if s.done != nil { close(s.done) } + s.eventMu.Unlock() + + s.drained = make(chan struct{}) + go func() { + s.remoteSend.Wait() + close(s.drained) + }() + select { + case <-s.drained: + case <-ctx.Done(): + s.closeErr = fmt.Errorf("%w: %s:%s: %v", ErrForwardStreamCloseTimeout, s.listenerID, s.pipelineID, ctx.Err()) + } }) + return s.closeErr } func (s *forwardLocalStream) sendEvent(event *clientpb.SpiteRequest) error { + return s.sendEventContext(context.Background(), event) +} + +func (s *forwardLocalStream) sendEventContext(ctx context.Context, event *clientpb.SpiteRequest) error { if event == nil { return nil } + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return err + } + timer := time.NewTimer(10 * time.Second) + defer timer.Stop() + slots := s.eventSlotPool() select { - case s.events <- event: + case <-slots: + s.eventMu.Lock() + if s.closed.Load() || ctx.Err() != nil { + s.eventMu.Unlock() + s.releaseEventSlot() + if err := ctx.Err(); err != nil { + return err + } + return io.ErrClosedPipe + } + s.events <- event + s.eventMu.Unlock() return nil case <-s.done: return io.ErrClosedPipe - case <-time.After(10 * time.Second): + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: return fmt.Errorf("forward stream %s:%s event queue full", s.listenerID, s.pipelineID) } } +func (s *forwardLocalStream) eventSlotPool() chan struct{} { + s.slotsOnce.Do(func() { + capacity := cap(s.events) + s.eventSlots = make(chan struct{}, capacity) + for i := len(s.events); i < capacity; i++ { + s.eventSlots <- struct{}{} + } + }) + return s.eventSlots +} + +func (s *forwardLocalStream) releaseEventSlot() { + s.eventSlotPool() <- struct{}{} +} + +func (s *forwardLocalStream) sendRemoteEvent(stream forwardrpc.ForwardListener_TaskStreamServer, event *clientpb.SpiteRequest) error { + s.eventMu.Lock() + if s.closed.Load() { + s.eventMu.Unlock() + return io.ErrClosedPipe + } + s.remoteSend.Add(1) + s.eventMu.Unlock() + + defer s.remoteSend.Done() + err := stream.Send(event) + return err +} + func metadataValue(ctx context.Context, key string) (string, error) { md, ok := metadata.FromIncomingContext(ctx) if !ok { diff --git a/server/listener/forward_backpressure_audit_test.go b/server/listener/forward_backpressure_audit_test.go new file mode 100644 index 000000000..be370ff41 --- /dev/null +++ b/server/listener/forward_backpressure_audit_test.go @@ -0,0 +1,219 @@ +package listener + +import ( + "context" + "errors" + "io" + "testing" + "time" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + "google.golang.org/grpc/metadata" +) + +func newAuditForwardLocalStream(requestCapacity, eventCapacity int) *forwardLocalStream { + return &forwardLocalStream{ + listenerID: "audit-listener", + pipelineID: "audit-pipeline", + requests: make(chan *clientpb.SpiteRequest, requestCapacity), + events: make(chan *clientpb.SpiteRequest, eventCapacity), + done: make(chan struct{}), + } +} + +type auditRequestTaskStream struct { + ctx context.Context + req *clientpb.SpiteRequest + delivered chan struct{} + sent bool +} + +func (s *auditRequestTaskStream) Send(*clientpb.SpiteRequest) error { return nil } +func (s *auditRequestTaskStream) Recv() (*clientpb.SpiteRequest, error) { + if !s.sent { + s.sent = true + close(s.delivered) + return s.req, nil + } + <-s.ctx.Done() + return nil, io.EOF +} +func (s *auditRequestTaskStream) SetHeader(metadata.MD) error { return nil } +func (s *auditRequestTaskStream) SendHeader(metadata.MD) error { return nil } +func (s *auditRequestTaskStream) SetTrailer(metadata.MD) {} +func (s *auditRequestTaskStream) Context() context.Context { return s.ctx } +func (s *auditRequestTaskStream) SendMsg(interface{}) error { return nil } +func (s *auditRequestTaskStream) RecvMsg(interface{}) error { return nil } + +func TestAuditForwardRegistryUsesBoundedQueues(t *testing.T) { + stream := mustOpenForwardLocalStream(t, newForwardStreamRegistry(), "audit-listener", "audit-pipeline") + if got := cap(stream.requests); got != 255 { + t.Fatalf("request queue capacity = %d, want 255", got) + } + if got := cap(stream.events); got != 255 { + t.Fatalf("event queue capacity = %d, want 255", got) + } +} + +func TestAuditForwardEventQueueBackpressurePreservesEvent(t *testing.T) { + stream := newAuditForwardLocalStream(1, 2) + first := &clientpb.SpiteRequest{ListenerId: "event-1"} + second := &clientpb.SpiteRequest{ListenerId: "event-2"} + third := &clientpb.SpiteRequest{ListenerId: "event-3"} + stream.events <- first + stream.events <- second + + started := make(chan struct{}) + resultCh := make(chan error, 1) + go func() { + close(started) + resultCh <- stream.sendEvent(third) + }() + <-started + + select { + case err := <-resultCh: + t.Fatalf("sendEvent returned before queue capacity was released: %v", err) + case <-time.After(25 * time.Millisecond): + } + + if got := <-stream.events; got != first { + t.Fatalf("first queued event = %#v, want %#v", got, first) + } + stream.releaseEventSlot() + select { + case err := <-resultCh: + if err != nil { + t.Fatalf("sendEvent failed after capacity was released: %v", err) + } + case <-time.After(time.Second): + t.Fatal("sendEvent stayed blocked after capacity was released") + } + if got := <-stream.events; got != second { + t.Fatalf("second queued event = %#v, want %#v", got, second) + } + if got := <-stream.events; got != third { + t.Fatalf("backpressured event = %#v, want %#v", got, third) + } +} + +func TestAuditForwardEventQueueCloseUnblocksProducer(t *testing.T) { + stream := newAuditForwardLocalStream(1, 1) + stream.events <- &clientpb.SpiteRequest{ListenerId: "queue-full"} + + started := make(chan struct{}) + resultCh := make(chan error, 1) + go func() { + close(started) + resultCh <- stream.sendEvent(&clientpb.SpiteRequest{ListenerId: "blocked-event"}) + }() + <-started + + select { + case err := <-resultCh: + t.Fatalf("sendEvent returned before close: %v", err) + case <-time.After(25 * time.Millisecond): + } + + stream.close() + select { + case err := <-resultCh: + if !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("sendEvent error = %v, want %v", err, io.ErrClosedPipe) + } + case <-time.After(time.Second): + t.Fatal("sendEvent stayed blocked after stream close") + } +} + +func TestAuditForwardRequestQueueBackpressurePreservesOrder(t *testing.T) { + stream := newAuditForwardLocalStream(1, 1) + first := &clientpb.SpiteRequest{ListenerId: "request-1"} + second := &clientpb.SpiteRequest{ListenerId: "request-2"} + stream.requests <- first + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + remote := &auditRequestTaskStream{ + ctx: ctx, + req: second, + delivered: make(chan struct{}), + } + serveResult := make(chan error, 1) + go func() { + serveResult <- stream.serve(remote) + }() + + select { + case <-remote.delivered: + case <-time.After(time.Second): + t.Fatal("remote request was not received") + } + select { + case err := <-serveResult: + t.Fatalf("serve returned while request producer should be backpressured: %v", err) + case <-time.After(25 * time.Millisecond): + } + + if got, err := stream.Recv(); err != nil || got != first { + t.Fatalf("first Recv = (%#v, %v), want (%#v, nil)", got, err, first) + } + secondResult := make(chan *clientpb.SpiteRequest, 1) + go func() { + got, _ := stream.Recv() + secondResult <- got + }() + select { + case got := <-secondResult: + if got != second { + t.Fatalf("second Recv = %#v, want %#v", got, second) + } + case <-time.After(time.Second): + t.Fatal("request producer stayed blocked after capacity was released") + } + + stream.close() + select { + case err := <-serveResult: + if err != nil { + t.Fatalf("serve error after close = %v, want nil", err) + } + case <-time.After(time.Second): + t.Fatal("serve did not return after close") + } + cancel() +} + +func TestAuditForwardRequestQueueCloseUnblocksProducer(t *testing.T) { + stream := newAuditForwardLocalStream(1, 1) + stream.requests <- &clientpb.SpiteRequest{ListenerId: "queue-full"} + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + remote := &auditRequestTaskStream{ + ctx: ctx, + req: &clientpb.SpiteRequest{ListenerId: "blocked-request"}, + delivered: make(chan struct{}), + } + serveResult := make(chan error, 1) + go func() { + serveResult <- stream.serve(remote) + }() + + select { + case <-remote.delivered: + case <-time.After(time.Second): + t.Fatal("remote request was not received") + } + stream.close() + + select { + case err := <-serveResult: + if err != nil { + t.Fatalf("serve error after close = %v, want nil", err) + } + case <-time.After(time.Second): + t.Fatal("serve stayed blocked after close") + } + cancel() +} diff --git a/server/listener/forward_closed_stream_audit_test.go b/server/listener/forward_closed_stream_audit_test.go new file mode 100644 index 000000000..7c4ca8fb8 --- /dev/null +++ b/server/listener/forward_closed_stream_audit_test.go @@ -0,0 +1,186 @@ +package listener + +import ( + "context" + "errors" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + "google.golang.org/grpc/metadata" +) + +func TestAuditForwardSendEventRejectsAfterClose(t *testing.T) { + accepted := 0 + for range 64 { + stream := newAuditForwardLocalStream(1, 1) + stream.close() + err := stream.sendEvent(&clientpb.SpiteRequest{ListenerId: "after-close"}) + if err == nil { + accepted++ + continue + } + if !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("sendEvent error = %v, want %v", err, io.ErrClosedPipe) + } + } + + // Both select branches are ready for a closed stream with an empty queue. + // Missing all invalid accepts therefore has probability 2^-64. + if accepted > 0 { + t.Fatalf("sendEvent accepted %d of 64 events after close", accepted) + } +} + +func TestForwardSendEventClosedErrorNeverQueuesEvent(t *testing.T) { + for i := 0; i < 2000; i++ { + stream := newAuditForwardLocalStream(1, 1) + start := make(chan struct{}) + errCh := make(chan error, 1) + closedCh := make(chan struct{}) + go func() { + <-start + errCh <- stream.sendEvent(&clientpb.SpiteRequest{ListenerId: "racing-event"}) + }() + go func() { + <-start + stream.close() + close(closedCh) + }() + close(start) + + err := <-errCh + <-closedCh + if !errors.Is(err, io.ErrClosedPipe) { + continue + } + select { + case event := <-stream.events: + t.Fatalf("iteration %d: sendEvent returned ErrClosedPipe after queueing %#v", i, event) + default: + } + } +} + +type barrierForwardTaskStream struct { + ctx context.Context + sendEntered chan struct{} + releaseSend chan struct{} + enterOnce sync.Once + sendCount atomic.Int32 +} + +func (s *barrierForwardTaskStream) Send(*clientpb.SpiteRequest) error { + s.enterOnce.Do(func() { close(s.sendEntered) }) + <-s.releaseSend + s.sendCount.Add(1) + return nil +} + +func (s *barrierForwardTaskStream) Recv() (*clientpb.SpiteRequest, error) { + <-s.ctx.Done() + return nil, io.EOF +} + +func (s *barrierForwardTaskStream) SetHeader(metadata.MD) error { return nil } +func (s *barrierForwardTaskStream) SendHeader(metadata.MD) error { return nil } +func (s *barrierForwardTaskStream) SetTrailer(metadata.MD) {} +func (s *barrierForwardTaskStream) Context() context.Context { return s.ctx } +func (s *barrierForwardTaskStream) SendMsg(interface{}) error { return nil } +func (s *barrierForwardTaskStream) RecvMsg(interface{}) error { return nil } + +func TestForwardCloseWaitsForInFlightRemoteSend(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stream := newAuditForwardLocalStream(1, 1) + remote := &barrierForwardTaskStream{ + ctx: ctx, + sendEntered: make(chan struct{}), + releaseSend: make(chan struct{}), + } + serveResult := make(chan error, 1) + go func() { serveResult <- stream.serve(remote) }() + + if err := stream.sendEvent(&clientpb.SpiteRequest{ListenerId: "in-flight"}); err != nil { + t.Fatalf("sendEvent failed: %v", err) + } + select { + case <-remote.sendEntered: + case <-time.After(time.Second): + t.Fatal("remote Send was not entered") + } + + closeReturned := make(chan struct{}) + go func() { + stream.close() + close(closeReturned) + }() + select { + case <-closeReturned: + t.Fatal("close returned while remote Send was still in flight") + case <-time.After(25 * time.Millisecond): + } + + close(remote.releaseSend) + select { + case <-closeReturned: + case <-time.After(time.Second): + t.Fatal("close did not return after remote Send completed") + } + if got := remote.sendCount.Load(); got != 1 { + t.Fatalf("remote Send count = %d, want 1", got) + } + if err := stream.sendEvent(&clientpb.SpiteRequest{ListenerId: "after-close"}); !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("sendEvent after close = %v, want ErrClosedPipe", err) + } + + select { + case err := <-serveResult: + if err != nil { + t.Fatalf("serve error after close = %v, want nil", err) + } + case <-time.After(time.Second): + t.Fatal("serve did not return after close") + } +} + +func TestForwardCloseContextReturnsWhenRemoteSendDoesNotStop(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stream := newAuditForwardLocalStream(1, 1) + remote := &barrierForwardTaskStream{ + ctx: ctx, + sendEntered: make(chan struct{}), + releaseSend: make(chan struct{}), + } + serveResult := make(chan error, 1) + go func() { serveResult <- stream.serve(remote) }() + + if err := stream.sendEvent(&clientpb.SpiteRequest{ListenerId: "in-flight-timeout"}); err != nil { + t.Fatalf("sendEvent failed: %v", err) + } + select { + case <-remote.sendEntered: + case <-time.After(time.Second): + t.Fatal("remote Send was not entered") + } + + closeCtx, closeCancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer closeCancel() + if err := stream.closeContext(closeCtx); !errors.Is(err, ErrForwardStreamCloseTimeout) { + t.Fatalf("closeContext error = %v, want ErrForwardStreamCloseTimeout", err) + } + + close(remote.releaseSend) + select { + case err := <-serveResult: + if err != nil { + t.Fatalf("serve error after close = %v, want nil", err) + } + case <-time.After(time.Second): + t.Fatal("serve did not return after timed close") + } +} diff --git a/server/listener/forward_dispatch_test.go b/server/listener/forward_dispatch_test.go index 829608e87..8084a4ea2 100644 --- a/server/listener/forward_dispatch_test.go +++ b/server/listener/forward_dispatch_test.go @@ -261,13 +261,12 @@ func TestForwardDispatchMissingConnectionCompletesTaskAcrossPollingPipelines(t * }) deadline := time.Now().Add(500 * time.Millisecond) - for time.Now().Before(deadline) { - if task.Finished() { - return - } - time.Sleep(10 * time.Millisecond) + for !task.IsClosed() && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if !task.IsClosed() || !task.Finished() { + t.Fatalf("task %s stayed running after %s forward dispatch lost the connection", task.TaskID(), tc.name) } - t.Fatalf("task %s stayed running after %s forward dispatch lost the connection", task.TaskID(), tc.name) }) } } diff --git a/server/listener/forward_rollback_test.go b/server/listener/forward_rollback_test.go new file mode 100644 index 000000000..73c648e29 --- /dev/null +++ b/server/listener/forward_rollback_test.go @@ -0,0 +1,386 @@ +package listener + +import ( + "context" + "crypto/tls" + "errors" + "net" + "net/http" + "sync" + "sync/atomic" + "testing" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" + "github.com/chainreactors/malice-network/helper/implanttypes" + "github.com/chainreactors/malice-network/server/internal/core" + "google.golang.org/grpc" +) + +var errListenerInitialization = errors.New("listener initialization failed") + +type rollbackForwardStream struct { + closeCalls atomic.Int32 + closeErr error +} + +func (*rollbackForwardStream) Send(*clientpb.SpiteResponse) error { return nil } +func (*rollbackForwardStream) Recv() (*clientpb.SpiteRequest, error) { + return nil, errors.New("unexpected Recv") +} +func (s *rollbackForwardStream) CloseSend() error { + s.closeCalls.Add(1) + return s.closeErr +} + +type rollbackPipelineRPC struct { + stream *rollbackForwardStream +} + +type rollbackListener struct { + closed chan struct{} + closeOnce sync.Once + closeCalls atomic.Int32 + closeErr error +} + +func newRollbackListener() *rollbackListener { + return &rollbackListener{closed: make(chan struct{})} +} + +func (l *rollbackListener) Accept() (net.Conn, error) { + <-l.closed + return nil, net.ErrClosed +} +func (l *rollbackListener) Close() error { + l.closeOnce.Do(func() { + l.closeCalls.Add(1) + close(l.closed) + }) + return l.closeErr +} +func (*rollbackListener) Addr() net.Addr { return rollbackAddr("rollback") } + +type rollbackAddr string + +func (a rollbackAddr) Network() string { return string(a) } +func (a rollbackAddr) String() string { return string(a) } + +func (c *rollbackPipelineRPC) OpenForwardStream(context.Context, core.Pipeline) (core.ForwardStream, error) { + return c.stream, nil +} +func (*rollbackPipelineRPC) Register(context.Context, *clientpb.RegisterSession, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*rollbackPipelineRPC) Checkin(context.Context, *implantpb.Ping, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*rollbackPipelineRPC) GetArtifact(context.Context, *clientpb.Artifact, ...grpc.CallOption) (*clientpb.Artifact, error) { + return nil, errors.New("artifact unavailable") +} + +func assertForwardRolledBack(t *testing.T, pipeline core.Pipeline, stream *rollbackForwardStream, start func() error) { + t.Helper() + if err := start(); !errors.Is(err, errListenerInitialization) { + t.Fatalf("Start error = %v, want %v", err, errListenerInitialization) + } + if got := stream.closeCalls.Load(); got != 1 { + t.Fatalf("forward stream CloseSend calls = %d, want 1", got) + } + if got := core.Forwarders.Get(core.PipelineRuntimeKey(pipeline.ToProtobuf().ListenerId, pipeline.ID())); got != nil { + t.Fatalf("uncommitted forward remained in registry: %#v", got) + } +} + +func assertRollbackState(t *testing.T, pipeline core.Pipeline, stream *rollbackForwardStream) { + t.Helper() + if got := stream.closeCalls.Load(); got != 1 { + t.Fatalf("forward stream CloseSend calls = %d, want 1", got) + } + if got := core.Forwarders.Get(core.PipelineRuntimeKey(pipeline.ToProtobuf().ListenerId, pipeline.ID())); got != nil { + t.Fatalf("uncommitted forward remained in registry: %#v", got) + } +} + +func addReplacementForward(t *testing.T, pipeline core.Pipeline, stream *rollbackForwardStream) *core.Forward { + t.Helper() + forward, err := core.NewForward(&rollbackPipelineRPC{stream: stream}, pipeline) + if err != nil { + t.Fatalf("create replacement forward: %v", err) + } + forward.ListenerId = pipeline.ToProtobuf().ListenerId + core.Forwarders.Add(forward) + return forward +} + +func TestTCPPipelineRollsBackForwardWhenListenFails(t *testing.T) { + oldListen := tcpListen + tcpListen = func(string, string) (net.Listener, error) { return nil, errListenerInitialization } + t.Cleanup(func() { tcpListen = oldListen }) + + stream := &rollbackForwardStream{} + rpc := &rollbackPipelineRPC{stream: stream} + pipeline, err := NewTcpPipeline(rpc, &clientpb.Pipeline{ + Name: "rollback-tcp", ListenerId: "rollback-listener", Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + assertForwardRolledBack(t, pipeline, stream, pipeline.Start) +} + +func TestTCPPipelineStartFailureDiscardsLocalForwardStream(t *testing.T) { + registry := newForwardStreamRegistry() + listenerID := "rollback-local-listener" + pipelineID := "rollback-local-tcp" + key := core.PipelineRuntimeKey(listenerID, pipelineID) + var failedGeneration *forwardLocalStream + oldListen := tcpListen + tcpListen = func(string, string) (net.Listener, error) { + registry.mu.Lock() + failedGeneration = registry.streams[key] + registry.mu.Unlock() + return nil, errListenerInitialization + } + t.Cleanup(func() { tcpListen = oldListen }) + + pipeline, err := NewTcpPipeline(&forwardPipelineRPC{listenerID: listenerID, registry: registry}, &clientpb.Pipeline{ + Name: pipelineID, ListenerId: listenerID, Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + if err := pipeline.Start(); !errors.Is(err, errListenerInitialization) { + t.Fatalf("Start error = %v, want errListenerInitialization", err) + } + if failedGeneration == nil { + t.Fatal("failed start did not create a forward stream") + } + if !failedGeneration.closed.Load() { + t.Fatal("failed start left its forward stream open") + } + fresh := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) + if fresh == failedGeneration { + t.Fatal("retry reused the failed forward stream generation") + } +} + +func TestTCPPipelineReturnsListenerAndForwardRollbackErrors(t *testing.T) { + oldListen := tcpListen + tcpListen = func(string, string) (net.Listener, error) { return nil, errListenerInitialization } + t.Cleanup(func() { tcpListen = oldListen }) + + errClose := errors.New("forward close failed") + stream := &rollbackForwardStream{closeErr: errClose} + pipeline, err := NewTcpPipeline(&rollbackPipelineRPC{stream: stream}, &clientpb.Pipeline{ + Name: "rollback-errors", ListenerId: "rollback-listener", Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + err = pipeline.Start() + if !errors.Is(err, errListenerInitialization) || !errors.Is(err, errClose) { + t.Fatalf("Start error = %v, want listener and rollback errors", err) + } +} + +func TestHTTPPipelineRollsBackForwardWhenListenFails(t *testing.T) { + oldListen := httpListen + httpListen = func(string, string) (net.Listener, error) { return nil, errListenerInitialization } + t.Cleanup(func() { httpListen = oldListen }) + + stream := &rollbackForwardStream{} + rpc := &rollbackPipelineRPC{stream: stream} + pipeline, err := NewHttpPipeline(rpc, &clientpb.Pipeline{ + Name: "rollback-http", ListenerId: "rollback-listener", Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: (&implanttypes.PipelineParams{}).String()}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + assertForwardRolledBack(t, pipeline, stream, pipeline.Start) +} + +func TestHTTPPipelineStartFailureDiscardsLocalForwardStream(t *testing.T) { + registry := newForwardStreamRegistry() + listenerID := "rollback-local-listener" + pipelineID := "rollback-local-http" + key := core.PipelineRuntimeKey(listenerID, pipelineID) + var failedGeneration *forwardLocalStream + oldListen := httpListen + httpListen = func(string, string) (net.Listener, error) { + registry.mu.Lock() + failedGeneration = registry.streams[key] + registry.mu.Unlock() + return nil, errListenerInitialization + } + t.Cleanup(func() { httpListen = oldListen }) + + pipeline, err := NewHttpPipeline(&forwardPipelineRPC{listenerID: listenerID, registry: registry}, &clientpb.Pipeline{ + Name: pipelineID, ListenerId: listenerID, Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: (&implanttypes.PipelineParams{}).String()}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + if err := pipeline.Start(); !errors.Is(err, errListenerInitialization) { + t.Fatalf("Start error = %v, want errListenerInitialization", err) + } + if failedGeneration == nil { + t.Fatal("failed start did not create a forward stream") + } + if !failedGeneration.closed.Load() { + t.Fatal("failed start left its forward stream open") + } + fresh := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) + if fresh == failedGeneration { + t.Fatal("retry reused the failed forward stream generation") + } +} + +func TestHTTPPipelineReturnsListenerAndForwardRollbackErrors(t *testing.T) { + oldListen := httpListen + httpListen = func(string, string) (net.Listener, error) { return nil, errListenerInitialization } + t.Cleanup(func() { httpListen = oldListen }) + + errClose := errors.New("forward close failed") + stream := &rollbackForwardStream{closeErr: errClose} + pipeline, err := NewHttpPipeline(&rollbackPipelineRPC{stream: stream}, &clientpb.Pipeline{ + Name: "rollback-http-errors", ListenerId: "rollback-listener", Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: (&implanttypes.PipelineParams{}).String()}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + err = pipeline.Start() + if !errors.Is(err, errListenerInitialization) || !errors.Is(err, errClose) { + t.Fatalf("Start error = %v, want listener and rollback errors", err) + } +} + +func TestTCPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { + oldListen, oldCmux := tcpListen, tcpStartCmux + errClose := errors.New("tcp listener close failed") + ln := newRollbackListener() + ln.closeErr = errClose + tcpListen = func(string, string) (net.Listener, error) { return ln, nil } + tcpStartCmux = func(net.Listener, *tls.Config, func(net.Conn), core.GoErrorHandler) (net.Listener, error) { + return nil, errListenerInitialization + } + t.Cleanup(func() { tcpListen, tcpStartCmux = oldListen, oldCmux }) + + stream := &rollbackForwardStream{} + pipeline, err := NewTcpPipeline(&rollbackPipelineRPC{stream: stream}, &clientpb.Pipeline{ + Name: "rollback-tcp-cmux", ListenerId: "rollback-listener", Tls: &clientpb.TLS{Enable: true}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + err = pipeline.Start() + if !errors.Is(err, errListenerInitialization) || !errors.Is(err, errClose) { + t.Fatalf("Start error = %v, want initialization and listener close errors", err) + } + assertRollbackState(t, pipeline, stream) + if got := ln.closeCalls.Load(); got != 1 { + t.Fatalf("listener Close calls = %d, want 1", got) + } +} + +func TestHTTPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { + oldListen, oldCmux := httpListen, httpStartWithCmux + errClose := errors.New("http listener close failed") + ln := newRollbackListener() + ln.closeErr = errClose + httpListen = func(string, string) (net.Listener, error) { return ln, nil } + httpStartWithCmux = func(*HTTPPipeline, net.Listener, *http.ServeMux) error { + return errListenerInitialization + } + t.Cleanup(func() { httpListen, httpStartWithCmux = oldListen, oldCmux }) + + stream := &rollbackForwardStream{} + pipeline, err := NewHttpPipeline(&rollbackPipelineRPC{stream: stream}, &clientpb.Pipeline{ + Name: "rollback-http-cmux", ListenerId: "rollback-listener", Tls: &clientpb.TLS{Enable: true, Cert: &clientpb.Cert{}}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: (&implanttypes.PipelineParams{}).String()}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + err = pipeline.Start() + if !errors.Is(err, errListenerInitialization) || !errors.Is(err, errClose) { + t.Fatalf("Start error = %v, want initialization and listener close errors", err) + } + assertRollbackState(t, pipeline, stream) + if got := ln.closeCalls.Load(); got != 1 { + t.Fatalf("listener Close calls = %d, want 1", got) + } +} + +func TestTCPPipelineFailedStartPreservesReplacementForward(t *testing.T) { + oldListen, oldCmux := tcpListen, tcpStartCmux + ln := newRollbackListener() + tcpListen = func(string, string) (net.Listener, error) { return ln, nil } + t.Cleanup(func() { tcpListen, tcpStartCmux = oldListen, oldCmux }) + + stream := &rollbackForwardStream{} + pipeline, err := NewTcpPipeline(&rollbackPipelineRPC{stream: stream}, &clientpb.Pipeline{ + Name: "rollback-tcp-replacement", ListenerId: "rollback-listener", Tls: &clientpb.TLS{Enable: true}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + key := core.PipelineRuntimeKey(pipeline.ListenerID, pipeline.ID()) + var replacement *core.Forward + tcpStartCmux = func(net.Listener, *tls.Config, func(net.Conn), core.GoErrorHandler) (net.Listener, error) { + replacement = addReplacementForward(t, pipeline, &rollbackForwardStream{}) + return nil, errListenerInitialization + } + t.Cleanup(func() { _ = core.Forwarders.Remove(key) }) + + if err := pipeline.Start(); !errors.Is(err, errListenerInitialization) { + t.Fatalf("Start error = %v, want %v", err, errListenerInitialization) + } + if got := core.Forwarders.Get(key); got != replacement { + t.Fatalf("replacement forward = %#v, want %#v", got, replacement) + } + if got := stream.closeCalls.Load(); got != 1 { + t.Fatalf("failed generation CloseSend calls = %d, want 1", got) + } +} + +func TestHTTPPipelineFailedStartPreservesReplacementForward(t *testing.T) { + oldListen, oldCmux := httpListen, httpStartWithCmux + ln := newRollbackListener() + httpListen = func(string, string) (net.Listener, error) { return ln, nil } + t.Cleanup(func() { httpListen, httpStartWithCmux = oldListen, oldCmux }) + + stream := &rollbackForwardStream{} + pipeline, err := NewHttpPipeline(&rollbackPipelineRPC{stream: stream}, &clientpb.Pipeline{ + Name: "rollback-http-replacement", ListenerId: "rollback-listener", Tls: &clientpb.TLS{Enable: true, Cert: &clientpb.Cert{}}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: (&implanttypes.PipelineParams{}).String()}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + key := core.PipelineRuntimeKey(pipeline.ListenerID, pipeline.ID()) + var replacement *core.Forward + httpStartWithCmux = func(*HTTPPipeline, net.Listener, *http.ServeMux) error { + replacement = addReplacementForward(t, pipeline, &rollbackForwardStream{}) + return errListenerInitialization + } + t.Cleanup(func() { _ = core.Forwarders.Remove(key) }) + + if err := pipeline.Start(); !errors.Is(err, errListenerInitialization) { + t.Fatalf("Start error = %v, want %v", err, errListenerInitialization) + } + if got := core.Forwarders.Get(key); got != replacement { + t.Fatalf("replacement forward = %#v, want %#v", got, replacement) + } + if got := stream.closeCalls.Load(); got != 1 { + t.Fatalf("failed generation CloseSend calls = %d, want 1", got) + } +} diff --git a/server/listener/forward_runtime_test.go b/server/listener/forward_runtime_test.go index 70d35b591..953fe2acf 100644 --- a/server/listener/forward_runtime_test.go +++ b/server/listener/forward_runtime_test.go @@ -19,9 +19,21 @@ import ( "github.com/chainreactors/malice-network/server/internal/configs" "github.com/chainreactors/malice-network/server/internal/core" "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" "gopkg.in/yaml.v3" ) +func mustOpenForwardLocalStream(t testing.TB, registry *forwardStreamRegistry, listenerID, pipelineID string) *forwardLocalStream { + t.Helper() + stream, err := registry.open(listenerID, pipelineID) + if err != nil { + t.Fatalf("open forward stream failed: %v", err) + } + return stream +} + func reserveForwardPort(t testing.TB) uint16 { t.Helper() ln, err := net.Listen("tcp", "127.0.0.1:0") @@ -190,7 +202,7 @@ func TestForwardStreamRegistryKeepsStreamUntilPipelineStop(t *testing.T) { } t.Cleanup(func() { _ = started.Close() }) - beforeDisconnect := registry.get(listenerID, pipelineID) + beforeDisconnect := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) taskStream := &mockTaskStream{ ctx: context.Background(), recvErr: io.EOF, @@ -198,7 +210,7 @@ func TestForwardStreamRegistryKeepsStreamUntilPipelineStop(t *testing.T) { if err := beforeDisconnect.serve(taskStream); err != nil { t.Fatalf("serve EOF error = %v", err) } - if afterDisconnect := registry.get(listenerID, pipelineID); afterDisconnect != beforeDisconnect { + if afterDisconnect := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID); afterDisconnect != beforeDisconnect { t.Fatal("TaskStream disconnect should keep the existing registry stream") } @@ -210,7 +222,7 @@ func TestForwardStreamRegistryKeepsStreamUntilPipelineStop(t *testing.T) { if status == nil || status.Status != consts.CtrlStatusSuccess { t.Fatalf("stop status = %#v, want success", status) } - afterStop := registry.get(listenerID, pipelineID) + afterStop := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) if afterStop == beforeDisconnect { t.Fatal("pipeline stop should remove the old registry stream") } @@ -224,6 +236,194 @@ func TestForwardStreamRegistryKeepsStreamUntilPipelineStop(t *testing.T) { } } +func TestForwardStreamRegistryRejectsOpenWhileRetiring(t *testing.T) { + registry := newForwardStreamRegistry() + oldStream := mustOpenForwardLocalStream(t, registry, "listener-retire", "pipeline-retire") + + release, err := registry.retire("listener-retire", "pipeline-retire") + if err != nil { + t.Fatalf("retire failed: %v", err) + } + if _, err := registry.open("listener-retire", "pipeline-retire"); !errors.Is(err, ErrForwardStreamClosing) { + t.Fatalf("open while retiring error = %v, want ErrForwardStreamClosing", err) + } + select { + case err := <-recvForwardStream(oldStream): + if !errors.Is(err, io.EOF) { + t.Fatalf("retired stream Recv error = %v, want EOF", err) + } + case <-time.After(time.Second): + t.Fatal("retired stream Recv did not unblock") + } + + release() + fresh, err := registry.open("listener-retire", "pipeline-retire") + if err != nil { + t.Fatalf("open after retirement failed: %v", err) + } + if fresh == oldStream { + t.Fatal("open after retirement reused the retired stream") + } +} + +func TestForwardStreamRegistryBlocksReuseUntilTimedOutSendDrains(t *testing.T) { + registry := newForwardStreamRegistry() + listenerID := "listener-poisoned" + pipelineID := "pipeline-poisoned" + oldStream := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + remote := &barrierForwardTaskStream{ + ctx: ctx, + sendEntered: make(chan struct{}), + releaseSend: make(chan struct{}), + } + serveResult := make(chan error, 1) + go func() { serveResult <- oldStream.serve(remote) }() + if err := oldStream.sendEvent(&clientpb.SpiteRequest{ListenerId: listenerID}); err != nil { + t.Fatalf("sendEvent failed: %v", err) + } + select { + case <-remote.sendEntered: + case <-time.After(time.Second): + t.Fatal("remote Send was not entered") + } + + closeCtx, closeCancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer closeCancel() + releaseRetirement, err := registry.retireContext(closeCtx, listenerID, pipelineID) + if !errors.Is(err, ErrForwardStreamCloseTimeout) { + t.Fatalf("retireContext error = %v, want ErrForwardStreamCloseTimeout", err) + } + releaseRetirement() + if _, err := registry.open(listenerID, pipelineID); !errors.Is(err, ErrForwardStreamPoisoned) { + t.Fatalf("open before old send drained error = %v, want ErrForwardStreamPoisoned", err) + } + + close(remote.releaseSend) + select { + case err := <-serveResult: + if err != nil { + t.Fatalf("serve error after release = %v", err) + } + case <-time.After(time.Second): + t.Fatal("serve did not return after release") + } + + deadline := time.Now().Add(time.Second) + for { + fresh, err := registry.open(listenerID, pipelineID) + if err == nil { + if fresh == oldStream { + t.Fatal("registry reused the drained old stream") + } + break + } + if !errors.Is(err, ErrForwardStreamPoisoned) { + t.Fatalf("open after drain error = %v", err) + } + if time.Now().After(deadline) { + t.Fatal("registry stayed poisoned after old send drained") + } + time.Sleep(time.Millisecond) + } +} + +func TestForwardLocalStreamAllowsOnlyOneActiveAttach(t *testing.T) { + stream := newForwardLocalStream(nil, "listener-attach", "pipeline-attach") + + releaseFirst, err := stream.attach() + if err != nil { + t.Fatalf("first attach failed: %v", err) + } + if _, err := stream.attach(); !errors.Is(err, ErrForwardStreamAlreadyAttached) { + t.Fatalf("second active attach error = %v, want ErrForwardStreamAlreadyAttached", err) + } + + releaseFirst() + releaseSecond, err := stream.attach() + if err != nil { + t.Fatalf("attach after release failed: %v", err) + } + defer releaseSecond() + + // A stale release must not detach the newer owner. + releaseFirst() + if _, err := stream.attach(); !errors.Is(err, ErrForwardStreamAlreadyAttached) { + t.Fatalf("attach after stale release error = %v, want ErrForwardStreamAlreadyAttached", err) + } +} + +func TestForwardTaskStreamRejectsSecondActiveAttach(t *testing.T) { + registry := newForwardStreamRegistry() + service := &forwardListenerService{registry: registry} + listenerID := "listener-task-attach" + pipelineID := "pipeline-task-attach" + firstCtx, cancelFirst := context.WithCancel(metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "listener_id", listenerID, + "pipeline_id", pipelineID, + ))) + first := &mockTaskStream{ctx: firstCtx} + firstResult := make(chan error, 1) + go func() { firstResult <- service.TaskStream(first) }() + + deadline := time.Now().Add(time.Second) + for { + local := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) + local.attachMu.Lock() + attached := local.attached + local.attachMu.Unlock() + if attached { + break + } + if time.Now().After(deadline) { + cancelFirst() + t.Fatal("first TaskStream did not become active") + } + time.Sleep(time.Millisecond) + } + + secondCtx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "listener_id", listenerID, + "pipeline_id", pipelineID, + )) + if err := service.TaskStream(&mockTaskStream{ctx: secondCtx}); status.Code(err) != codes.AlreadyExists { + t.Fatalf("second TaskStream error = %v, want code %s", err, codes.AlreadyExists) + } + + cancelFirst() + select { + case err := <-firstResult: + if err != nil { + t.Fatalf("first TaskStream error = %v", err) + } + case <-time.After(time.Second): + t.Fatal("first TaskStream did not stop after context cancellation") + } +} + +func TestForwardPipelineRPCDoesNotCreateStreamForCanceledContext(t *testing.T) { + registry := newForwardStreamRegistry() + rpc := &forwardPipelineRPC{listenerID: "listener-canceled", registry: registry} + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := rpc.Register(ctx, &clientpb.RegisterSession{ + ListenerId: "listener-canceled", + PipelineId: "pipeline-canceled", + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Register error = %v, want context.Canceled", err) + } + + registry.mu.Lock() + streamCount := len(registry.streams) + registry.mu.Unlock() + if streamCount != 0 { + t.Fatalf("registry stream count = %d, want 0", streamCount) + } +} + func TestForwardStreamRegistryCleanupOnListenerClose(t *testing.T) { listenerID := "forward-close-listener" pipelineID := "forward-close-pipeline" @@ -249,12 +449,12 @@ func TestForwardStreamRegistryCleanupOnListenerClose(t *testing.T) { Secure: &clientpb.Secure{}, } lns.pipelines.Add(NewCustomPipeline(pipeline)) - beforeClose := registry.get(listenerID, pipelineID) + beforeClose := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) if err := lns.Close(); err != nil { t.Fatalf("listener Close failed: %v", err) } - afterClose := registry.get(listenerID, pipelineID) + afterClose := mustOpenForwardLocalStream(t, registry, listenerID, pipelineID) if afterClose == beforeClose { t.Fatal("listener Close should remove the old registry stream") } diff --git a/server/listener/http.go b/server/listener/http.go index 7c16128cd..b02c0ac20 100644 --- a/server/listener/http.go +++ b/server/listener/http.go @@ -14,6 +14,8 @@ import ( "net" "net/http" "strconv" + "sync" + "time" "github.com/chainreactors/logs" "github.com/chainreactors/malice-network/server/internal/certutils" @@ -22,6 +24,13 @@ import ( cryptostream "github.com/chainreactors/malice-network/server/internal/stream" ) +var ( + httpListen = net.Listen + httpStartWithCmux = func(pipeline *HTTPPipeline, ln net.Listener, mux *http.ServeMux) error { + return pipeline.startWithCmux(ln, mux) + } +) + func NewHttpPipeline(rpc pipelineRPCClient, pipeline *clientpb.Pipeline) (*HTTPPipeline, error) { http := pipeline.GetHttp() @@ -47,14 +56,17 @@ func NewHttpPipeline(rpc pipelineRPCClient, pipeline *clientpb.Pipeline) (*HTTPP } type HTTPPipeline struct { - srv net.Listener - rpc pipelineRPCClient - Name string - Port uint16 - Host string - Enable bool - Target []string - CertName string + stateMu sync.RWMutex + starting bool + startDone chan struct{} + srv net.Listener + rpc pipelineRPCClient + Name string + Port uint16 + Host string + Enable bool + Target []string + CertName string *core.PipelineConfig Headers map[string][]string ErrorPage []byte @@ -63,6 +75,9 @@ type HTTPPipeline struct { } func (pipeline *HTTPPipeline) ToProtobuf() *clientpb.Pipeline { + pipeline.stateMu.RLock() + enabled := pipeline.Enable + pipeline.stateMu.RUnlock() params := (&implanttypes.PipelineParams{ Headers: pipeline.Headers, ErrorPage: string(pipeline.ErrorPage), @@ -72,7 +87,7 @@ func (pipeline *HTTPPipeline) ToProtobuf() *clientpb.Pipeline { p := &clientpb.Pipeline{ Name: pipeline.Name, - Enable: pipeline.Enable, + Enable: enabled, Type: consts.HTTPPipeline, ListenerId: pipeline.ListenerID, Parser: pipeline.Parser, @@ -98,42 +113,48 @@ func (pipeline *HTTPPipeline) ID() string { } func (pipeline *HTTPPipeline) Close() error { + pipeline.stateMu.Lock() pipeline.Enable = false - if pipeline.srv != nil { - ln := pipeline.srv - pipeline.srv = nil - if err := ln.Close(); err != nil && !errors.Is(err, net.ErrClosed) { - return err - } + if pipeline.starting { + done := pipeline.startDone + pipeline.stateMu.Unlock() + <-done return nil } - return nil + ln := pipeline.srv + pipeline.srv = nil + pipeline.stateMu.Unlock() + if ln == nil { + return nil + } + return closePipelineListener(ln) } -func (pipeline *HTTPPipeline) Start() error { - if pipeline.Enable { +func (pipeline *HTTPPipeline) Start() (err error) { + if !pipeline.beginStart() { return nil } + forward, err := core.NewForward(pipeline.rpc, pipeline) if err != nil { + pipeline.abortStart() return err } forward.ListenerId = pipeline.ListenerID - core.Forwarders.Add(forward) - core.GoGuarded("http-forward-recv:"+pipeline.Name, func() error { - for { - msg, err := forward.Stream.Recv() - if err != nil { - if !pipeline.Enable { - return nil - } - return fmt.Errorf("http pipeline %s forward recv: %w", pipeline.Name, err) + committed := false + registered := false + defer func() { + if !committed { + if registered { + err = errors.Join(err, pipeline.rollbackCommittedStart(forward)) + } else { + err = errors.Join(err, forward.Abort()) } - dispatchForwardTaskRequest("http", pipeline.Name, msg) + pipeline.abortStart() } - }, pipeline.runtimeErrorHandler("forward recv loop")) + }() - ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", pipeline.Host, pipeline.Port)) + ln, err := httpListen("tcp", fmt.Sprintf("%s:%d", pipeline.Host, pipeline.Port)) if err != nil { return err } @@ -141,14 +162,16 @@ func (pipeline *HTTPPipeline) Start() error { mux := http.NewServeMux() mux.HandleFunc("/", pipeline.handler) + if !pipeline.commitStart(ln, forward) { + return closePipelineListener(ln) + } + registered = true if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable && pipeline.TLSConfig.Cert != nil { - err := pipeline.startWithCmux(ln, mux) + err := httpStartWithCmux(pipeline, ln, mux) if err != nil { return err } } else { - // 非 TLS 模式，使用原有逻辑 - pipeline.srv = ln server := NewHTTPServer(mux) core.GoGuarded("http-serve:"+pipeline.Name, func() error { if err := serveHTTP(server, ln); err != nil && err != http.ErrServerClosed && !errors.Is(err, net.ErrClosed) { @@ -157,13 +180,105 @@ func (pipeline *HTTPPipeline) Start() error { return nil }, pipeline.runtimeErrorHandler("serve loop")) } + if !pipeline.finishStart() { + return nil + } + committed = true + pipeline.startForwardRecv(forward) logs.Log.Infof("pipeline.http - start host=%s port=%d parser=%s tls=%t", pipeline.Host, pipeline.Port, pipeline.Parser, pipeline.TLSConfig.Enable) - pipeline.Enable = true return nil } +func (pipeline *HTTPPipeline) enabled() bool { + pipeline.stateMu.RLock() + defer pipeline.stateMu.RUnlock() + return pipeline.Enable +} + +func (pipeline *HTTPPipeline) beginStart() bool { + for { + pipeline.stateMu.Lock() + if pipeline.starting { + done := pipeline.startDone + pipeline.stateMu.Unlock() + <-done + continue + } + if pipeline.Enable { + pipeline.stateMu.Unlock() + return false + } + pipeline.Enable = true + pipeline.starting = true + pipeline.startDone = make(chan struct{}) + pipeline.stateMu.Unlock() + return true + } +} + +func (pipeline *HTTPPipeline) abortStart() { + pipeline.stateMu.Lock() + pipeline.starting = false + pipeline.Enable = false + done := pipeline.startDone + pipeline.startDone = nil + if done != nil { + close(done) + } + pipeline.stateMu.Unlock() +} + +func (pipeline *HTTPPipeline) commitStart(ln net.Listener, forward *core.Forward) bool { + pipeline.stateMu.Lock() + defer pipeline.stateMu.Unlock() + if !pipeline.Enable { + return false + } + pipeline.srv = ln + core.Forwarders.Add(forward) + return true +} + +func (pipeline *HTTPPipeline) finishStart() bool { + pipeline.stateMu.Lock() + defer pipeline.stateMu.Unlock() + if !pipeline.Enable { + return false + } + pipeline.starting = false + done := pipeline.startDone + pipeline.startDone = nil + close(done) + return true +} + +func (pipeline *HTTPPipeline) rollbackCommittedStart(forward *core.Forward) error { + forwardErr := core.Forwarders.RemoveIfSame(forward.RuntimeKey(), forward) + pipeline.stateMu.Lock() + ln := pipeline.srv + pipeline.srv = nil + pipeline.Enable = false + pipeline.stateMu.Unlock() + return errors.Join(forwardErr, closePipelineListener(ln)) +} + +func (pipeline *HTTPPipeline) startForwardRecv(forward *core.Forward) { + core.GoGuarded("http-forward-recv:"+pipeline.Name, func() error { + for { + msg, err := forward.Stream.Recv() + if err != nil { + if !pipeline.enabled() { + return nil + } + return fmt.Errorf("http pipeline %s forward recv: %w", pipeline.Name, err) + } + dispatchForwardTaskRequest("http", pipeline.Name, msg) + } + }, pipeline.runtimeErrorHandler("forward recv loop")) +} + // startWithCmux 使用 cmux 实现 HTTP TLS 和非 TLS 的端口复用 func (pipeline *HTTPPipeline) startWithCmux(ln net.Listener, mux *http.ServeMux) error { // 获取 TLS 配置 @@ -177,9 +292,6 @@ func (pipeline *HTTPPipeline) startWithCmux(ln net.Listener, mux *http.ServeMux) return err } - // 保存服务器引用用于关闭 - pipeline.srv = ln - return nil } @@ -313,6 +425,10 @@ func (h *httpReadWriter) Close() error { return nil } +func (h *httpReadWriter) SetReadDeadline(deadline time.Time) error { + return http.NewResponseController(h.writer).SetReadDeadline(deadline) +} + func (h *httpReadWriter) RemoteAddr() net.Addr { return h.remoteAddr } @@ -348,10 +464,7 @@ func (pipeline *HTTPPipeline) runtimeErrorHandler(scope string) core.GoErrorHand return core.CombineErrorHandlers( core.LogGuardedError(label), func(err error) { - pipeline.Enable = false - if pipeline.srv != nil { - _ = pipeline.srv.Close() - } + _ = pipeline.Close() if core.EventBroker != nil { core.EventBroker.Publish(core.Event{ EventType: consts.EventListener, diff --git a/server/listener/http_pipeline_e2e_test.go b/server/listener/http_pipeline_e2e_test.go index b0cf6dd7b..42c036ab5 100644 --- a/server/listener/http_pipeline_e2e_test.go +++ b/server/listener/http_pipeline_e2e_test.go @@ -5,7 +5,9 @@ import ( "context" "crypto/sha256" "encoding/binary" + "errors" "io" + "net" "net/http" "net/http/httptest" "testing" @@ -100,6 +102,42 @@ func TestHTTPPipelineHandlerDeliversLargeScreenshotSizedResponse(t *testing.T) { } } +func TestHTTPReadWriterReadDeadlineInterruptsSlowBody(t *testing.T) { + readResult := make(chan error, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rw := &httpReadWriter{body: r.Body, writer: w} + if err := rw.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil { + readResult <- err + return + } + _, err := rw.Read(make([]byte, 1)) + readResult <- err + })) + defer server.Close() + + conn, err := net.Dial("tcp", server.Listener.Addr().String()) + if err != nil { + t.Fatalf("dial HTTP server: %v", err) + } + defer conn.Close() + if _, err := io.WriteString(conn, "POST / HTTP/1.1\r\nHost: example.test\r\nContent-Length: 1\r\nConnection: close\r\n\r\n"); err != nil { + t.Fatalf("write partial HTTP request: %v", err) + } + + select { + case err := <-readResult: + if err == nil { + t.Fatal("slow request body read returned without a deadline error") + } + var netErr net.Error + if !errors.As(err, &netErr) || !netErr.Timeout() { + t.Fatalf("slow request body read error = %v, want timeout", err) + } + case <-time.After(2 * time.Second): + t.Fatal("slow request body read was not interrupted by its deadline") + } +} + func newHTTPProbePipeline(listenerID, name string) *HTTPPipeline { return &HTTPPipeline{ Name: name, @@ -123,7 +161,7 @@ func serveMaleficHTTPPacket(t testing.TB, pipeline *HTTPPipeline, packet []byte) t.Helper() req := httptest.NewRequest(http.MethodPost, "http://example.test/", bytes.NewReader(packet)) - resp := httptest.NewRecorder() + resp := &deadlineResponseRecorder{ResponseRecorder: httptest.NewRecorder()} pipeline.handler(resp, req) if resp.Code >= 400 { @@ -131,6 +169,14 @@ func serveMaleficHTTPPacket(t testing.TB, pipeline *HTTPPipeline, packet []byte) } } +type deadlineResponseRecorder struct { + *httptest.ResponseRecorder +} + +func (r *deadlineResponseRecorder) SetReadDeadline(time.Time) error { + return nil +} + func maleficBinaryResponsePacket(t testing.TB, rawID uint32, payload []byte) []byte { t.Helper() diff --git a/server/listener/lifecycle_concurrency_audit_test.go b/server/listener/lifecycle_concurrency_audit_test.go new file mode 100644 index 000000000..29f7a2cb3 --- /dev/null +++ b/server/listener/lifecycle_concurrency_audit_test.go @@ -0,0 +1,75 @@ +package listener + +import ( + "sync" + "testing" + + "github.com/chainreactors/malice-network/server/internal/core" +) + +func runAuditConcurrentClose(t *testing.T, closeFn func() error) { + t.Helper() + + const ( + workers = 8 + iterations = 1000 + ) + + start := make(chan struct{}) + errs := make(chan error, workers*iterations) + var wg sync.WaitGroup + wg.Add(workers) + for range workers { + go func() { + defer wg.Done() + <-start + for range iterations { + if err := closeFn(); err != nil { + errs <- err + } + } + }() + } + + close(start) + wg.Wait() + close(errs) + for err := range errs { + t.Errorf("concurrent Close returned an error: %v", err) + } +} + +func TestAuditTCPPipelineConcurrentClose(t *testing.T) { + pipeline := &TCPPipeline{Enable: true} + runAuditConcurrentClose(t, pipeline.Close) +} + +func TestAuditHTTPPipelineConcurrentClose(t *testing.T) { + pipeline := &HTTPPipeline{Enable: true} + runAuditConcurrentClose(t, pipeline.Close) +} + +func TestAuditREMConcurrentClose(t *testing.T) { + pipeline := &REM{Enable: true} + runAuditConcurrentClose(t, pipeline.Close) +} + +func TestAuditListenerConcurrentClose(t *testing.T) { + lns := &listener{ + Name: "audit-listener-close", + pipelines: core.NewPipelines(), + websites: make(map[string]*Website), + } + + listenerGlobalMu.Lock() + oldListener := currentListener + currentListener = lns + listenerGlobalMu.Unlock() + t.Cleanup(func() { + listenerGlobalMu.Lock() + currentListener = oldListener + listenerGlobalMu.Unlock() + }) + + runAuditConcurrentClose(t, lns.Close) +} diff --git a/server/listener/lifecycle_linearization_test.go b/server/listener/lifecycle_linearization_test.go new file mode 100644 index 000000000..c6ee4f0d4 --- /dev/null +++ b/server/listener/lifecycle_linearization_test.go @@ -0,0 +1,385 @@ +package listener + +import ( + "context" + "crypto/tls" + "errors" + "net" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" + remhelper "github.com/chainreactors/malice-network/helper/third/rem" + "github.com/chainreactors/malice-network/server/internal/core" + "google.golang.org/grpc" +) + +var ( + errFirstStart = errors.New("first start failed") + errSecondStart = errors.New("second start failed") +) + +type linearizedStartRPC struct { + calls atomic.Int32 + firstEntered chan struct{} + releaseFirst chan struct{} +} + +type committedStartStream struct { + done chan struct{} + once sync.Once +} + +func newCommittedStartStream() *committedStartStream { + return &committedStartStream{done: make(chan struct{})} +} + +func (*committedStartStream) Send(*clientpb.SpiteResponse) error { return nil } +func (s *committedStartStream) Recv() (*clientpb.SpiteRequest, error) { + <-s.done + return nil, context.Canceled +} +func (s *committedStartStream) CloseSend() error { + s.once.Do(func() { close(s.done) }) + return nil +} + +type committedStartRPC struct { + stream *committedStartStream +} + +func (c *committedStartRPC) OpenForwardStream(context.Context, core.Pipeline) (core.ForwardStream, error) { + return c.stream, nil +} +func (*committedStartRPC) Register(context.Context, *clientpb.RegisterSession, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*committedStartRPC) Checkin(context.Context, *implantpb.Ping, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*committedStartRPC) GetArtifact(context.Context, *clientpb.Artifact, ...grpc.CallOption) (*clientpb.Artifact, error) { + return nil, errors.New("artifact unavailable") +} + +func (c *linearizedStartRPC) OpenForwardStream(context.Context, core.Pipeline) (core.ForwardStream, error) { + if c.calls.Add(1) == 1 { + close(c.firstEntered) + <-c.releaseFirst + return nil, errFirstStart + } + return nil, errSecondStart +} +func (*linearizedStartRPC) Register(context.Context, *clientpb.RegisterSession, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*linearizedStartRPC) Checkin(context.Context, *implantpb.Ping, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*linearizedStartRPC) GetArtifact(context.Context, *clientpb.Artifact, ...grpc.CallOption) (*clientpb.Artifact, error) { + return nil, errors.New("artifact unavailable") +} + +func assertCloseWaitsForStart(t *testing.T, start, closeFn func() error, entered <-chan struct{}, release func()) { + t.Helper() + firstResult := make(chan error, 1) + go func() { firstResult <- start() }() + + select { + case <-entered: + case <-time.After(time.Second): + t.Fatal("first Start did not enter initialization") + } + + closeResult := make(chan error, 1) + go func() { closeResult <- closeFn() }() + select { + case err := <-closeResult: + t.Fatalf("Close returned before the in-flight Start completed: %v", err) + case <-time.After(25 * time.Millisecond): + } + + secondResult := make(chan error, 1) + go func() { secondResult <- start() }() + select { + case err := <-secondResult: + t.Fatalf("second Start returned while the previous generation was still initializing: %v", err) + case <-time.After(25 * time.Millisecond): + } + + release() + if err := <-firstResult; !errors.Is(err, errFirstStart) { + t.Fatalf("first Start error = %v, want %v", err, errFirstStart) + } + if err := <-closeResult; err != nil { + t.Fatalf("Close error = %v, want nil", err) + } + if err := <-secondResult; !errors.Is(err, errSecondStart) { + t.Fatalf("waiting second Start error = %v, want %v", err, errSecondStart) + } + if err := start(); !errors.Is(err, errSecondStart) { + t.Fatalf("post-Close Start error = %v, want %v", err, errSecondStart) + } +} + +func TestTCPPipelineCloseLinearizesWithInFlightStart(t *testing.T) { + rpc := &linearizedStartRPC{firstEntered: make(chan struct{}), releaseFirst: make(chan struct{})} + pipeline, err := NewTcpPipeline(rpc, &clientpb.Pipeline{ + Name: "linearized-tcp", Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + assertCloseWaitsForStart(t, pipeline.Start, pipeline.Close, rpc.firstEntered, func() { close(rpc.releaseFirst) }) +} + +func TestHTTPPipelineCloseLinearizesWithInFlightStart(t *testing.T) { + rpc := &linearizedStartRPC{firstEntered: make(chan struct{}), releaseFirst: make(chan struct{})} + pipeline, err := NewHttpPipeline(rpc, &clientpb.Pipeline{ + Name: "linearized-http", Tls: &clientpb.TLS{}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: "{}"}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + assertCloseWaitsForStart(t, pipeline.Start, pipeline.Close, rpc.firstEntered, func() { close(rpc.releaseFirst) }) +} + +func TestTCPPipelineStartsServingAfterForwardAndStateCommit(t *testing.T) { + oldListen, oldCmux := tcpListen, tcpStartCmux + ln := newRollbackListener() + tcpListen = func(string, string) (net.Listener, error) { return ln, nil } + t.Cleanup(func() { tcpListen, tcpStartCmux = oldListen, oldCmux }) + + const listenerID = "committed-listener" + const pipelineID = "committed-tcp" + key := core.PipelineRuntimeKey(listenerID, pipelineID) + _ = core.Forwarders.Remove(key) + t.Cleanup(func() { _ = core.Forwarders.Remove(key) }) + + stream := newCommittedStartStream() + pipeline, err := NewTcpPipeline(&committedStartRPC{stream: stream}, &clientpb.Pipeline{ + Name: pipelineID, ListenerId: listenerID, Tls: &clientpb.TLS{Enable: true}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + tcpStartCmux = func(got net.Listener, _ *tls.Config, _ func(net.Conn), _ core.GoErrorHandler) (net.Listener, error) { + pipeline.stateMu.RLock() + stateCommitted := pipeline.Enable && pipeline.starting && pipeline.ln == got + pipeline.stateMu.RUnlock() + if core.Forwarders.Get(key) == nil || !stateCommitted { + t.Errorf("TCP serving started before forward/state commit: forward=%v stateCommitted=%t", core.Forwarders.Get(key) != nil, stateCommitted) + } + return got, nil + } + + if err := pipeline.Start(); err != nil { + t.Fatalf("Start failed: %v", err) + } +} + +func TestHTTPPipelineStartsServingAfterForwardAndStateCommit(t *testing.T) { + oldListen, oldCmux := httpListen, httpStartWithCmux + ln := newRollbackListener() + httpListen = func(string, string) (net.Listener, error) { return ln, nil } + t.Cleanup(func() { httpListen, httpStartWithCmux = oldListen, oldCmux }) + + const listenerID = "committed-listener" + const pipelineID = "committed-http" + key := core.PipelineRuntimeKey(listenerID, pipelineID) + _ = core.Forwarders.Remove(key) + t.Cleanup(func() { _ = core.Forwarders.Remove(key) }) + + stream := newCommittedStartStream() + pipeline, err := NewHttpPipeline(&committedStartRPC{stream: stream}, &clientpb.Pipeline{ + Name: pipelineID, ListenerId: listenerID, Tls: &clientpb.TLS{Enable: true, Cert: &clientpb.Cert{}}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: "{}"}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + httpStartWithCmux = func(gotPipeline *HTTPPipeline, got net.Listener, _ *http.ServeMux) error { + gotPipeline.stateMu.RLock() + stateCommitted := gotPipeline.Enable && gotPipeline.starting && gotPipeline.srv == got + gotPipeline.stateMu.RUnlock() + if core.Forwarders.Get(key) == nil || !stateCommitted { + t.Errorf("HTTP serving started before forward/state commit: forward=%v stateCommitted=%t", core.Forwarders.Get(key) != nil, stateCommitted) + } + return nil + } + + if err := pipeline.Start(); err != nil { + t.Fatalf("Start failed: %v", err) + } +} + +func TestTCPPipelineServeFailureKeepsStartInFlight(t *testing.T) { + oldListen, oldCmux := tcpListen, tcpStartCmux + ln := newRollbackListener() + tcpListen = func(string, string) (net.Listener, error) { return ln, nil } + t.Cleanup(func() { tcpListen, tcpStartCmux = oldListen, oldCmux }) + + entered := make(chan struct{}) + release := make(chan struct{}) + var calls atomic.Int32 + tcpStartCmux = func(net.Listener, *tls.Config, func(net.Conn), core.GoErrorHandler) (net.Listener, error) { + if calls.Add(1) == 1 { + close(entered) + <-release + return nil, errFirstStart + } + return nil, errSecondStart + } + pipeline, err := NewTcpPipeline(&committedStartRPC{stream: newCommittedStartStream()}, &clientpb.Pipeline{ + Name: "serve-failure-tcp", ListenerId: "serve-failure-listener", Tls: &clientpb.TLS{Enable: true}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{Host: "127.0.0.1"}}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + t.Cleanup(func() { _ = core.Forwarders.Remove(core.PipelineRuntimeKey(pipeline.ListenerID, pipeline.ID())) }) + + assertCloseWaitsForStart(t, pipeline.Start, pipeline.Close, entered, func() { close(release) }) +} + +func TestHTTPPipelineServeFailureKeepsStartInFlight(t *testing.T) { + oldListen, oldCmux := httpListen, httpStartWithCmux + ln := newRollbackListener() + httpListen = func(string, string) (net.Listener, error) { return ln, nil } + t.Cleanup(func() { httpListen, httpStartWithCmux = oldListen, oldCmux }) + + entered := make(chan struct{}) + release := make(chan struct{}) + var calls atomic.Int32 + httpStartWithCmux = func(*HTTPPipeline, net.Listener, *http.ServeMux) error { + if calls.Add(1) == 1 { + close(entered) + <-release + return errFirstStart + } + return errSecondStart + } + pipeline, err := NewHttpPipeline(&committedStartRPC{stream: newCommittedStartStream()}, &clientpb.Pipeline{ + Name: "serve-failure-http", ListenerId: "serve-failure-listener", Tls: &clientpb.TLS{Enable: true, Cert: &clientpb.Cert{}}, Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{Host: "127.0.0.1", Params: "{}"}}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + t.Cleanup(func() { _ = core.Forwarders.Remove(core.PipelineRuntimeKey(pipeline.ListenerID, pipeline.ID())) }) + + assertCloseWaitsForStart(t, pipeline.Start, pipeline.Close, entered, func() { close(release) }) +} + +func TestREMCloseLinearizesWithInFlightStart(t *testing.T) { + oldListen := remConsoleListen + oldClose := remConsoleClose + defer func() { + remConsoleListen = oldListen + remConsoleClose = oldClose + }() + + entered := make(chan struct{}) + release := make(chan struct{}) + var listenCalls atomic.Int32 + var closeCalls atomic.Int32 + var firstListenFinished atomic.Bool + var overlappingClose atomic.Bool + remConsoleListen = func(*remhelper.RemConsole) error { + if listenCalls.Add(1) == 1 { + close(entered) + <-release + firstListenFinished.Store(true) + return errFirstStart + } + return errSecondStart + } + remConsoleClose = func(*remhelper.RemConsole) error { + if !firstListenFinished.Load() { + overlappingClose.Store(true) + } + closeCalls.Add(1) + return nil + } + + pipeline := &REM{Name: "linearized-rem", con: &remhelper.RemConsole{}} + assertCloseWaitsForStart(t, pipeline.Start, pipeline.Close, entered, func() { close(release) }) + if overlappingClose.Load() { + t.Fatal("REM Close overlapped the in-flight Listen") + } + if got := closeCalls.Load(); got != 0 { + t.Fatalf("REM close calls = %d, want 0 after failed starts", got) + } +} + +func TestREMCloseWaitsForSuccessfulListenCleanup(t *testing.T) { + oldListen := remConsoleListen + oldClose := remConsoleClose + defer func() { + remConsoleListen = oldListen + remConsoleClose = oldClose + }() + + entered := make(chan struct{}) + release := make(chan struct{}) + listenFinished := make(chan struct{}) + var closeCalls atomic.Int32 + var overlappingClose atomic.Bool + remConsoleListen = func(*remhelper.RemConsole) error { + close(entered) + <-release + close(listenFinished) + return nil + } + remConsoleClose = func(*remhelper.RemConsole) error { + select { + case <-listenFinished: + default: + overlappingClose.Store(true) + } + closeCalls.Add(1) + return nil + } + + pipeline := &REM{Name: "linearized-rem-success", con: &remhelper.RemConsole{}} + startResult := make(chan error, 1) + go func() { startResult <- pipeline.Start() }() + select { + case <-entered: + case <-time.After(time.Second): + t.Fatal("REM Start did not enter Listen") + } + + closeResult := make(chan error, 1) + go func() { closeResult <- pipeline.Close() }() + select { + case err := <-closeResult: + t.Fatalf("REM Close returned before Listen cleanup: %v", err) + case <-time.After(25 * time.Millisecond): + } + close(release) + if err := <-startResult; err != nil { + t.Fatalf("canceled REM Start error = %v, want nil", err) + } + if err := <-closeResult; err != nil { + t.Fatalf("REM Close error = %v, want nil", err) + } + if overlappingClose.Load() { + t.Fatal("REM Console Close overlapped Listen") + } + if got := closeCalls.Load(); got != 1 { + t.Fatalf("REM Console Close calls = %d, want 1", got) + } + if err := pipeline.Close(); err != nil { + t.Fatalf("repeated REM Close error = %v", err) + } + if got := closeCalls.Load(); got != 1 { + t.Fatalf("REM Console Close calls after repeated Close = %d, want 1", got) + } +} diff --git a/server/listener/listener.go b/server/listener/listener.go index 374751571..006393fe3 100644 --- a/server/listener/listener.go +++ b/server/listener/listener.go @@ -25,10 +25,35 @@ import ( ) var ( - Listener *listener + currentListener *listener + listenerGlobalMu sync.Mutex // ListenerSessions 在 listener 层维护的 Sessions map (rawID -> Session) ) +func setCurrentListener(lns *listener) { + listenerGlobalMu.Lock() + currentListener = lns + listenerGlobalMu.Unlock() +} + +func clearCurrentListener(lns *listener) { + listenerGlobalMu.Lock() + if currentListener == lns { + currentListener = nil + } + listenerGlobalMu.Unlock() +} + +func CloseCurrentListener() error { + listenerGlobalMu.Lock() + lns := currentListener + listenerGlobalMu.Unlock() + if lns == nil { + return nil + } + return lns.Close() +} + var openListenerJobStream = func(client listenerrpc.ListenerRPCClient, ctx context.Context) (listenerrpc.ListenerRPC_JobStreamClient, error) { return client.JobStream(ctx) } @@ -73,7 +98,7 @@ func NewListener(clientConf *mtls.ClientConfig, cfg *configs.ListenerConfig, ser return err } core.GoGuarded("listener-job-stream:"+lns.ID(), lns.Handler, core.LogGuardedError("listener-job-stream:"+lns.ID())) - Listener = lns + setCurrentListener(lns) for _, tcpPipeline := range cfg.TcpPipelines { pipeline, err := tcpPipeline.ToProtobuf(lns.Name) @@ -267,9 +292,7 @@ func (lns *listener) close() error { } } - if Listener == lns { - Listener = nil - } + clearCurrentListener(lns) return errors.Join(errs...) } @@ -526,9 +549,17 @@ func (lns *listener) startPipeline(pipelinepb *clientpb.Pipeline) (core.Pipeline return nil, fmt.Errorf("pipeline is nil") } - // Idempotency: if pipeline already exists locally, treat start as a no-op. + // Keep active pipelines idempotent, but replace a runtime that previously + // failed and marked itself disabled. if existing := lns.pipelines.Get(pipelinepb.Name); existing != nil { - return existing, nil + state := existing.ToProtobuf() + if state != nil && state.Enable { + return existing, nil + } + if err := errors.Join(existing.Close(), lns.cleanupForwardPipelineRuntime(existing.ID())); err != nil { + return nil, fmt.Errorf("cleanup stale pipeline %s: %w", existing.ID(), err) + } + lns.pipelines.Delete(existing.ID()) } var p core.Pipeline @@ -575,10 +606,14 @@ func (lns *listener) cleanupForwardPipelineRuntime(pipelineID string) error { if pipelineID == "" { return nil } + var release func() + var streamErr error if rpc, ok := lns.pipelineRPC.(*forwardPipelineRPC); ok { - rpc.removeStream(lns.ID(), pipelineID) + release, streamErr = rpc.retireStream(lns.ID(), pipelineID) + defer release() } - return core.Forwarders.Remove(core.PipelineRuntimeKey(lns.ID(), pipelineID)) + forwardErr := core.Forwarders.Remove(core.PipelineRuntimeKey(lns.ID(), pipelineID)) + return errors.Join(streamErr, forwardErr) } func (lns *listener) handleStartWebsite(job *clientpb.Job) error { @@ -747,13 +782,16 @@ func (lns *listener) handleStartRem(job *clientpb.Job) error { // Idempotency: REM already started in this listener process. if existing := lns.pipelines.Get(pipe.Name); existing != nil { remPipeline, ok := existing.(*REM) - if ok && remPipeline.Enable { + if ok && remPipeline.enabled() { // Still healthy — just sync its current state. _, err := lns.Rpc.SyncPipeline(lns.Context(), existing.ToProtobuf()) return err } // Dead pipeline (crashed via runtimeErrorHandler) — remove the stale // entry so we can create a fresh one below. + if err := existing.Close(); err != nil { + return fmt.Errorf("cleanup stale REM pipeline %s: %w", existing.ID(), err) + } lns.pipelines.Delete(existing.ID()) } @@ -767,9 +805,9 @@ func (lns *listener) handleStartRem(job *clientpb.Job) error { return err } - _, err = lns.Rpc.SyncPipeline(lns.Context(), rem.ToProtobuf()) - if err != nil { - return err + _, syncErr := lns.Rpc.SyncPipeline(lns.Context(), rem.ToProtobuf()) + if syncErr != nil { + return errors.Join(syncErr, rem.Close()) } lns.pipelines.Add(rem) diff --git a/server/listener/listener_runtime_recovery_test.go b/server/listener/listener_runtime_recovery_test.go new file mode 100644 index 000000000..b6278e3d2 --- /dev/null +++ b/server/listener/listener_runtime_recovery_test.go @@ -0,0 +1,112 @@ +package listener + +import ( + "testing" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + "github.com/chainreactors/malice-network/server/internal/core" +) + +func TestStartPipelineReplacesDisabledRuntime(t *testing.T) { + lns := &listener{ + Name: "runtime-recovery-listener", + pipelines: core.NewPipelines(), + } + existing := NewCustomPipeline(&clientpb.Pipeline{ + Name: "runtime-recovery-pipeline", + ListenerId: lns.Name, + Enable: false, + Body: &clientpb.Pipeline_Custom{ + Custom: &clientpb.CustomPipeline{Name: "runtime-recovery-pipeline", ListenerId: lns.Name}, + }, + }) + lns.pipelines.Add(existing) + + replacement, err := lns.startPipeline(&clientpb.Pipeline{ + Name: existing.ID(), + ListenerId: lns.Name, + Enable: true, + Body: &clientpb.Pipeline_Custom{ + Custom: &clientpb.CustomPipeline{Name: existing.ID(), ListenerId: lns.Name}, + }, + }) + if err != nil { + t.Fatalf("startPipeline failed: %v", err) + } + if replacement == existing { + t.Fatal("startPipeline reused a disabled runtime") + } + if got := lns.pipelines.Get(existing.ID()); got != replacement { + t.Fatalf("registered runtime = %#v, want replacement %#v", got, replacement) + } + if state := replacement.ToProtobuf(); state == nil || !state.Enable { + t.Fatalf("replacement state = %#v, want enabled", state) + } +} + +func TestStartPipelineKeepsActiveRuntimeIdempotent(t *testing.T) { + lns := &listener{ + Name: "runtime-idempotent-listener", + pipelines: core.NewPipelines(), + } + existing := NewCustomPipeline(&clientpb.Pipeline{ + Name: "runtime-idempotent-pipeline", + ListenerId: lns.Name, + Enable: true, + Body: &clientpb.Pipeline_Custom{ + Custom: &clientpb.CustomPipeline{Name: "runtime-idempotent-pipeline", ListenerId: lns.Name}, + }, + }) + lns.pipelines.Add(existing) + + got, err := lns.startPipeline(existing.ToProtobuf()) + if err != nil { + t.Fatalf("startPipeline failed: %v", err) + } + if got != existing { + t.Fatalf("startPipeline runtime = %#v, want existing %#v", got, existing) + } +} + +func TestStartPipelineRemovesDisabledRuntimeForward(t *testing.T) { + lns := &listener{ + Name: "runtime-forward-listener", + pipelines: core.NewPipelines(), + } + existing := NewCustomPipeline(&clientpb.Pipeline{ + Name: "runtime-forward-pipeline", + ListenerId: lns.Name, + Enable: false, + Body: &clientpb.Pipeline_Custom{ + Custom: &clientpb.CustomPipeline{Name: "runtime-forward-pipeline", ListenerId: lns.Name}, + }, + }) + lns.pipelines.Add(existing) + + stream := &rollbackForwardStream{} + forward, err := core.NewForward(&rollbackPipelineRPC{stream: stream}, existing) + if err != nil { + t.Fatalf("create stale forward: %v", err) + } + forward.ListenerId = lns.Name + key := core.PipelineRuntimeKey(lns.Name, existing.ID()) + core.Forwarders.Add(forward) + t.Cleanup(func() { _ = core.Forwarders.Remove(key) }) + + if _, err := lns.startPipeline(&clientpb.Pipeline{ + Name: existing.ID(), + ListenerId: lns.Name, + Enable: true, + Body: &clientpb.Pipeline_Custom{ + Custom: &clientpb.CustomPipeline{Name: existing.ID(), ListenerId: lns.Name}, + }, + }); err != nil { + t.Fatalf("startPipeline failed: %v", err) + } + if got := core.Forwarders.Get(key); got != nil { + t.Fatalf("stale forward remained registered: %#v", got) + } + if got := stream.closeCalls.Load(); got != 1 { + t.Fatalf("stale forward CloseSend calls = %d, want 1", got) + } +} diff --git a/server/listener/pipeline_start_close_audit_test.go b/server/listener/pipeline_start_close_audit_test.go new file mode 100644 index 000000000..0c5e199ec --- /dev/null +++ b/server/listener/pipeline_start_close_audit_test.go @@ -0,0 +1,105 @@ +package listener + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" + "github.com/chainreactors/malice-network/helper/implanttypes" + "github.com/chainreactors/malice-network/server/internal/core" + "google.golang.org/grpc" +) + +var errAuditForwardUnavailable = errors.New("audit forward unavailable") + +type auditUnavailablePipelineRPC struct{} + +func (*auditUnavailablePipelineRPC) OpenForwardStream(context.Context, core.Pipeline) (core.ForwardStream, error) { + return nil, errAuditForwardUnavailable +} +func (*auditUnavailablePipelineRPC) Register(context.Context, *clientpb.RegisterSession, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*auditUnavailablePipelineRPC) Checkin(context.Context, *implantpb.Ping, ...grpc.CallOption) (*clientpb.Empty, error) { + return &clientpb.Empty{}, nil +} +func (*auditUnavailablePipelineRPC) GetArtifact(context.Context, *clientpb.Artifact, ...grpc.CallOption) (*clientpb.Artifact, error) { + return nil, errAuditForwardUnavailable +} + +func runAuditConcurrentStartClose(t *testing.T, startFn, closeFn func() error) { + t.Helper() + + const iterations = 2000 + start := make(chan struct{}) + errCh := make(chan error, 2) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + for range iterations { + if err := startFn(); !errors.Is(err, errAuditForwardUnavailable) { + errCh <- err + return + } + } + }() + go func() { + defer wg.Done() + <-start + for range iterations { + if err := closeFn(); err != nil { + errCh <- err + return + } + } + }() + + close(start) + wg.Wait() + close(errCh) + for err := range errCh { + t.Fatalf("concurrent Start/Close returned an unexpected error: %v", err) + } +} + +func TestAuditTCPPipelineConcurrentStartClose(t *testing.T) { + rpc := &auditUnavailablePipelineRPC{} + pipeline, err := NewTcpPipeline(rpc, &clientpb.Pipeline{ + Name: "audit-tcp-start-close", + ListenerId: "audit-listener", + Tls: &clientpb.TLS{}, + Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Tcp{Tcp: &clientpb.TCPPipeline{ + Host: "127.0.0.1", + Port: 0, + }}, + }) + if err != nil { + t.Fatalf("NewTcpPipeline failed: %v", err) + } + runAuditConcurrentStartClose(t, pipeline.Start, pipeline.Close) +} + +func TestAuditHTTPPipelineConcurrentStartClose(t *testing.T) { + rpc := &auditUnavailablePipelineRPC{} + pipeline, err := NewHttpPipeline(rpc, &clientpb.Pipeline{ + Name: "audit-http-start-close", + ListenerId: "audit-listener", + Tls: &clientpb.TLS{}, + Secure: &clientpb.Secure{}, + Body: &clientpb.Pipeline_Http{Http: &clientpb.HTTPPipeline{ + Host: "127.0.0.1", + Port: 0, + Params: (&implanttypes.PipelineParams{}).String(), + }}, + }) + if err != nil { + t.Fatalf("NewHttpPipeline failed: %v", err) + } + runAuditConcurrentStartClose(t, pipeline.Start, pipeline.Close) +} diff --git a/server/listener/rem.go b/server/listener/rem.go index ef83125ac..65a26a4df 100644 --- a/server/listener/rem.go +++ b/server/listener/rem.go @@ -26,7 +26,38 @@ var remHealthCheck = func(client listenerrpc.ListenerRPCClient, ctx context.Cont return err } -var remSleep = time.Sleep +var remConsoleListen = func(con *rem.RemConsole) error { + return con.Listen(con.ConsoleURL) +} + +var remConsoleClose = func(con *rem.RemConsole) error { + return con.Close() +} + +var remConsoleAccept = func(con *rem.RemConsole) (*agent.Agent, error) { + return con.Accept() +} + +var remConsoleHandler = func(con *rem.RemConsole, ag *agent.Agent) { + con.Handler(ag) +} + +const ( + remHealthCheckTimeout = 10 * time.Second + remAcceptRetryMin = 10 * time.Millisecond + remAcceptRetryMax = time.Second + // remStartCloseTimeout bounds how long Close waits for an in-flight Start to + // abort, so shutdown cannot wedge if the underlying console Listen hangs. + // Console.Listen is normally a fast local bind; this only fires when it is + // stuck on an unresponsive remote link, so it is kept small to keep shutdown + // responsive rather than sized for network latency. + remStartCloseTimeout = 10 * time.Second +) + +type remStartState struct { + done chan struct{} + cleanupErr error +} func NewRem(rpc listenerrpc.ListenerRPCClient, pipeline *clientpb.Pipeline) (*REM, error) { remConfig := pipeline.GetRem() @@ -54,13 +85,19 @@ func NewRem(rpc listenerrpc.ListenerRPCClient, pipeline *clientpb.Pipeline) (*RE } type REM struct { - con *rem.RemConsole - rpc listenerrpc.ListenerRPCClient - remConfig *clientpb.REM - ListenerID string - Name string - Enable bool - CertName string + stateMu sync.RWMutex + starting bool + startState *remStartState + runCtx context.Context + runCancel context.CancelFunc + healthInterval time.Duration + con *rem.RemConsole + rpc listenerrpc.ListenerRPCClient + remConfig *clientpb.REM + ListenerID string + Name string + Enable bool + CertName string *core.PipelineConfig ownAgents sync.Map // agent.ID → struct{}: tracks agents belonging to this pipeline } @@ -70,22 +107,31 @@ func (rem *REM) ID() string { } func (rem *REM) Start() error { - if rem.Enable { + if !rem.beginStart() { return nil } - err := rem.con.Listen(rem.con.ConsoleURL) + err := remConsoleListen(rem.con) if err != nil { + rem.abortStart(nil) return err } - rem.Enable = true + if !rem.enabled() { + cleanupErr := remConsoleClose(rem.con) + rem.abortStart(cleanupErr) + return cleanupErr + } logs.Log.Important(rem.con.Link()) - core.GoGuarded("rem-accept:"+rem.Name, rem.acceptLoop, rem.runtimeErrorHandler("accept loop")) - core.GoGuarded("rem-health:"+rem.Name, rem.healthLoop, rem.runtimeErrorHandler("health loop")) + if !rem.commitStart() { + cleanupErr := remConsoleClose(rem.con) + rem.abortStart(cleanupErr) + return cleanupErr + } return nil } func (rem *REM) ToProtobuf() *clientpb.Pipeline { + enabled := rem.enabled() link := rem.getLink() subscribe := rem.getSubscribe() host := "" @@ -121,7 +167,7 @@ func (rem *REM) ToProtobuf() *clientpb.Pipeline { return &clientpb.Pipeline{ Name: rem.Name, - Enable: rem.Enable, + Enable: enabled, ListenerId: rem.ListenerID, Parser: parserName, Type: consts.RemPipeline, @@ -181,55 +227,166 @@ func (rem *REM) getSubscribe() (subscribe string) { } func (rem *REM) Close() error { + rem.stateMu.Lock() + wasActive := rem.Enable rem.Enable = false - if rem.con == nil { + cancel := rem.runCancel + rem.runCancel = nil + rem.runCtx = nil + if rem.starting { + startState := rem.startState + rem.stateMu.Unlock() + if cancel != nil { + cancel() + } + select { + case <-startState.done: + return startState.cleanupErr + case <-time.After(remStartCloseTimeout): + return fmt.Errorf("rem %s close timed out waiting for startup to abort", rem.Name) + } + } + rem.stateMu.Unlock() + if cancel != nil { + cancel() + } + if !wasActive || rem.con == nil { return nil } - return rem.con.Close() + return remConsoleClose(rem.con) +} + +func (rem *REM) enabled() bool { + rem.stateMu.RLock() + defer rem.stateMu.RUnlock() + return rem.Enable +} + +func (rem *REM) beginStart() bool { + for { + rem.stateMu.Lock() + if rem.starting { + done := rem.startState.done + rem.stateMu.Unlock() + <-done + continue + } + if rem.Enable { + rem.stateMu.Unlock() + return false + } + rem.Enable = true + rem.starting = true + rem.startState = &remStartState{done: make(chan struct{})} + rem.runCtx, rem.runCancel = context.WithCancel(context.Background()) + rem.stateMu.Unlock() + return true + } +} + +func (rem *REM) abortStart(cleanupErr error) { + rem.stateMu.Lock() + rem.starting = false + rem.Enable = false + startState := rem.startState + rem.startState = nil + cancel := rem.runCancel + rem.runCancel = nil + rem.runCtx = nil + if startState != nil { + startState.cleanupErr = cleanupErr + close(startState.done) + } + rem.stateMu.Unlock() + if cancel != nil { + cancel() + } } -func (rem *REM) acceptLoop() error { - for rem.Enable { - ag, err := rem.con.Accept() +func (rem *REM) commitStart() bool { + rem.stateMu.Lock() + defer rem.stateMu.Unlock() + if !rem.Enable { + return false + } + runCtx := rem.runCtx + core.GoGuarded("rem-accept:"+rem.Name, func() error { + return rem.acceptLoopContext(runCtx) + }, rem.runtimeErrorHandler("accept loop")) + core.GoGuarded("rem-health:"+rem.Name, func() error { + return rem.healthLoopContext(runCtx) + }, rem.runtimeErrorHandler("health loop")) + rem.starting = false + startState := rem.startState + rem.startState = nil + close(startState.done) + return true +} + +func (rem *REM) acceptLoopContext(runCtx context.Context) error { + return rem.acceptLoopContextWithBackoff(runCtx, remAcceptRetryMin, remAcceptRetryMax) +} + +func (rem *REM) acceptLoopContextWithBackoff(runCtx context.Context, retryDelay, retryMax time.Duration) error { + if runCtx == nil { + runCtx = context.Background() + } + retryMinimum := retryDelay + for rem.enabled() { + ag, err := remConsoleAccept(rem.con) if err != nil { - if !rem.Enable { + if !rem.enabled() || runCtx.Err() != nil { return nil } // Accept errors are typically transient (timeout, client disconnect). // Log and continue rather than killing the entire pipeline — the next // client reconnect should succeed once the simplex channel is healthy. logs.Log.Errorf("rem %s accept error (will retry): %v", rem.Name, err) + if !waitREMRetry(runCtx, retryDelay) { + return nil + } + retryDelay = nextREMRetryDelay(retryDelay, retryMax) continue } + if !rem.enabled() || runCtx.Err() != nil { + return nil + } + retryDelay = retryMinimum rem.ownAgents.Store(ag.ID, struct{}{}) // Trigger an immediate health check so the new agent's PivotingContext // is created in DB right away instead of waiting for the periodic loop. - if err := remHealthCheck(rem.rpc, context.Background(), rem.ToProtobuf()); err != nil { + if err := rem.healthCheck(runCtx); err != nil { logs.Log.Warnf("rem %s post-accept health check failed: %v", rem.Name, err) } + if !rem.enabled() || runCtx.Err() != nil { + rem.ownAgents.Delete(ag.ID) + return nil + } core.GoGuarded("rem-agent:"+rem.Name, func() error { - rem.con.Handler(ag) - rem.ownAgents.Delete(ag.ID) + rem.handleAgent(ag) return nil }, core.LogGuardedError("rem-agent:"+rem.Name)) } return nil } -func (rem *REM) healthLoop() error { +func (rem *REM) healthLoopContext(runCtx context.Context) error { const ( healthFailureThreshold = 3 opHealthDegraded = "health-check-failed" opHealthRecovered = "health-check-recovered" ) + if runCtx == nil { + runCtx = context.Background() + } consecutiveFailures := 0 unhealthy := false - for rem.Enable { - if err := remHealthCheck(rem.rpc, context.Background(), rem.ToProtobuf()); err != nil { + for rem.enabled() { + if err := rem.healthCheck(runCtx); err != nil { consecutiveFailures++ logs.Log.Errorf("rem %s health check failed (%d/%d): %v", rem.Name, consecutiveFailures, healthFailureThreshold, err) if consecutiveFailures >= healthFailureThreshold && !unhealthy { @@ -258,20 +415,64 @@ func (rem *REM) healthLoop() error { consecutiveFailures = 0 unhealthy = false } - remSleep(30 * time.Second) + interval := rem.healthInterval + if interval <= 0 { + interval = 30 * time.Second + } + timer := time.NewTimer(interval) + select { + case <-timer.C: + case <-runCtx.Done(): + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return nil + } } return nil } +func (rem *REM) healthCheck(runCtx context.Context) error { + if runCtx == nil { + runCtx = context.Background() + } + ctx, cancel := context.WithTimeout(runCtx, remHealthCheckTimeout) + defer cancel() + return remHealthCheck(rem.rpc, ctx, rem.ToProtobuf()) +} + +func nextREMRetryDelay(current, maximum time.Duration) time.Duration { + if current >= maximum || current > maximum/2 { + return maximum + } + return current * 2 +} + +func waitREMRetry(runCtx context.Context, delay time.Duration) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + return true + case <-runCtx.Done(): + return false + } +} + +func (rem *REM) handleAgent(ag *agent.Agent) { + defer rem.ownAgents.Delete(ag.ID) + remConsoleHandler(rem.con, ag) +} + func (rem *REM) runtimeErrorHandler(scope string) core.GoErrorHandler { label := fmt.Sprintf("rem pipeline %s %s", rem.Name, scope) return core.CombineErrorHandlers( core.LogGuardedError(label), func(err error) { - rem.Enable = false - if rem.con != nil { - _ = rem.con.Close() - } + _ = rem.Close() if core.EventBroker != nil { core.EventBroker.Publish(core.Event{ EventType: consts.EventListener, @@ -303,10 +504,10 @@ func (lns *listener) handlerRemAgentCtrl(job *clientpb.Job) error { rem.(*REM).ownAgents.Store(a.ID, struct{}{}) job.Body = &clientpb.Job_RemAgent{ RemAgent: &clientpb.REMAgent{ - Id: a.Name(), + Id: a.Name(), InboundSide: a.InboundSide, - Local: a.LocalURL.String(), - Remote: a.RemoteURL.String(), + Local: a.LocalURL.String(), + Remote: a.RemoteURL.String(), }, } return nil diff --git a/server/listener/rem_runtime_test.go b/server/listener/rem_runtime_test.go index 58e8ec4cb..6b8e66089 100644 --- a/server/listener/rem_runtime_test.go +++ b/server/listener/rem_runtime_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "sync" "testing" "time" @@ -11,8 +12,201 @@ import ( "github.com/chainreactors/IoM-go/proto/services/listenerrpc" remhelper "github.com/chainreactors/malice-network/helper/third/rem" "github.com/chainreactors/malice-network/server/internal/core" + "github.com/chainreactors/rem/agent" ) +func setREMTestRunContext(rem *REM) context.Context { + runCtx, cancel := context.WithCancel(context.Background()) + rem.stateMu.Lock() + rem.runCtx = runCtx + rem.runCancel = cancel + rem.stateMu.Unlock() + return runCtx +} + +func TestREMStartClosePreservesConsoleCleanupError(t *testing.T) { + oldListen := remConsoleListen + oldClose := remConsoleClose + defer func() { + remConsoleListen = oldListen + remConsoleClose = oldClose + }() + + entered := make(chan struct{}) + release := make(chan struct{}) + cleanupErr := errors.New("console cleanup failed") + remConsoleListen = func(*remhelper.RemConsole) error { + close(entered) + <-release + return nil + } + remConsoleClose = func(*remhelper.RemConsole) error { return cleanupErr } + + rem := &REM{Name: "rem-cleanup", con: &remhelper.RemConsole{}} + startResult := make(chan error, 1) + go func() { startResult <- rem.Start() }() + <-entered + + closeResult := make(chan error, 1) + go func() { closeResult <- rem.Close() }() + deadline := time.Now().Add(time.Second) + for rem.enabled() && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if rem.enabled() { + t.Fatal("Close did not cancel the starting generation") + } + close(release) + + if err := <-startResult; !errors.Is(err, cleanupErr) { + t.Fatalf("Start error = %v, want cleanup error", err) + } + if err := <-closeResult; !errors.Is(err, cleanupErr) { + t.Fatalf("Close error = %v, want cleanup error", err) + } +} + +func TestREMCloseCancelsHealthRPCContextWithDeadline(t *testing.T) { + oldHealthCheck := remHealthCheck + defer func() { remHealthCheck = oldHealthCheck }() + + healthEntered := make(chan struct{}) + healthReturned := make(chan error, 1) + releaseHealth := make(chan struct{}) + defer close(releaseHealth) + var enteredOnce sync.Once + remHealthCheck = func(_ listenerrpc.ListenerRPCClient, ctx context.Context, _ *clientpb.Pipeline) error { + if _, ok := ctx.Deadline(); !ok { + healthReturned <- errors.New("health context has no deadline") + return nil + } + enteredOnce.Do(func() { close(healthEntered) }) + select { + case <-ctx.Done(): + healthReturned <- ctx.Err() + case <-releaseHealth: + healthReturned <- errors.New("health check released without cancellation") + } + return nil + } + rem := &REM{Name: "rem-health-cancel", Enable: true, remConfig: &clientpb.REM{}} + runCtx := setREMTestRunContext(rem) + loopDone := make(chan error, 1) + go func() { loopDone <- rem.healthLoopContext(runCtx) }() + select { + case <-healthEntered: + case err := <-healthReturned: + t.Fatalf("health check returned before Close: %v", err) + case <-time.After(time.Second): + t.Fatal("health check did not start") + } + + if err := rem.Close(); err != nil { + t.Fatalf("Close error = %v", err) + } + select { + case err := <-healthReturned: + if !errors.Is(err, context.Canceled) { + t.Fatalf("health context error = %v, want context.Canceled", err) + } + case <-time.After(250 * time.Millisecond): + t.Fatal("Close did not cancel the in-flight health RPC") + } + select { + case err := <-loopDone: + if err != nil { + t.Fatalf("health loop error = %v", err) + } + case <-time.After(250 * time.Millisecond): + t.Fatal("health loop did not stop after Close") + } +} + +func TestREMAcceptErrorsUseInterruptibleBoundedBackoff(t *testing.T) { + oldAccept := remConsoleAccept + + calls := make(chan time.Time, 8) + remConsoleAccept = func(*remhelper.RemConsole) (*agent.Agent, error) { + calls <- time.Now() + return nil, errors.New("temporary accept failure") + } + rem := &REM{ + Name: "rem-accept-backoff", + Enable: true, + } + runCtx := setREMTestRunContext(rem) + done := make(chan error, 1) + go func() { done <- rem.acceptLoopContextWithBackoff(runCtx, 20*time.Millisecond, 30*time.Millisecond) }() + + drained := false + // Stop the accept loop before restoring the seam on every exit path + // (including t.Fatal): a late retry restored to the default Accept would + // dereference the nil console and SIGSEGV the whole package. + defer func() { + _ = rem.Close() + if !drained { + select { + case <-done: + case <-time.After(2 * time.Second): + } + } + remConsoleAccept = oldAccept + }() + + first := <-calls + select { + case second := <-calls: + if delay := second.Sub(first); delay < 15*time.Millisecond { + t.Fatalf("first retry delay = %v, want backoff", delay) + } + case <-time.After(2 * time.Second): + t.Fatal("second Accept did not occur within bounded backoff") + } + select { + case <-calls: + case <-time.After(2 * time.Second): + t.Fatal("third Accept did not occur within maximum backoff") + } + + startedClose := time.Now() + if err := rem.Close(); err != nil { + t.Fatalf("Close error = %v", err) + } + select { + case err := <-done: + drained = true + if err != nil { + t.Fatalf("acceptLoop error = %v", err) + } + if elapsed := time.Since(startedClose); elapsed > time.Second { + t.Fatalf("acceptLoop cancellation took %v", elapsed) + } + case <-time.After(2 * time.Second): + t.Fatal("Close did not interrupt Accept retry backoff") + } +} + +func TestREMAgentOwnershipDeletedWhenHandlerPanics(t *testing.T) { + oldHandler := remConsoleHandler + defer func() { remConsoleHandler = oldHandler }() + remConsoleHandler = func(*remhelper.RemConsole, *agent.Agent) { panic("handler panic") } + + rem := &REM{Name: "rem-handler", con: &remhelper.RemConsole{}} + ag := &agent.Agent{ID: "agent-panic"} + rem.ownAgents.Store(ag.ID, struct{}{}) + err := core.RunGuarded("rem-agent", func() error { + rem.handleAgent(ag) + return nil + }, func(error) {}) + var panicErr *core.PanicError + if !errors.As(err, &panicErr) { + t.Fatalf("handler error = %v, want PanicError", err) + } + if _, ok := rem.ownAgents.Load(ag.ID); ok { + t.Fatal("agent ownership remained after handler panic") + } +} + func TestREMGetLinkFallsBackWhenRuntimePanics(t *testing.T) { rem := &REM{ Name: "rem-a", @@ -27,14 +221,11 @@ func TestREMGetLinkFallsBackWhenRuntimePanics(t *testing.T) { func TestREMHealthLoopPanicBecomesGuardedError(t *testing.T) { oldHealthCheck := remHealthCheck - oldSleep := remSleep remHealthCheck = func(listenerrpc.ListenerRPCClient, context.Context, *clientpb.Pipeline) error { panic("health panic") } - remSleep = func(time.Duration) {} defer func() { remHealthCheck = oldHealthCheck - remSleep = oldSleep }() rem := &REM{ @@ -42,8 +233,12 @@ func TestREMHealthLoopPanicBecomesGuardedError(t *testing.T) { Enable: true, remConfig: &clientpb.REM{}, } + runCtx := setREMTestRunContext(rem) + defer rem.Close() - err := core.RunGuarded("rem-health", rem.healthLoop, func(error) {}) + err := core.RunGuarded("rem-health", func() error { + return rem.healthLoopContext(runCtx) + }, func(error) {}) var panicErr *core.PanicError if !errors.As(err, &panicErr) { t.Fatalf("expected PanicError, got %T", err) @@ -52,12 +247,10 @@ func TestREMHealthLoopPanicBecomesGuardedError(t *testing.T) { func TestREMHealthLoopPublishesDegradedAndRecoveredEvents(t *testing.T) { oldHealthCheck := remHealthCheck - oldSleep := remSleep oldBroker := core.EventBroker oldTicker := core.GlobalTicker defer func() { remHealthCheck = oldHealthCheck - remSleep = oldSleep core.EventBroker = oldBroker core.GlobalTicker = oldTicker }() @@ -68,7 +261,10 @@ func TestREMHealthLoopPublishesDegradedAndRecoveredEvents(t *testing.T) { broker := core.NewBroker() defer broker.Stop() - sub := broker.Subscribe() + sub, err := broker.Subscribe() + if err != nil { + t.Fatalf("Subscribe error = %v", err) + } defer broker.Unsubscribe(sub) readyDeadline := time.After(2 * time.Second) @@ -96,18 +292,18 @@ func TestREMHealthLoopPublishesDegradedAndRecoveredEvents(t *testing.T) { } return nil } - remSleep = func(time.Duration) {} - rem := &REM{ - Name: "rem-health", - Enable: true, - ListenerID: "listener-a", - remConfig: &clientpb.REM{}, + Name: "rem-health", + Enable: true, + ListenerID: "listener-a", + remConfig: &clientpb.REM{}, + healthInterval: time.Millisecond, } done := make(chan error, 1) + runCtx := setREMTestRunContext(rem) go func() { - done <- rem.healthLoop() + done <- rem.healthLoopContext(runCtx) }() deadline := time.After(2 * time.Second) @@ -121,7 +317,7 @@ func TestREMHealthLoopPublishesDegradedAndRecoveredEvents(t *testing.T) { degraded = true case "health-check-recovered": recovered = true - rem.Enable = false + _ = rem.Close() } case err := <-done: if err != nil { @@ -135,7 +331,7 @@ func TestREMHealthLoopPublishesDegradedAndRecoveredEvents(t *testing.T) { } } - rem.Enable = false + _ = rem.Close() select { case err := <-done: if err != nil { diff --git a/server/listener/tcp.go b/server/listener/tcp.go index 650d04c82..420cbb0d1 100644 --- a/server/listener/tcp.go +++ b/server/listener/tcp.go @@ -7,6 +7,7 @@ import ( "fmt" "io" "net" + "sync" "github.com/chainreactors/IoM-go/consts" "github.com/chainreactors/IoM-go/proto/client/clientpb" @@ -20,6 +21,11 @@ import ( cryptostream "github.com/chainreactors/malice-network/server/internal/stream" ) +var ( + tcpListen = net.Listen + tcpStartCmux = StartCmuxTCPListener +) + func NewTcpPipeline(rpc pipelineRPCClient, pipeline *clientpb.Pipeline) (*TCPPipeline, error) { tcp := pipeline.GetTcp() @@ -36,22 +42,28 @@ func NewTcpPipeline(rpc pipelineRPCClient, pipeline *clientpb.Pipeline) (*TCPPip } type TCPPipeline struct { - ln net.Listener - rpc pipelineRPCClient - Name string - Port uint16 - Host string - Enable bool - Target []string - CertName string - parser *parser.MessageParser + stateMu sync.RWMutex + starting bool + startDone chan struct{} + ln net.Listener + rpc pipelineRPCClient + Name string + Port uint16 + Host string + Enable bool + Target []string + CertName string + parser *parser.MessageParser *core.PipelineConfig } func (pipeline *TCPPipeline) ToProtobuf() *clientpb.Pipeline { + pipeline.stateMu.RLock() + enabled := pipeline.Enable + pipeline.stateMu.RUnlock() p := &clientpb.Pipeline{ Name: pipeline.Name, - Enable: pipeline.Enable, + Enable: enabled, Type: consts.TCPPipeline, ListenerId: pipeline.ListenerID, Parser: pipeline.Parser, @@ -76,75 +88,154 @@ func (pipeline *TCPPipeline) ID() string { } func (pipeline *TCPPipeline) Close() error { + pipeline.stateMu.Lock() pipeline.Enable = false - if pipeline.ln == nil { + if pipeline.starting { + done := pipeline.startDone + pipeline.stateMu.Unlock() + <-done return nil } ln := pipeline.ln pipeline.ln = nil - err := ln.Close() - if err != nil && !errors.Is(err, net.ErrClosed) { - return err + pipeline.stateMu.Unlock() + if ln == nil { + return nil } - return nil + return closePipelineListener(ln) } -func (pipeline *TCPPipeline) Start() error { - if pipeline.Enable { +func (pipeline *TCPPipeline) Start() (err error) { + if !pipeline.beginStart() { return nil } + forward, err := core.NewForward(pipeline.rpc, pipeline) if err != nil { + pipeline.abortStart() return err } forward.ListenerId = pipeline.ListenerID - core.Forwarders.Add(forward) - core.GoGuarded("tcp-forward-recv:"+pipeline.Name, func() error { - for { - msg, err := forward.Stream.Recv() - if err != nil { - if !pipeline.Enable { - return nil - } - return fmt.Errorf("tcp pipeline %s forward recv: %w", pipeline.Name, err) + committed := false + registered := false + defer func() { + if !committed { + if registered { + err = errors.Join(err, pipeline.rollbackCommittedStart(forward)) + } else { + err = errors.Join(err, forward.Abort()) } - dispatchForwardTaskRequest("tcp", pipeline.Name, msg) + pipeline.abortStart() } - }, pipeline.runtimeErrorHandler("forward recv loop")) + }() - pipeline.ln, err = pipeline.handler() + ln, tlsConfig, err := pipeline.prepareListener() if err != nil { return err } + if !pipeline.commitStart(ln, forward) { + return closePipelineListener(ln) + } + registered = true + if err := pipeline.startServing(ln, tlsConfig); err != nil { + return err + } + if !pipeline.finishStart() { + return nil + } + committed = true + pipeline.startForwardRecv(forward) logs.Log.Infof("pipeline.tcp - start host=%s port=%d parser=%s tls=%t", pipeline.Host, pipeline.Port, pipeline.Parser, pipeline.TLSConfig.Enable) - pipeline.Enable = true return nil } -func (pipeline *TCPPipeline) handler() (net.Listener, error) { - ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", pipeline.Host, pipeline.Port)) - if err != nil { - return nil, err +func (pipeline *TCPPipeline) enabled() bool { + pipeline.stateMu.RLock() + defer pipeline.stateMu.RUnlock() + return pipeline.Enable +} + +func (pipeline *TCPPipeline) beginStart() bool { + for { + pipeline.stateMu.Lock() + if pipeline.starting { + done := pipeline.startDone + pipeline.stateMu.Unlock() + <-done + continue + } + if pipeline.Enable { + pipeline.stateMu.Unlock() + return false + } + pipeline.Enable = true + pipeline.starting = true + pipeline.startDone = make(chan struct{}) + pipeline.stateMu.Unlock() + return true } +} - // 如果启用了 TLS，使用 cmux 实现 TLS 和非 TLS 的端口复用 - if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable { - return pipeline.handleWithCmux(ln) +func (pipeline *TCPPipeline) abortStart() { + pipeline.stateMu.Lock() + pipeline.starting = false + pipeline.Enable = false + done := pipeline.startDone + pipeline.startDone = nil + if done != nil { + close(done) } + pipeline.stateMu.Unlock() +} - // 非 TLS 模式，使用原有逻辑 - core.GoGuarded("tcp-accept:"+pipeline.Name, func() error { - return pipeline.startAcceptLoop(ln, "tcp pipeline") - }, pipeline.runtimeErrorHandler("accept loop")) - return ln, nil +func (pipeline *TCPPipeline) commitStart(ln net.Listener, forward *core.Forward) bool { + pipeline.stateMu.Lock() + defer pipeline.stateMu.Unlock() + if !pipeline.Enable { + return false + } + pipeline.ln = ln + core.Forwarders.Add(forward) + return true } -// handleWithCmux 使用 cmux 实现 TLS 和非 TLS 的端口复用 -func (pipeline *TCPPipeline) handleWithCmux(ln net.Listener) (net.Listener, error) { +func (pipeline *TCPPipeline) finishStart() bool { + pipeline.stateMu.Lock() + defer pipeline.stateMu.Unlock() + if !pipeline.Enable { + return false + } + pipeline.starting = false + done := pipeline.startDone + pipeline.startDone = nil + close(done) + return true +} + +func (pipeline *TCPPipeline) startForwardRecv(forward *core.Forward) { + core.GoGuarded("tcp-forward-recv:"+pipeline.Name, func() error { + for { + msg, err := forward.Stream.Recv() + if err != nil { + if !pipeline.enabled() { + return nil + } + return fmt.Errorf("tcp pipeline %s forward recv: %w", pipeline.Name, err) + } + dispatchForwardTaskRequest("tcp", pipeline.Name, msg) + } + }, pipeline.runtimeErrorHandler("forward recv loop")) +} + +func (pipeline *TCPPipeline) prepareListener() (net.Listener, *tls.Config, error) { + ln, err := tcpListen("tcp", fmt.Sprintf("%s:%d", pipeline.Host, pipeline.Port)) + if err != nil { + return nil, nil, err + } + var tlsConfig *tls.Config - if pipeline.TLSConfig.Cert != nil { - var err error + if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable && pipeline.TLSConfig.Cert != nil { if pipeline.TLSConfig.MTLS && pipeline.TLSConfig.CA != nil { tlsConfig, err = certutils.GetMTlsConfig(pipeline.TLSConfig.Cert, pipeline.TLSConfig.CA) logs.Log.Infof("pipeline.tcp - mtls_enabled pipeline=%s", pipeline.Name) @@ -152,11 +243,42 @@ func (pipeline *TCPPipeline) handleWithCmux(ln net.Listener) (net.Listener, erro tlsConfig, err = certutils.GetTlsConfig(pipeline.TLSConfig.Cert) } if err != nil { - return nil, err + return nil, nil, errors.Join(err, closePipelineListener(ln)) } } + return ln, tlsConfig, nil +} + +func (pipeline *TCPPipeline) startServing(ln net.Listener, tlsConfig *tls.Config) error { + if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable { + _, err := tcpStartCmux(ln, tlsConfig, pipeline.HandleConnection, pipeline.runtimeErrorHandler("cmux")) + return err + } + core.GoGuarded("tcp-accept:"+pipeline.Name, func() error { + return pipeline.startAcceptLoop(ln, "tcp pipeline") + }, pipeline.runtimeErrorHandler("accept loop")) + return nil +} - return StartCmuxTCPListener(ln, tlsConfig, pipeline.HandleConnection, pipeline.runtimeErrorHandler("cmux")) +func closePipelineListener(ln net.Listener) error { + if ln == nil { + return nil + } + err := ln.Close() + if errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + +func (pipeline *TCPPipeline) rollbackCommittedStart(forward *core.Forward) error { + forwardErr := core.Forwarders.RemoveIfSame(forward.RuntimeKey(), forward) + pipeline.stateMu.Lock() + ln := pipeline.ln + pipeline.ln = nil + pipeline.Enable = false + pipeline.stateMu.Unlock() + return errors.Join(forwardErr, closePipelineListener(ln)) } // startAcceptLoop 启动连接接受循环 (用于非 cmux 模式) @@ -165,7 +287,7 @@ func (pipeline *TCPPipeline) startAcceptLoop(ln net.Listener, logPrefix string) for { conn, err := ln.Accept() if err != nil { - if !pipeline.Enable || errors.Is(err, net.ErrClosed) { + if !pipeline.enabled() || errors.Is(err, net.ErrClosed) { logs.Log.Importantf("%s already disable, break accept", ln.Addr().String()) return nil } @@ -253,10 +375,7 @@ func (pipeline *TCPPipeline) runtimeErrorHandler(scope string) core.GoErrorHandl return core.CombineErrorHandlers( core.LogGuardedError(label), func(err error) { - pipeline.Enable = false - if pipeline.ln != nil { - _ = pipeline.ln.Close() - } + _ = pipeline.Close() if core.EventBroker != nil { core.EventBroker.Publish(core.Event{ EventType: consts.EventListener, diff --git a/server/rpc/generic.go b/server/rpc/generic.go index b1b026cef..4d788b435 100644 --- a/server/rpc/generic.go +++ b/server/rpc/generic.go @@ -322,14 +322,15 @@ func (rpc *Server) GenericInternal(ctx context.Context, req proto.Message, expec } func loadPipelineStreamForSession(sess *core.Session) (interface{}, bool) { - if sess == nil || sess.PipelineID == "" { + listenerID, pipelineID := sess.RoutingSnapshot() + if pipelineID == "" { return nil, false } - if streamVal, ok := pipelinesCh.Load(core.PipelineRuntimeKey(sess.ListenerID, sess.PipelineID)); ok { + if streamVal, ok := pipelinesCh.Load(core.PipelineRuntimeKey(listenerID, pipelineID)); ok { return streamVal, true } - if sess.ListenerID != "" && runtimePipelineNameCount(sess.PipelineID) <= 1 { - return pipelinesCh.Load(sess.PipelineID) + if listenerID != "" && runtimePipelineNameCount(pipelineID) <= 1 { + return pipelinesCh.Load(pipelineID) } return nil, false } diff --git a/server/rpc/generic_runtime_test.go b/server/rpc/generic_runtime_test.go index 5be2bcbd1..65e6855a4 100644 --- a/server/rpc/generic_runtime_test.go +++ b/server/rpc/generic_runtime_test.go @@ -34,16 +34,19 @@ func waitEventBrokerReady(t testing.TB, broker interface{ TryPublish(core.Event) } func subscribeEventBrokerReady(t testing.TB, broker interface { - Subscribe() chan core.Event + Subscribe() (chan core.Event, error) Unsubscribe(chan core.Event) TryPublish(core.Event) error }) chan core.Event { t.Helper() - sub := broker.Subscribe() + sub, err := broker.Subscribe() + if err != nil { + t.Fatalf("Subscribe error = %v", err) + } readyOp := "subscriber-ready" deadline := time.After(2 * time.Second) for { - err := broker.TryPublish(core.Event{EventType: "test", Op: readyOp}) + err = broker.TryPublish(core.Event{EventType: "test", Op: readyOp}) if err != nil && !errors.Is(err, core.ErrEventBrokerQueueFull) { t.Fatalf("publish subscriber readiness event: %v", err) } diff --git a/server/rpc/grpc.go b/server/rpc/grpc.go index 37d0b98b5..df69fea5f 100644 --- a/server/rpc/grpc.go +++ b/server/rpc/grpc.go @@ -162,24 +162,3 @@ func NewServer() *Server { // todo event return &Server{} } - -//func (rpc *Server) genericHandler(ctx context.Context, req *GenericRequest) (proto.Message, error) { -// spite, err := req.NewSpite(req.Message) -// if err != nil { -// logs.Log.Errorf(err.Error()) -// return nil, err -// } -// data, err := req.Session.RequestAndWait( -// &clientpb.SpiteSession{SessionId: req.Session.ID, TaskId: req.Task.Id, Spite: spite}, -// pipelinesCh[req.Session.PipelineID], -// consts.MinTimeout) -// if err != nil { -// return nil, err -// } -// req.Session.DeleteResp(req.Task.Id) -// resp, err := types.ParseSpite(data) -// if err != nil { -// return nil, err -// } -// return resp, nil -//} diff --git a/server/rpc/rpc-addon.go b/server/rpc/rpc-addon.go index f1b4d00c4..abe3cbd33 100644 --- a/server/rpc/rpc-addon.go +++ b/server/rpc/rpc-addon.go @@ -3,6 +3,7 @@ package rpc import ( "context" "errors" + "github.com/chainreactors/IoM-go/consts" "github.com/chainreactors/IoM-go/proto/client/clientpb" implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" @@ -34,23 +35,9 @@ func applyAddonsResponse(sess *core.Session, spite *implantpb.Spite, replace boo return } if replace { - sess.Addons = nil - } - seen := make(map[string]bool, len(sess.Addons)) - for _, a := range sess.Addons { - if a != nil && a.GetName() != "" { - seen[a.GetName()] = true - } - } - for _, a := range exts.GetAddons() { - if a == nil || a.GetName() == "" { - continue - } - if seen[a.GetName()] { - continue - } - seen[a.GetName()] = true - sess.Addons = append(sess.Addons, a) + sess.ReplaceAddons(exts.GetAddons()) + } else { + sess.MergeAddons(exts.GetAddons()) } sess.SaveAndNotify("") } @@ -61,19 +48,11 @@ func applyAddonLoad(sess *core.Session, req *implantpb.LoadAddon) { if req == nil || req.GetName() == "" { return } - for _, a := range sess.Addons { - if a != nil && a.GetName() == req.GetName() { - a.Type = req.GetType() - a.Depend = req.GetDepend() - sess.SaveAndNotify("") - return - } - } - sess.Addons = append(sess.Addons, &implantpb.Addon{ + sess.MergeAddons([]*implantpb.Addon{{ Name: req.GetName(), Depend: req.GetDepend(), Type: req.GetType(), - }) + }}) sess.SaveAndNotify("") } @@ -82,14 +61,7 @@ func (rpc *Server) ExecuteAddon(ctx context.Context, req *implantpb.ExecuteAddon return nil, types.ErrMissingRequestField } if session, err := getSession(ctx); err == nil { - hasAddon := false - for _, addon := range session.Addons { - if addon.Name == req.Addon { - hasAddon = true - break - } - } - if !hasAddon { + if !session.HasAddon(req.Addon) { return nil, errors.New("addon not found, please load_addon first") } } diff --git a/server/rpc/rpc-events.go b/server/rpc/rpc-events.go index 29a1abcb5..008f6fc1f 100644 --- a/server/rpc/rpc-events.go +++ b/server/rpc/rpc-events.go @@ -14,23 +14,50 @@ import ( ) func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsServer) error { - events := core.EventBroker.Subscribe() - if err := stream.SendHeader(metadata.Pairs(consts.EventStreamReadyHeader, "1")); err != nil { - core.EventBroker.Unsubscribe(events) + broker := core.EventBroker + if broker == nil { + return core.ErrEventBrokerUnavailable + } + requestMetadata, _ := metadata.FromIncomingContext(stream.Context()) + replayHistory := len(requestMetadata.Get(consts.EventStreamHistoryReplayHeader)) > 0 + var events chan core.Event + var history []core.Event + var err error + if replayHistory { + events, history, err = broker.SubscribeWithHistory() + } else { + events, err = broker.Subscribe() + } + if err != nil { return err } clientID := core.GetCurrentID() defer func() { logs.Log.Infof("client: %d disconnected", clientID) core.Clients.Remove(int(clientID)) - core.EventBroker.Unsubscribe(events) + broker.Unsubscribe(events) }() + responseHeader := metadata.Pairs(consts.EventStreamReadyHeader, "true") + if replayHistory { + responseHeader.Set(consts.EventStreamHistoryCountHeader, strconv.Itoa(len(history))) + } + if err := stream.SendHeader(responseHeader); err != nil { + return err + } + for _, event := range history { + if err := stream.Send(event.ToProtobuf()); err != nil { + return err + } + } for { select { case <-stream.Context().Done(): return nil - case event := <-events: + case event, ok := <-events: + if !ok { + return nil + } pb := event.ToProtobuf() err := stream.Send(pb) if err != nil { @@ -42,22 +69,29 @@ func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsSe } func (rpc *Server) EventsV2(req *clientpb.EventSubscription, stream clientrpc.MaliceRPC_EventsV2Server) error { - events := core.EventBroker.SubscribeV2(core.EventSubscription{ + broker := core.EventBroker + if broker == nil { + return core.ErrEventBrokerUnavailable + } + events, err := broker.SubscribeV2(core.EventSubscription{ StreamID: req.GetStreamId(), AfterSequence: req.GetAfterSequence(), Replay: req.GetReplay(), Topics: req.GetTopics(), IncludeHeartbeats: req.GetIncludeHeartbeats(), }) + if err != nil { + return err + } if err := stream.SendHeader(metadata.Pairs(consts.EventStreamReadyHeader, "2")); err != nil { - core.EventBroker.UnsubscribeV2(events) + broker.UnsubscribeV2(events) return err } clientID := core.GetCurrentID() defer func() { logs.Log.Infof("client: %d disconnected", clientID) core.Clients.Remove(int(clientID)) - core.EventBroker.UnsubscribeV2(events) + broker.UnsubscribeV2(events) }() for { diff --git a/server/rpc/rpc-filesystem.go b/server/rpc/rpc-filesystem.go index f04263720..c738c080e 100644 --- a/server/rpc/rpc-filesystem.go +++ b/server/rpc/rpc-filesystem.go @@ -20,7 +20,7 @@ func (rpc *Server) Ls(ctx context.Context, req *implantpb.Request) (*clientpb.Ta func (rpc *Server) Cd(ctx context.Context, req *implantpb.Request) (*clientpb.Task, error) { return rpc.AssertAndHandleWithSession(ctx, req, consts.ModuleCd, types.MsgResponse, func(greq *GenericRequest, spite *implantpb.Spite) { if output := spite.GetResponse().GetOutput(); output != "" { - greq.Session.WorkDir = output + greq.Session.SetWorkDir(output) _ = greq.Session.SaveAndNotify("") } }) diff --git a/server/rpc/rpc-forward-listener.go b/server/rpc/rpc-forward-listener.go index 27cdff3c6..856cb15a0 100644 --- a/server/rpc/rpc-forward-listener.go +++ b/server/rpc/rpc-forward-listener.go @@ -475,7 +475,8 @@ func deliverForwardTaskResponse(event *clientpb.SpiteRequest) error { if err := sess.Save(); err != nil { logs.Log.Errorf("save session %s reborn state failed: %s", sess.ID, err.Error()) } - sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true) + target, _, pipelineID := sess.ConnectionSnapshot() + sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), target, pipelineID), true, true) } taskID := event.Spite.TaskId if taskID == 0 && event.Task != nil { diff --git a/server/rpc/rpc-implant.go b/server/rpc/rpc-implant.go index 607d26ec9..f6dedc648 100644 --- a/server/rpc/rpc-implant.go +++ b/server/rpc/rpc-implant.go @@ -42,8 +42,9 @@ func (rpc *Server) Register(ctx context.Context, req *clientpb.RegisterSession) if err != nil { return nil, err } - sess.Publish(consts.CtrlSessionRegister, fmt.Sprintf("new session %s from %s start at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true) - logs.Log.Importantf("new session %s from %s", sess.ID, sess.PipelineID) + target, _, pipelineID := sess.ConnectionSnapshot() + sess.Publish(consts.CtrlSessionRegister, fmt.Sprintf("new session %s from %s start at %s", sess.Abstract(), target, pipelineID), true, true) + logs.Log.Importantf("new session %s from %s", sess.ID, pipelineID) core.Sessions.Add(sess) } else { if err := validateReRegisterPipeline(req); err != nil { @@ -52,7 +53,8 @@ func (rpc *Server) Register(ctx context.Context, req *clientpb.RegisterSession) logs.Log.Infof("session %s re-register", sess.ID) sess.SetLastCheckin(getTimestamp(ctx)) sess.Update(req) - sess.Publish(consts.CtrlSessionUpdate, fmt.Sprintf("%s from %s re-registered at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true) + target, _, pipelineID := sess.ConnectionSnapshot() + sess.Publish(consts.CtrlSessionUpdate, fmt.Sprintf("%s from %s re-registered at %s", sess.Abstract(), target, pipelineID), true, true) core.Sessions.Add(sess) } @@ -147,7 +149,8 @@ func (rpc *Server) Checkin(ctx context.Context, req *implantpb.Ping) (*clientpb. logs.Log.Errorf("save session %s checkin failed: %s", sess.ID, err.Error()) } if reborn { - sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true) + target, _, pipelineID := sess.ConnectionSnapshot() + sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), target, pipelineID), true, true) } sess.Publish(consts.CtrlSessionCheckin, "", false, false) @@ -179,8 +182,7 @@ func (rpc *Server) Sleep(ctx context.Context, req *implantpb.Timer) (*clientpb.T greq.HandlerResponse(ch, types.MsgEmpty) if session, err := getSession(ctx); err == nil { - session.Jitter = req.Jitter - session.Expression = req.Expression + session.SetTimer(req.Expression, req.Jitter) if err := session.SaveAndNotify(""); err != nil { return nil, err } @@ -330,7 +332,8 @@ func bindPollingRunning(sessionID string) bool { func sendBindPing(sess *core.Session) error { streamVal, ok := loadPipelineStreamForSession(sess) if !ok || streamVal == nil { - return fmt.Errorf("bind pipeline %s unavailable for session %s", sess.PipelineID, sess.ID) + _, pipelineID := sess.RoutingSnapshot() + return fmt.Errorf("bind pipeline %s unavailable for session %s", pipelineID, sess.ID) } return sess.Request( &clientpb.SpiteRequest{Session: sess.ToProtobufLite(), Task: nil, Spite: types.BuildPingSpite()}, @@ -547,7 +550,7 @@ func (rpc *Server) triggerKeyExchange(ctx context.Context, sess *core.Session) e return } sess.SecureManager.ResetCounters() - sess.UpdateKeyPairFieldsAndPushCtrl(resp.PublicKey, keyPair.Private) + sess.UpdateKeyPair(resp.PublicKey, keyPair.Private) }) return err } diff --git a/server/rpc/rpc-listener.go b/server/rpc/rpc-listener.go index f149c4e7d..c80efe1aa 100644 --- a/server/rpc/rpc-listener.go +++ b/server/rpc/rpc-listener.go @@ -121,7 +121,8 @@ func (rpc *Server) SpiteStream(stream listenerrpc.ListenerRPC_SpiteStreamServer) if err := sess.Save(); err != nil { logs.Log.Errorf("save session %s reborn state failed: %s", sess.ID, err.Error()) } - sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), sess.Target, sess.PipelineID), true, true) + target, _, pipelineID := sess.ConnectionSnapshot() + sess.Publish(consts.CtrlSessionReborn, fmt.Sprintf("session %s from %s reborn at %s", sess.Abstract(), target, pipelineID), true, true) } if size := proto.Size(msg.Spite); size <= 1000 { diff --git a/server/rpc/rpc-module.go b/server/rpc/rpc-module.go index 0ef0f8045..5189f5937 100644 --- a/server/rpc/rpc-module.go +++ b/server/rpc/rpc-module.go @@ -19,20 +19,7 @@ func applyModulesResponse(sess *core.Session, spite *implantpb.Spite, appendOnly if modules == nil { return } - if appendOnly { - sess.Modules = append(sess.Modules, modules.Modules...) - if bundleMap := modules.GetBundleMap(); len(bundleMap) > 0 { - if sess.BundleMap == nil { - sess.BundleMap = make(map[string]string) - } - for k, v := range bundleMap { - sess.BundleMap[k] = v - } - } - } else { - sess.Modules = modules.Modules - sess.BundleMap = modules.GetBundleMap() - } + sess.ApplyModules(modules, appendOnly) sess.SaveAndNotify("") } diff --git a/server/rpc/rpc-session.go b/server/rpc/rpc-session.go index 4b5ce75f1..f5f9cf2ab 100644 --- a/server/rpc/rpc-session.go +++ b/server/rpc/rpc-session.go @@ -96,7 +96,7 @@ func (rpc *Server) SessionManage(ctx context.Context, req *clientpb.BasicUpdateS case "note": session, err := core.Sessions.Get(req.SessionId) if err == nil { - session.Note = req.Arg + session.SetNote(req.Arg) if err = session.SaveAndNotify(fmt.Sprintf("session %s note updated to %s", req.SessionId, req.Arg)); err != nil { return nil, err } @@ -109,7 +109,7 @@ func (rpc *Server) SessionManage(ctx context.Context, req *clientpb.BasicUpdateS case "group": session, err := core.Sessions.Get(req.SessionId) if err == nil { - session.Group = req.Arg + session.SetGroup(req.Arg) if err = session.SaveAndNotify(fmt.Sprintf("session %s group updated to %s", req.SessionId, req.Arg)); err != nil { return nil, err } diff --git a/server/rpc/rpc_addon_runtime_test.go b/server/rpc/rpc_addon_runtime_test.go index 37b76d621..633971d7b 100644 --- a/server/rpc/rpc_addon_runtime_test.go +++ b/server/rpc/rpc_addon_runtime_test.go @@ -26,9 +26,9 @@ func TestAddonHandlersRejectNilRequest(t *testing.T) { func TestApplyAddonsResponseDeduplicatesAndPersistsAddons(t *testing.T) { env := newRPCTestEnv(t) sess := env.seedSession(t, "rpc-addon-list", "rpc-addon-list-pipe", true) - sess.Addons = []*implantpb.Addon{ + sess.ReplaceAddons([]*implantpb.Addon{ {Name: "seatbelt", Type: "bof", Depend: "exec"}, - } + }) applyAddonsResponse(sess, &implantpb.Spite{ Body: &implantpb.Spite_Addons{ @@ -42,13 +42,14 @@ func TestApplyAddonsResponseDeduplicatesAndPersistsAddons(t *testing.T) { }, }, false) - if len(sess.Addons) != 2 { - t.Fatalf("runtime addons = %#v, want 2 unique addons", sess.Addons) + addons := sess.AddonsSnapshot() + if len(addons) != 2 { + t.Fatalf("runtime addons = %#v, want 2 unique addons", addons) } - if addon := findAddon(sess.Addons, "seatbelt"); addon == nil || addon.GetType() != "bof" { + if addon := findAddon(addons, "seatbelt"); addon == nil || addon.GetType() != "bof" { t.Fatalf("runtime seatbelt addon = %#v, want bof addon", addon) } - if addon := findAddon(sess.Addons, "sharpview"); addon == nil || addon.GetType() != "assembly" { + if addon := findAddon(addons, "sharpview"); addon == nil || addon.GetType() != "assembly" { t.Fatalf("runtime sharpview addon = %#v, want assembly addon", addon) } @@ -64,18 +65,19 @@ func TestApplyAddonsResponseDeduplicatesAndPersistsAddons(t *testing.T) { func TestApplyAddonLoadDeduplicatesRepeatedLoadsAndRefreshesMetadata(t *testing.T) { env := newRPCTestEnv(t) sess := env.seedSession(t, "rpc-addon-load", "rpc-addon-load-pipe", true) - sess.Addons = []*implantpb.Addon{ + sess.ReplaceAddons([]*implantpb.Addon{ {Name: "seatbelt", Type: "bof", Depend: "exec"}, - } + }) applyAddonLoad(sess, &implantpb.LoadAddon{Name: "seatbelt", Type: "assembly", Depend: "execute_dll"}) applyAddonLoad(sess, &implantpb.LoadAddon{Name: "sharpview", Type: "assembly", Depend: "exec"}) applyAddonLoad(sess, &implantpb.LoadAddon{Name: "sharpview", Type: "assembly", Depend: "exec"}) - if len(sess.Addons) != 2 { - t.Fatalf("runtime addons = %#v, want 2 unique addons", sess.Addons) + addons := sess.AddonsSnapshot() + if len(addons) != 2 { + t.Fatalf("runtime addons = %#v, want 2 unique addons", addons) } - seatbelt := findAddon(sess.Addons, "seatbelt") + seatbelt := findAddon(addons, "seatbelt") if seatbelt == nil || seatbelt.GetType() != "assembly" || seatbelt.GetDepend() != "execute_dll" { t.Fatalf("runtime seatbelt addon = %#v, want refreshed metadata", seatbelt) } diff --git a/server/rpc/rpc_events_runtime_test.go b/server/rpc/rpc_events_runtime_test.go new file mode 100644 index 000000000..844681777 --- /dev/null +++ b/server/rpc/rpc_events_runtime_test.go @@ -0,0 +1,299 @@ +package rpc + +import ( + "context" + "errors" + "io" + "strconv" + "sync" + "testing" + "time" + + "github.com/chainreactors/IoM-go/consts" + "github.com/chainreactors/IoM-go/proto/client/clientpb" + "github.com/chainreactors/malice-network/server/internal/core" + "google.golang.org/grpc/metadata" +) + +type captureEventsStream struct { + ctx context.Context + headerMu sync.Mutex + header metadata.MD + headerReady chan struct{} + headerOnce sync.Once + events chan *clientpb.Event + sendGate <-chan struct{} + sendStarted chan struct{} + sendOnce sync.Once +} + +func newCaptureEventsStream(ctx context.Context) *captureEventsStream { + return &captureEventsStream{ + ctx: ctx, + headerReady: make(chan struct{}), + events: make(chan *clientpb.Event, 4), + } +} + +func (s *captureEventsStream) Send(event *clientpb.Event) error { + if s.sendGate != nil { + s.sendOnce.Do(func() { close(s.sendStarted) }) + select { + case <-s.sendGate: + case <-s.ctx.Done(): + return s.ctx.Err() + } + } + select { + case s.events <- event: + return nil + case <-s.ctx.Done(): + return s.ctx.Err() + } +} + +func (s *captureEventsStream) SetHeader(metadata.MD) error { return nil } + +func (s *captureEventsStream) SendHeader(header metadata.MD) error { + s.headerMu.Lock() + s.header = header.Copy() + s.headerMu.Unlock() + s.headerOnce.Do(func() { close(s.headerReady) }) + return nil +} + +func (s *captureEventsStream) SetTrailer(metadata.MD) {} +func (s *captureEventsStream) Context() context.Context { return s.ctx } +func (s *captureEventsStream) SendMsg(interface{}) error { return nil } +func (s *captureEventsStream) RecvMsg(interface{}) error { return io.EOF } + +func (s *captureEventsStream) responseHeader() metadata.MD { + s.headerMu.Lock() + defer s.headerMu.Unlock() + return s.header.Copy() +} + +func newEventTestBroker(t *testing.T) interface { + TryPublish(core.Event) error + GetAll() []*core.Event + Stop() +} { + t.Helper() + oldBroker := core.EventBroker + broker := core.NewBroker() + waitEventBrokerReady(t, broker) + t.Cleanup(func() { + broker.Stop() + core.EventBroker = oldBroker + }) + return broker +} + +func waitHistorySize(t *testing.T, broker interface{ GetAll() []*core.Event }, size int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for len(broker.GetAll()) != size { + if time.Now().After(deadline) { + t.Fatalf("event history size = %d, want %d", len(broker.GetAll()), size) + } + time.Sleep(time.Millisecond) + } +} + +func TestEventsReplaysHistoryWhenNegotiated(t *testing.T) { + broker := newEventTestBroker(t) + if err := broker.TryPublish(core.Event{ + EventType: consts.EventBroadcast, + Op: "history", + Important: true, + }); err != nil { + t.Fatalf("publish history: %v", err) + } + waitHistorySize(t, broker, 1) + + baseContext, cancel := context.WithCancel(context.Background()) + ctx := metadata.NewIncomingContext(baseContext, metadata.Pairs( + consts.EventStreamHistoryReplayHeader, "true", + )) + stream := newCaptureEventsStream(ctx) + result := make(chan error, 1) + go func() { result <- NewServer().Events(&clientpb.Empty{}, stream) }() + + select { + case event := <-stream.events: + if event.Op != "history" { + t.Fatalf("replayed event op = %q, want history", event.Op) + } + case <-time.After(2 * time.Second): + t.Fatal("negotiated event history was not replayed") + } + if values := stream.responseHeader().Get(consts.EventStreamHistoryCountHeader); len(values) != 1 || values[0] != "1" { + t.Fatalf("history count header = %v, want [1]", values) + } + cancel() + select { + case err := <-result: + if err != nil { + t.Fatalf("Events returned error: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Events did not stop after context cancellation") + } +} + +func TestEventsKeepsLegacyClientsOnLiveOnlyStream(t *testing.T) { + broker := newEventTestBroker(t) + if err := broker.TryPublish(core.Event{ + EventType: consts.EventBroadcast, + Op: "history", + Important: true, + }); err != nil { + t.Fatalf("publish history: %v", err) + } + waitHistorySize(t, broker, 1) + + ctx, cancel := context.WithCancel(context.Background()) + stream := newCaptureEventsStream(ctx) + result := make(chan error, 1) + go func() { result <- NewServer().Events(&clientpb.Empty{}, stream) }() + select { + case <-stream.headerReady: + case <-time.After(2 * time.Second): + t.Fatal("legacy event stream header was not sent") + } + if values := stream.responseHeader().Get(consts.EventStreamHistoryCountHeader); len(values) != 0 { + t.Fatalf("legacy history count header = %v, want absent", values) + } + if err := broker.TryPublish(core.Event{EventType: consts.EventBroadcast, Op: "live"}); err != nil { + t.Fatalf("publish live event: %v", err) + } + select { + case event := <-stream.events: + if event.Op != "live" { + t.Fatalf("legacy first event op = %q, want live", event.Op) + } + case <-time.After(2 * time.Second): + t.Fatal("legacy live event was not delivered") + } + cancel() + select { + case err := <-result: + if err != nil { + t.Fatalf("Events returned error: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Events did not stop after context cancellation") + } +} + +func TestEventsStopsCleanlyWhenBrokerCloses(t *testing.T) { + broker := newEventTestBroker(t) + stream := newCaptureEventsStream(context.Background()) + result := make(chan error, 1) + go func() { result <- NewServer().Events(&clientpb.Empty{}, stream) }() + select { + case <-stream.headerReady: + case <-time.After(2 * time.Second): + t.Fatal("event stream header was not sent") + } + + broker.Stop() + select { + case err := <-result: + if err != nil { + t.Fatalf("Events returned error after broker close: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Events did not stop after broker closure") + } +} + +func TestEventsKeepsLegacyStreamAliveAfterSubscriberOverflow(t *testing.T) { + broker := newEventTestBroker(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stream := newCaptureEventsStream(ctx) + stream.events = make(chan *clientpb.Event, 256) + stream.sendStarted = make(chan struct{}) + sendGate := make(chan struct{}) + stream.sendGate = sendGate + result := make(chan error, 1) + go func() { result <- NewServer().Events(&clientpb.Empty{}, stream) }() + select { + case <-stream.headerReady: + case <-time.After(2 * time.Second): + t.Fatal("event stream header was not sent") + } + + const eventCount = 64 + for i := 0; i < eventCount; i++ { + event := core.Event{ + EventType: consts.EventBroadcast, + Op: "overflow-" + strconv.Itoa(i), + Important: true, + } + deadline := time.Now().Add(2 * time.Second) + for { + err := broker.TryPublish(event) + if err == nil { + break + } + if !errors.Is(err, core.ErrEventBrokerQueueFull) { + t.Fatalf("publish overflow event %d: %v", i, err) + } + if time.Now().After(deadline) { + t.Fatalf("publish overflow event %d stayed queue-full", i) + } + time.Sleep(time.Millisecond) + } + } + select { + case <-stream.sendStarted: + case <-time.After(2 * time.Second): + t.Fatal("event stream Send did not block") + } + lastOp := "overflow-" + strconv.Itoa(eventCount-1) + deadline := time.Now().Add(2 * time.Second) + for { + history := broker.GetAll() + if len(history) > 0 && history[len(history)-1].Op == lastOp { + break + } + if time.Now().After(deadline) { + t.Fatal("broker did not process overflow burst") + } + time.Sleep(time.Millisecond) + } + + close(sendGate) + marker := core.Event{EventType: consts.EventBroadcast, Op: "after-overflow"} + deadline = time.Now().Add(2 * time.Second) + for { + if err := broker.TryPublish(marker); err != nil && !errors.Is(err, core.ErrEventBrokerQueueFull) { + t.Fatalf("publish marker: %v", err) + } + select { + case event := <-stream.events: + if event.Op == marker.Op { + goto markerReceived + } + case err := <-result: + t.Fatalf("Events stopped after legacy overflow: %v", err) + case <-time.After(time.Millisecond): + } + if time.Now().After(deadline) { + t.Fatal("legacy stream did not recover after subscriber overflow") + } + } + +markerReceived: + cancel() + select { + case err := <-result: + if err != nil && !errors.Is(err, context.Canceled) { + t.Fatalf("Events returned error after cancellation: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Events did not stop after context cancellation") + } +} diff --git a/server/rpc/rpc_events_v2_test.go b/server/rpc/rpc_events_v2_test.go index f399be76a..a00a81484 100644 --- a/server/rpc/rpc_events_v2_test.go +++ b/server/rpc/rpc_events_v2_test.go @@ -93,6 +93,7 @@ func TestGetEventHonorsLatestEventLimit(t *testing.T) { Important: true, }) } + waitHistorySize(t, core.EventBroker, 3) limited, err := (&Server{}).GetEvent(context.Background(), &clientpb.Int{Limit: 2}) if err != nil { diff --git a/server/rpc/rpc_forward_listener_test.go b/server/rpc/rpc_forward_listener_test.go index 9374d629d..10fc8069c 100644 --- a/server/rpc/rpc_forward_listener_test.go +++ b/server/rpc/rpc_forward_listener_test.go @@ -7,6 +7,7 @@ import ( "net" "os" "path/filepath" + "sync" "sync/atomic" "testing" "time" @@ -247,6 +248,63 @@ func TestEnsureForwardTaskStreamDoesNotOpenDuplicateConcurrentStreams(t *testing cancel() } +func TestResetForwardListenerRuntimesWaitsForInFlightStop(t *testing.T) { + resetForwardListenerRuntimes() + withIsolatedListenersAndJobs(t) + + const listenerID = "forward-stop-reset-race" + core.Listeners.Add(core.NewListener(listenerID, "127.0.0.1")) + stopEntered := make(chan struct{}) + releaseStop := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseStop) }) } + t.Cleanup(release) + + runtime := &forwardListenerRuntime{ + listenerID: listenerID, + ownsListener: true, + cancel: func() { + close(stopEntered) + <-releaseStop + }, + } + forwardListenerRuntimes.Store(listenerID, runtime) + + stopDone := make(chan struct{}) + go func() { + runtime.stop() + close(stopDone) + }() + select { + case <-stopEntered: + case <-time.After(time.Second): + t.Fatal("runtime stop did not reach the deterministic barrier") + } + + resetDone := make(chan struct{}) + go func() { + resetForwardListenerRuntimes() + close(resetDone) + }() + select { + case <-resetDone: + t.Fatal("reset returned while the runtime stop was still in flight") + case <-time.After(25 * time.Millisecond): + } + + release() + select { + case <-stopDone: + case <-time.After(time.Second): + t.Fatal("runtime stop did not finish after releasing the barrier") + } + select { + case <-resetDone: + case <-time.After(time.Second): + t.Fatal("reset did not finish after runtime stop completed") + } +} + type blockingForwardTaskClient struct { calls atomic.Int32 started chan int32 diff --git a/server/rpc/rpc_implant_session_handlers_test.go b/server/rpc/rpc_implant_session_handlers_test.go index dd25a1341..3036b6da7 100644 --- a/server/rpc/rpc_implant_session_handlers_test.go +++ b/server/rpc/rpc_implant_session_handlers_test.go @@ -72,8 +72,9 @@ func TestSleepDispatchesRequestAndUpdatesSession(t *testing.T) { if err != nil { t.Fatalf("Sleep failed: %v", err) } - if sess.Expression != req.Expression || sess.Jitter != req.Jitter { - t.Fatalf("session timer = %q/%v, want %q/%v", sess.Expression, sess.Jitter, req.Expression, req.Jitter) + expression, jitter := sess.TimerSnapshot() + if expression != req.Expression || jitter != req.Jitter { + t.Fatalf("session timer = %q/%v, want %q/%v", expression, jitter, req.Expression, req.Jitter) } select { @@ -91,10 +92,7 @@ func TestSleepDispatchesRequestAndUpdatesSession(t *testing.T) { deliverTaskResponse(t, sess, task.TaskId, &implantpb.Spite{ Body: &implantpb.Spite_Empty{Empty: &implantpb.Empty{}}, }) - waitForCondition(t, 2*time.Second, func() bool { - stored := sess.Tasks.Get(task.TaskId) - return stored != nil && stored.Finished() - }, "sleep task to finish") + waitForTaskDone(t, sess.Tasks.Get(task.TaskId), "sleep task") waitForCondition(t, 2*time.Second, func() bool { _, pending := sess.GetResp(task.TaskId) return !pending @@ -133,6 +131,7 @@ func TestKeepaliveEnablesSessionAfterResponse(t *testing.T) { Name: consts.ModuleKeepalive, Body: &implantpb.Spite_Common{Common: &implantpb.CommonBody{Name: consts.ModuleKeepalive}}, }) + waitForTaskDone(t, sess.Tasks.Get(task.TaskId), "keepalive task") waitForCondition(t, 2*time.Second, sess.IsKeepaliveEnabled, "keepalive to become enabled") waitForCondition(t, 2*time.Second, func() bool { _, pending := sess.GetResp(task.TaskId) @@ -163,12 +162,14 @@ func TestInfoUpdatesSessionSysinfoFromResponse(t *testing.T) { Process: &implantpb.Process{Name: "agent.bin"}, }}, }) + waitForTaskDone(t, sess.Tasks.Get(task.TaskId), "info task") waitForCondition(t, 2*time.Second, func() bool { _, pending := sess.GetResp(task.TaskId) return !pending }, "info response handler to finish cleanup") - if sess.Os == nil || sess.Os.Name != "linux" || sess.Os.Arch != "x64" { - t.Fatalf("session OS = %#v, want normalized linux/x64", sess.Os) + pb := sess.ToProtobufLite() + if pb.Os == nil || pb.Os.Name != "linux" || pb.Os.Arch != "x64" { + t.Fatalf("session sysinfo = %#v, want linux/x64", pb.Os) } } @@ -444,3 +445,21 @@ func waitForCondition(t testing.TB, timeout time.Duration, cond func() bool, des } t.Fatalf("timed out waiting for %s", description) } + +func waitForTaskDone(t testing.TB, task *core.Task, description string) { + t.Helper() + if task == nil { + t.Fatalf("%s was not stored", description) + } + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if task.IsClosed() { + if !task.Finished() { + t.Fatalf("%s closed but stayed running", description) + } + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s to finish", description) +} diff --git a/server/rpc/rpc_runtime_test.go b/server/rpc/rpc_runtime_test.go index 71a8a2427..e8b4148b4 100644 --- a/server/rpc/rpc_runtime_test.go +++ b/server/rpc/rpc_runtime_test.go @@ -124,7 +124,10 @@ func TestPollingPublishesSessionErrorAndStopsRuntime(t *testing.T) { broker := core.NewBroker() defer broker.Stop() waitEventBrokerReady(t, broker) - sub := broker.Subscribe() + sub, err := broker.Subscribe() + if err != nil { + t.Fatalf("Subscribe error = %v", err) + } defer broker.Unsubscribe(sub) core.EventBroker = broker @@ -140,7 +143,7 @@ func TestPollingPublishesSessionErrorAndStopsRuntime(t *testing.T) { } core.Sessions.Add(sess) - _, err := (&Server{}).Polling(context.Background(), &clientpb.Polling{ + _, err = (&Server{}).Polling(context.Background(), &clientpb.Polling{ SessionId: sess.ID, Id: "polling-1", Interval: 1, diff --git a/server/testsupport/mock_transport.go b/server/testsupport/mock_transport.go index 73f9cd43f..f79fdfb79 100644 --- a/server/testsupport/mock_transport.go +++ b/server/testsupport/mock_transport.go @@ -56,9 +56,7 @@ func StartInProcessListener(t testing.TB, h *ControlPlaneHarness, listenerName s } t.Cleanup(func() { - if serverlistener.Listener != nil { - _ = serverlistener.Listener.Close() - } + _ = serverlistener.CloseCurrentListener() }) } diff --git a/server/testsupport/real_implant.go b/server/testsupport/real_implant.go index 79e0a56bc..b575baa8d 100644 --- a/server/testsupport/real_implant.go +++ b/server/testsupport/real_implant.go @@ -372,10 +372,8 @@ func (r *RealImplant) Close() error { } } - if listener.Listener != nil { - if err := listener.Listener.Close(); err != nil { - errs = append(errs, fmt.Sprintf("close listener: %v", err)) - } + if err := listener.CloseCurrentListener(); err != nil { + errs = append(errs, fmt.Sprintf("close listener: %v", err)) } for _, path := range []string{r.BinaryPath, r.ProfilePath, r.AuthPath} {