From e55cb1ca5f2c3672421f8b0e7e51b51caba57cfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sat, 11 Jul 2026 03:28:54 +0800 Subject: [PATCH 01/29] test(audit): add vulnerability reproduction coverage --- .../session_request_timeout_audit_test.go | 84 +++++++ .../malefic/untrusted_input_audit_test.go | 19 ++ .../malefic/untrusted_input_fuzz_test.go | 90 ++++++++ .../parser/untrusted_length_audit_test.go | 55 +++++ .../forward_backpressure_audit_test.go | 218 ++++++++++++++++++ .../forward_closed_stream_audit_test.go | 33 +++ .../lifecycle_concurrency_audit_test.go | 71 ++++++ .../pipeline_start_close_audit_test.go | 107 +++++++++ 8 files changed, 677 insertions(+) create mode 100644 server/internal/core/session_request_timeout_audit_test.go create mode 100644 server/internal/parser/malefic/untrusted_input_audit_test.go create mode 100644 server/internal/parser/malefic/untrusted_input_fuzz_test.go create mode 100644 server/internal/parser/untrusted_length_audit_test.go create mode 100644 server/listener/forward_backpressure_audit_test.go create mode 100644 server/listener/forward_closed_stream_audit_test.go create mode 100644 server/listener/lifecycle_concurrency_audit_test.go create mode 100644 server/listener/pipeline_start_close_audit_test.go diff --git a/server/internal/core/session_request_timeout_audit_test.go b/server/internal/core/session_request_timeout_audit_test.go new file mode 100644 index 000000000..7d48dd1d5 --- /dev/null +++ b/server/internal/core/session_request_timeout_audit_test.go @@ -0,0 +1,84 @@ +//go:build audit + +package core + +import ( + "errors" + "testing" + "time" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" + implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" +) + +func TestAuditSessionRequestAndWaitHonorsTimeout(t *testing.T) { + const taskID uint32 = 4242 + const requestTimeout = 25 * time.Millisecond + + sess := newTestSession("audit-request-timeout") + defer sess.Cancel() + + sent := make(chan struct{}) + stream := &testServerStream{ + sendMsg: func(interface{}) error { + close(sent) + return nil + }, + } + + type result struct { + spite *implantpb.Spite + err error + } + resultCh := make(chan result, 1) + exited := make(chan struct{}) + go func() { + defer close(exited) + spite, err := sess.RequestAndWait(&clientpb.SpiteRequest{ + Task: &clientpb.Task{TaskId: taskID}, + }, stream, requestTimeout) + resultCh <- result{spite: spite, err: err} + }() + defer func() { + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + for { + select { + case <-exited: + sess.RemoveResp(taskID) + return + case <-deadline.C: + sess.RemoveResp(taskID) + t.Errorf("RequestAndWait goroutine did not exit during test cleanup") + return + default: + } + + if respCh, ok := sess.GetResp(taskID); ok { + select { + case respCh <- &implantpb.Spite{Name: "audit-cleanup"}: + default: + } + } + time.Sleep(time.Millisecond) + } + }() + + select { + case <-sent: + case <-time.After(time.Second): + t.Fatal("request was not sent") + } + + select { + case got := <-resultCh: + if !errors.Is(got.err, ErrImplantSendTimeout) { + t.Fatalf("RequestAndWait error = %v, want %v", got.err, ErrImplantSendTimeout) + } + if got.spite != nil { + t.Fatalf("RequestAndWait response = %#v, want nil", got.spite) + } + case <-time.After(10 * requestTimeout): + t.Fatalf("RequestAndWait did not return within %s; the timeout argument is not honored", 10*requestTimeout) + } +} diff --git a/server/internal/parser/malefic/untrusted_input_audit_test.go b/server/internal/parser/malefic/untrusted_input_audit_test.go new file mode 100644 index 000000000..2f0e51325 --- /dev/null +++ b/server/internal/parser/malefic/untrusted_input_audit_test.go @@ -0,0 +1,19 @@ +//go:build audit + +package malefic + +import ( + "bytes" + "testing" +) + +func TestAuditMaleficParserReadHeader_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") + } +} 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..97dfbb13a --- /dev/null +++ b/server/internal/parser/malefic/untrusted_input_fuzz_test.go @@ -0,0 +1,90 @@ +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 + } + // A complete header may be rejected by current or future boundary checks. + // If it is accepted, its decoded fields must still be internally consistent. + if err != nil { + return + } + + wantSID := binary.LittleEndian.Uint32(data[MsgSessionStart:MsgSessionEnd]) + wantLength := binary.LittleEndian.Uint32(data[MsgSessionEnd:HeaderLength]) + 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 && decodedLength > fuzzMaxPayloadLength { + return + } + } + + _, _ = p.Parse(payload) + }) +} diff --git a/server/internal/parser/untrusted_length_audit_test.go b/server/internal/parser/untrusted_length_audit_test.go new file mode 100644 index 000000000..0ddbad324 --- /dev/null +++ b/server/internal/parser/untrusted_length_audit_test.go @@ -0,0 +1,55 @@ +//go:build audit + +package parser + +import ( + "encoding/binary" + "io" + "testing" + + "github.com/chainreactors/IoM-go/consts" + "github.com/chainreactors/malice-network/server/internal/parser/malefic" +) + +type auditHeaderThenEOFConn struct { + header []byte + offset int + maxReadRequest int +} + +func (c *auditHeaderThenEOFConn) 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 *auditHeaderThenEOFConn) Write(p []byte) (int, error) { return len(p), nil } +func (c *auditHeaderThenEOFConn) Close() error { return nil } + +func TestAuditReadPacket_RejectsOversizedDeclarationBeforePayloadAllocation(t *testing.T) { + const configuredLimit = uint32(1024) + declaredLength := configuredLimit + uint32(consts.KB)*16 + 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) + } + p.WithMaxPacketLength(configuredLimit) + conn := &auditHeaderThenEOFConn{header: header} + + _, _, _ = p.ReadPacket(conn) + if conn.maxReadRequest > malefic.HeaderLength { + t.Fatalf("ReadPacket requested a %d-byte payload buffer for an oversized declaration; want rejection after the %d-byte header", conn.maxReadRequest, malefic.HeaderLength) + } +} diff --git a/server/listener/forward_backpressure_audit_test.go b/server/listener/forward_backpressure_audit_test.go new file mode 100644 index 000000000..0655bc967 --- /dev/null +++ b/server/listener/forward_backpressure_audit_test.go @@ -0,0 +1,218 @@ +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 := newForwardStreamRegistry().get("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) + } + 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..f50f7a0d0 --- /dev/null +++ b/server/listener/forward_closed_stream_audit_test.go @@ -0,0 +1,33 @@ +//go:build audit + +package listener + +import ( + "errors" + "io" + "testing" + + "github.com/chainreactors/IoM-go/proto/client/clientpb" +) + +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) + } +} diff --git a/server/listener/lifecycle_concurrency_audit_test.go b/server/listener/lifecycle_concurrency_audit_test.go new file mode 100644 index 000000000..51c447a50 --- /dev/null +++ b/server/listener/lifecycle_concurrency_audit_test.go @@ -0,0 +1,71 @@ +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), + } + + oldListener := Listener + Listener = lns + t.Cleanup(func() { + Listener = oldListener + }) + + runAuditConcurrentClose(t, lns.Close) +} 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..a9e395a7b --- /dev/null +++ b/server/listener/pipeline_start_close_audit_test.go @@ -0,0 +1,107 @@ +//go:build audit + +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) +} From 30767dd4df1430e1c9d8a65a38e3f9b032d82c4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sat, 11 Jul 2026 05:26:28 +0800 Subject: [PATCH 02/29] fix(parser): harden untrusted payload reads --- server/internal/parser/malefic/parser.go | 26 +++-- server/internal/parser/malefic/parser_test.go | 25 ++++- ..._audit_test.go => untrusted_input_test.go} | 8 +- server/internal/parser/parser.go | 34 +++++- .../parser/untrusted_length_audit_test.go | 55 ---------- .../internal/parser/untrusted_length_test.go | 100 ++++++++++++++++++ 6 files changed, 172 insertions(+), 76 deletions(-) rename server/internal/parser/malefic/{untrusted_input_audit_test.go => untrusted_input_test.go} (61%) delete mode 100644 server/internal/parser/untrusted_length_audit_test.go create mode 100644 server/internal/parser/untrusted_length_test.go diff --git a/server/internal/parser/malefic/parser.go b/server/internal/parser/malefic/parser.go index 79c8b2732..736768840 100644 --- a/server/internal/parser/malefic/parser.go +++ b/server/internal/parser/malefic/parser.go @@ -34,11 +34,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 +104,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) { diff --git a/server/internal/parser/malefic/parser_test.go b/server/internal/parser/malefic/parser_test.go index c98e24151..1186d7984 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) } } diff --git a/server/internal/parser/malefic/untrusted_input_audit_test.go b/server/internal/parser/malefic/untrusted_input_test.go similarity index 61% rename from server/internal/parser/malefic/untrusted_input_audit_test.go rename to server/internal/parser/malefic/untrusted_input_test.go index 2f0e51325..011dd10f0 100644 --- a/server/internal/parser/malefic/untrusted_input_audit_test.go +++ b/server/internal/parser/malefic/untrusted_input_test.go @@ -1,13 +1,12 @@ -//go:build audit - package malefic import ( "bytes" + "strings" "testing" ) -func TestAuditMaleficParserReadHeader_RejectsMaxDeclaredLength(t *testing.T) { +func TestMaleficParserReadHeader_RejectsMaxDeclaredLength(t *testing.T) { p := NewMaleficParser() p.MaxPacketLength = 1024 header := buildHeader(DefaultStartDelimiter, 0x10203040, ^uint32(0)) @@ -16,4 +15,7 @@ func TestAuditMaleficParserReadHeader_RejectsMaxDeclaredLength(t *testing.T) { 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..2f53544c3 100644 --- a/server/internal/parser/parser.go +++ b/server/internal/parser/parser.go @@ -57,6 +57,8 @@ type MessageParser struct { PacketParser } +const payloadReadChunkSize = 64 * 1024 + // WithSecure 为 MessageParser 添加安全支持 func (mp *MessageParser) WithSecure(keyPair *clientpb.KeyPair) { switch mp.Implant { @@ -80,8 +82,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 +95,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 +104,32 @@ func (parser *MessageParser) ReadPacket(conn io.ReadWriteCloser) (uint32, *impla return sessionId, msg, err } +func readPayload(conn io.Reader, length uint32) ([]byte, error) { + if length == 0 { + return []byte{}, nil + } + + chunkLength := payloadReadChunkSize + if uint32(chunkLength) > length { + chunkLength = int(length) + } + chunk := make([]byte, chunkLength) + payload := make([]byte, 0, chunkLength) + remaining := uint64(length) + for remaining > 0 { + readLength := len(chunk) + if uint64(readLength) > remaining { + readLength = int(remaining) + } + if _, err := io.ReadFull(conn, chunk[:readLength]); err != nil { + return nil, err + } + payload = append(payload, chunk[:readLength]...) + remaining -= uint64(readLength) + } + 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/untrusted_length_audit_test.go b/server/internal/parser/untrusted_length_audit_test.go deleted file mode 100644 index 0ddbad324..000000000 --- a/server/internal/parser/untrusted_length_audit_test.go +++ /dev/null @@ -1,55 +0,0 @@ -//go:build audit - -package parser - -import ( - "encoding/binary" - "io" - "testing" - - "github.com/chainreactors/IoM-go/consts" - "github.com/chainreactors/malice-network/server/internal/parser/malefic" -) - -type auditHeaderThenEOFConn struct { - header []byte - offset int - maxReadRequest int -} - -func (c *auditHeaderThenEOFConn) 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 *auditHeaderThenEOFConn) Write(p []byte) (int, error) { return len(p), nil } -func (c *auditHeaderThenEOFConn) Close() error { return nil } - -func TestAuditReadPacket_RejectsOversizedDeclarationBeforePayloadAllocation(t *testing.T) { - const configuredLimit = uint32(1024) - declaredLength := configuredLimit + uint32(consts.KB)*16 + 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) - } - p.WithMaxPacketLength(configuredLimit) - conn := &auditHeaderThenEOFConn{header: header} - - _, _, _ = p.ReadPacket(conn) - if conn.maxReadRequest > malefic.HeaderLength { - t.Fatalf("ReadPacket requested a %d-byte payload buffer for an oversized declaration; want rejection after the %d-byte header", conn.maxReadRequest, malefic.HeaderLength) - } -} diff --git a/server/internal/parser/untrusted_length_test.go b/server/internal/parser/untrusted_length_test.go new file mode 100644 index 000000000..4f41da43c --- /dev/null +++ b/server/internal/parser/untrusted_length_test.go @@ -0,0 +1,100 @@ +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_DeclaredLengthUsesBoundedReadBuffer(t *testing.T) { + const configuredLimit = uint32(1024) + declaredLength := uint32(payloadReadChunkSize * 4) + + 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 > payloadReadChunkSize { + t.Fatalf("ReadPacket requested a %d-byte read buffer, want at most %d", conn.maxReadRequest, payloadReadChunkSize) + } +} + +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, payloadReadChunkSize*2) + 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") + } +} From 46281bef20a187cf541756ccb3b1757a67f5435a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sat, 11 Jul 2026 05:26:34 +0800 Subject: [PATCH 03/29] fix(listener): linearize pipeline lifecycle --- server/internal/core/forward.go | 61 ++++- server/internal/core/forward_runtime_test.go | 56 +++++ server/listener/forward.go | 63 ++++- .../forward_backpressure_audit_test.go | 1 + .../forward_closed_stream_audit_test.go | 119 +++++++++- server/listener/forward_dispatch_test.go | 11 +- server/listener/forward_rollback_test.go | 181 +++++++++++++++ server/listener/http.go | 160 +++++++++---- .../lifecycle_concurrency_audit_test.go | 10 +- .../listener/lifecycle_linearization_test.go | 217 ++++++++++++++++++ server/listener/listener.go | 41 +++- .../pipeline_start_close_audit_test.go | 2 - server/listener/rem.go | 173 +++++++++++--- server/listener/rem_runtime_test.go | 20 +- server/listener/tcp.go | 148 +++++++++--- server/testsupport/mock_transport.go | 4 +- server/testsupport/real_implant.go | 6 +- 17 files changed, 1125 insertions(+), 148 deletions(-) create mode 100644 server/listener/forward_rollback_test.go create mode 100644 server/listener/lifecycle_linearization_test.go diff --git a/server/internal/core/forward.go b/server/internal/core/forward.go index 3c98cd46c..bc5e59df3 100644 --- a/server/internal/core/forward.go +++ b/server/internal/core/forward.go @@ -129,6 +129,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 +139,32 @@ 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 + count int Pipeline ListenerId string Stream ForwardStream @@ -164,9 +172,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 { @@ -192,10 +203,31 @@ 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) }) } +// 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 { + if f == nil { + return nil + } + f.abortOnce.Do(func() { + f.shutdown() + if stream, ok := f.Stream.(interface{ CloseSend() error }); ok { + f.abortErr = stream.CloseSend() + } + if f.handlerDone != nil { + <-f.handlerDone + } + }) + return f.abortErr +} + func (f *Forward) Count() int { return f.count } @@ -211,7 +243,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 +302,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..933ab9f70 100644 --- a/server/internal/core/forward_runtime_test.go +++ b/server/internal/core/forward_runtime_test.go @@ -37,6 +37,30 @@ func (testForwardStream) Recv() (*clientpb.SpiteRequest, error) { return nil, errors.New("not used") } +type abortForwardStream struct { + closeCalls atomic.Int32 +} + +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 +76,38 @@ 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 TestForwardHandlerReturnsStreamSendError(t *testing.T) { want := errors.New("stream send failed") forward := &Forward{ diff --git a/server/listener/forward.go b/server/listener/forward.go index 0bf66c63c..5a466879c 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" @@ -73,7 +74,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 } @@ -243,6 +244,11 @@ type forwardLocalStream struct { events chan *clientpb.SpiteRequest done chan struct{} closeOnce sync.Once + closed atomic.Bool + eventMu sync.Mutex + slotsOnce sync.Once + eventSlots chan struct{} + remoteSend sync.WaitGroup } func (s *forwardLocalStream) Send(resp *clientpb.SpiteResponse) error { @@ -297,13 +303,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 } @@ -328,9 +339,13 @@ func (s *forwardLocalStream) serve(stream forwardrpc.ForwardListener_TaskStreamS func (s *forwardLocalStream) close() { s.closeOnce.Do(func() { + s.eventMu.Lock() + s.closed.Store(true) if s.done != nil { close(s.done) } + s.eventMu.Unlock() + s.remoteSend.Wait() }) } @@ -338,16 +353,56 @@ func (s *forwardLocalStream) sendEvent(event *clientpb.SpiteRequest) error { if event == nil { return nil } + 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() { + s.eventMu.Unlock() + s.releaseEventSlot() + return io.ErrClosedPipe + } + s.events <- event + s.eventMu.Unlock() return nil case <-s.done: return io.ErrClosedPipe - case <-time.After(10 * time.Second): + 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 index 0655bc967..8ca109069 100644 --- a/server/listener/forward_backpressure_audit_test.go +++ b/server/listener/forward_backpressure_audit_test.go @@ -80,6 +80,7 @@ func TestAuditForwardEventQueueBackpressurePreservesEvent(t *testing.T) { 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 { diff --git a/server/listener/forward_closed_stream_audit_test.go b/server/listener/forward_closed_stream_audit_test.go index f50f7a0d0..f3509e7ce 100644 --- a/server/listener/forward_closed_stream_audit_test.go +++ b/server/listener/forward_closed_stream_audit_test.go @@ -1,13 +1,16 @@ -//go:build audit - 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) { @@ -31,3 +34,115 @@ func TestAuditForwardSendEventRejectsAfterClose(t *testing.T) { 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") + } +} 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..17f8a6667 --- /dev/null +++ b/server/listener/forward_rollback_test.go @@ -0,0 +1,181 @@ +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 +} + +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 nil +} + +type rollbackPipelineRPC struct { + stream *rollbackForwardStream +} + +type rollbackListener struct { + closed chan struct{} + closeOnce sync.Once + closeCalls atomic.Int32 +} + +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 nil +} +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 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 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 TestTCPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { + oldListen, oldCmux := tcpListen, tcpStartCmux + ln := newRollbackListener() + 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) + } + assertForwardRolledBack(t, pipeline, stream, pipeline.Start) + 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 + ln := newRollbackListener() + 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) + } + assertForwardRolledBack(t, pipeline, stream, pipeline.Start) + if got := ln.closeCalls.Load(); got != 1 { + t.Fatalf("listener Close calls = %d, want 1", got) + } +} diff --git a/server/listener/http.go b/server/listener/http.go index 7c16128cd..203f91891 100644 --- a/server/listener/http.go +++ b/server/listener/http.go @@ -14,6 +14,7 @@ import ( "net" "net/http" "strconv" + "sync" "github.com/chainreactors/logs" "github.com/chainreactors/malice-network/server/internal/certutils" @@ -22,6 +23,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 +55,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 +74,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 +86,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 +112,46 @@ 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 + } + ln := pipeline.srv + pipeline.srv = nil + pipeline.stateMu.Unlock() + if ln == nil { return nil } + if err := ln.Close(); err != nil && !errors.Is(err, net.ErrClosed) { + return err + } return nil } func (pipeline *HTTPPipeline) Start() error { - if pipeline.Enable { + 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) - } - dispatchForwardTaskRequest("http", pipeline.Name, msg) + committed := false + defer func() { + if !committed { + _ = forward.Abort() + 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 } @@ -142,13 +160,13 @@ func (pipeline *HTTPPipeline) Start() error { mux.HandleFunc("/", pipeline.handler) if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable && pipeline.TLSConfig.Cert != nil { - err := pipeline.startWithCmux(ln, mux) + err := httpStartWithCmux(pipeline, ln, mux) if err != nil { + _ = ln.Close() 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 +175,83 @@ func (pipeline *HTTPPipeline) Start() error { return nil }, pipeline.runtimeErrorHandler("serve loop")) } + if !pipeline.commitStart(ln, forward) { + _ = ln.Close() + return nil + } + committed = true 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) + 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")) + pipeline.starting = false + done := pipeline.startDone + pipeline.startDone = nil + close(done) + return true +} + // startWithCmux 使用 cmux 实现 HTTP TLS 和非 TLS 的端口复用 func (pipeline *HTTPPipeline) startWithCmux(ln net.Listener, mux *http.ServeMux) error { // 获取 TLS 配置 @@ -177,9 +265,6 @@ func (pipeline *HTTPPipeline) startWithCmux(ln net.Listener, mux *http.ServeMux) return err } - // 保存服务器引用用于关闭 - pipeline.srv = ln - return nil } @@ -348,10 +433,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/lifecycle_concurrency_audit_test.go b/server/listener/lifecycle_concurrency_audit_test.go index 51c447a50..29f7a2cb3 100644 --- a/server/listener/lifecycle_concurrency_audit_test.go +++ b/server/listener/lifecycle_concurrency_audit_test.go @@ -61,10 +61,14 @@ func TestAuditListenerConcurrentClose(t *testing.T) { websites: make(map[string]*Website), } - oldListener := Listener - Listener = lns + listenerGlobalMu.Lock() + oldListener := currentListener + currentListener = lns + listenerGlobalMu.Unlock() t.Cleanup(func() { - Listener = oldListener + 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..ffefc730f --- /dev/null +++ b/server/listener/lifecycle_linearization_test.go @@ -0,0 +1,217 @@ +package listener + +import ( + "context" + "errors" + "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{} +} + +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 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 07b8051eb..9ae6ef982 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) @@ -217,13 +242,21 @@ type listener struct { websites map[string]*Website shutdown func() error retireOnce sync.Once + closeOnce sync.Once + closeErr error } func (lns *listener) Close() error { if lns == nil { return nil } + lns.closeOnce.Do(func() { + lns.closeErr = lns.close() + }) + return lns.closeErr +} +func (lns *listener) close() error { var errs []error for _, pipeline := range lns.pipelines.ToProtobuf().GetPipelines() { @@ -259,9 +292,7 @@ func (lns *listener) Close() error { } } - if Listener == lns { - Listener = nil - } + clearCurrentListener(lns) return errors.Join(errs...) } diff --git a/server/listener/pipeline_start_close_audit_test.go b/server/listener/pipeline_start_close_audit_test.go index a9e395a7b..0c5e199ec 100644 --- a/server/listener/pipeline_start_close_audit_test.go +++ b/server/listener/pipeline_start_close_audit_test.go @@ -1,5 +1,3 @@ -//go:build audit - package listener import ( diff --git a/server/listener/rem.go b/server/listener/rem.go index ef83125ac..90f6da03c 100644 --- a/server/listener/rem.go +++ b/server/listener/rem.go @@ -26,7 +26,13 @@ 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() +} func NewRem(rpc listenerrpc.ListenerRPCClient, pipeline *clientpb.Pipeline) (*REM, error) { remConfig := pipeline.GetRem() @@ -54,13 +60,18 @@ 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 + startDone chan struct{} + stopCh chan struct{} + 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 +81,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() return err } - rem.Enable = true + if !rem.enabled() { + _ = remConsoleClose(rem.con) + rem.abortStart() + return nil + } 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() { + _ = remConsoleClose(rem.con) + rem.abortStart() + return nil + } return nil } func (rem *REM) ToProtobuf() *clientpb.Pipeline { + enabled := rem.enabled() link := rem.getLink() subscribe := rem.getSubscribe() host := "" @@ -121,7 +141,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,18 +201,95 @@ func (rem *REM) getSubscribe() (subscribe string) { } func (rem *REM) Close() error { + rem.stateMu.Lock() + wasActive := rem.Enable rem.Enable = false - if rem.con == nil { + stopCh := rem.stopCh + rem.stopCh = nil + if rem.starting { + done := rem.startDone + rem.stateMu.Unlock() + if stopCh != nil { + close(stopCh) + } + <-done return nil } - return rem.con.Close() + rem.stateMu.Unlock() + if stopCh != nil { + close(stopCh) + } + if !wasActive || rem.con == nil { + return nil + } + 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.startDone + rem.stateMu.Unlock() + <-done + continue + } + if rem.Enable { + rem.stateMu.Unlock() + return false + } + rem.Enable = true + rem.starting = true + rem.startDone = make(chan struct{}) + rem.stopCh = make(chan struct{}) + rem.stateMu.Unlock() + return true + } +} + +func (rem *REM) abortStart() { + rem.stateMu.Lock() + rem.starting = false + rem.Enable = false + done := rem.startDone + rem.startDone = nil + stopCh := rem.stopCh + rem.stopCh = nil + if done != nil { + close(done) + } + rem.stateMu.Unlock() + if stopCh != nil { + close(stopCh) + } +} + +func (rem *REM) commitStart() bool { + rem.stateMu.Lock() + defer rem.stateMu.Unlock() + if !rem.Enable { + return false + } + core.GoGuarded("rem-accept:"+rem.Name, rem.acceptLoop, rem.runtimeErrorHandler("accept loop")) + core.GoGuarded("rem-health:"+rem.Name, rem.healthLoop, rem.runtimeErrorHandler("health loop")) + rem.starting = false + done := rem.startDone + rem.startDone = nil + close(done) + return true } func (rem *REM) acceptLoop() error { - for rem.Enable { + for rem.enabled() { ag, err := rem.con.Accept() if err != nil { - if !rem.Enable { + if !rem.enabled() { return nil } // Accept errors are typically transient (timeout, client disconnect). @@ -228,7 +325,8 @@ func (rem *REM) healthLoop() error { consecutiveFailures := 0 unhealthy := false - for rem.Enable { + stopCh := rem.stopSignal() + for rem.enabled() { if err := remHealthCheck(rem.rpc, context.Background(), rem.ToProtobuf()); err != nil { consecutiveFailures++ logs.Log.Errorf("rem %s health check failed (%d/%d): %v", rem.Name, consecutiveFailures, healthFailureThreshold, err) @@ -258,20 +356,41 @@ 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 <-stopCh: + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + return nil + } } return nil } +func (rem *REM) stopSignal() <-chan struct{} { + rem.stateMu.Lock() + defer rem.stateMu.Unlock() + if rem.stopCh == nil { + rem.stopCh = make(chan struct{}) + } + return rem.stopCh +} + 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 +422,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..62d15db38 100644 --- a/server/listener/rem_runtime_test.go +++ b/server/listener/rem_runtime_test.go @@ -27,14 +27,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{ @@ -52,12 +49,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 }() @@ -96,13 +91,12 @@ 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) @@ -121,7 +115,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 +129,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..7915e9273 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,12 +88,20 @@ 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 + pipeline.stateMu.Unlock() + if ln == nil { + return nil + } err := ln.Close() if err != nil && !errors.Is(err, net.ErrClosed) { return err @@ -90,20 +110,90 @@ func (pipeline *TCPPipeline) Close() error { } func (pipeline *TCPPipeline) Start() error { - if pipeline.Enable { + if !pipeline.beginStart() { return nil } + forward, err := core.NewForward(pipeline.rpc, pipeline) if err != nil { + pipeline.abortStart() return err } forward.ListenerId = pipeline.ListenerID + committed := false + defer func() { + if !committed { + _ = forward.Abort() + pipeline.abortStart() + } + }() + + ln, err := pipeline.handler() + if err != nil { + return err + } + if !pipeline.commitStart(ln, forward) { + _ = ln.Close() + return nil + } + committed = true + logs.Log.Infof("pipeline.tcp - start host=%s port=%d parser=%s tls=%t", + pipeline.Host, pipeline.Port, pipeline.Parser, pipeline.TLSConfig.Enable) + return nil +} + +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 + } +} + +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() +} + +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) core.GoGuarded("tcp-forward-recv:"+pipeline.Name, func() error { for { msg, err := forward.Stream.Recv() if err != nil { - if !pipeline.Enable { + if !pipeline.enabled() { return nil } return fmt.Errorf("tcp pipeline %s forward recv: %w", pipeline.Name, err) @@ -111,26 +201,27 @@ func (pipeline *TCPPipeline) Start() error { dispatchForwardTaskRequest("tcp", pipeline.Name, msg) } }, pipeline.runtimeErrorHandler("forward recv loop")) - - pipeline.ln, err = pipeline.handler() - if err != nil { - return err - } - 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 + pipeline.starting = false + done := pipeline.startDone + pipeline.startDone = nil + close(done) + return true } func (pipeline *TCPPipeline) handler() (net.Listener, error) { - ln, err := net.Listen("tcp", fmt.Sprintf("%s:%d", pipeline.Host, pipeline.Port)) + ln, err := tcpListen("tcp", fmt.Sprintf("%s:%d", pipeline.Host, pipeline.Port)) if err != nil { return nil, err } // 如果启用了 TLS，使用 cmux 实现 TLS 和非 TLS 的端口复用 if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable { - return pipeline.handleWithCmux(ln) + cmuxListener, err := pipeline.handleWithCmux(ln) + if err != nil { + _ = ln.Close() + return nil, err + } + return cmuxListener, nil } // 非 TLS 模式，使用原有逻辑 @@ -156,7 +247,7 @@ func (pipeline *TCPPipeline) handleWithCmux(ln net.Listener) (net.Listener, erro } } - return StartCmuxTCPListener(ln, tlsConfig, pipeline.HandleConnection, pipeline.runtimeErrorHandler("cmux")) + return tcpStartCmux(ln, tlsConfig, pipeline.HandleConnection, pipeline.runtimeErrorHandler("cmux")) } // startAcceptLoop 启动连接接受循环 (用于非 cmux 模式) @@ -165,7 +256,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 +344,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/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} { From 60b37f7820c087bfb6d1ee80d64201576f509e01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sat, 11 Jul 2026 05:26:40 +0800 Subject: [PATCH 04/29] fix(core): synchronize session sysinfo state --- server/internal/core/session.go | 72 +++++++++------- .../session_request_timeout_audit_test.go | 84 ------------------- server/internal/core/session_test.go | 55 ++++++++++++ server/rpc/grpc.go | 21 ----- server/rpc/rpc-filesystem.go | 2 +- .../rpc/rpc_implant_session_handlers_test.go | 36 ++++++-- 6 files changed, 128 insertions(+), 142 deletions(-) delete mode 100644 server/internal/core/session_request_timeout_audit_test.go diff --git a/server/internal/core/session.go b/server/internal/core/session.go index 247338260..c96cbc9c9 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 + sysInfoMu sync.RWMutex Ctx context.Context Cancel context.CancelFunc @@ -355,6 +354,8 @@ func (w *SpiteStreamWriter) finish(err error) { } func (s *Session) Abstract() string { + s.sysInfoMu.RLock() + defer s.sysInfoMu.RUnlock() if s.Os == nil { return fmt.Sprintf("%s(%s)", s.Name, s.ID) } else { @@ -553,7 +554,8 @@ func (s *Session) isAlived() bool { } func (s *Session) ToProtobuf() *clientpb.Session { - return &clientpb.Session{ + s.sysInfoMu.RLock() + pb := &clientpb.Session{ Type: s.Type, SessionId: s.ID, RawId: s.RawID, @@ -565,8 +567,8 @@ func (s *Session) ToProtobuf() *clientpb.Session { Target: s.Target, PipelineId: s.PipelineID, ListenerId: s.ListenerID, - Os: s.Os, - Process: s.Process, + Os: cloneSysInfoOS(s.Os), + Process: cloneSysInfoProcess(s.Process), LastCheckin: s.LastCheckinUnix(), Filepath: s.SessionContext.Filepath, Workdir: s.SessionContext.WorkDir, @@ -574,15 +576,19 @@ func (s *Session) ToProtobuf() *clientpb.Session { 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(), } + s.sysInfoMu.RUnlock() + pb.Tasks = s.Tasks.ToProtobuf() + return pb } func (s *Session) ToProtobufLite() *clientpb.Session { + s.sysInfoMu.RLock() + defer s.sysInfoMu.RUnlock() return &clientpb.Session{ Type: s.Type, SessionId: s.ID, @@ -595,8 +601,8 @@ func (s *Session) ToProtobufLite() *clientpb.Session { Target: s.Target, PipelineId: s.PipelineID, ListenerId: s.ListenerID, - Os: s.Os, - Process: s.Process, + Os: cloneSysInfoOS(s.Os), + Process: cloneSysInfoProcess(s.Process), LastCheckin: s.LastCheckinUnix(), Filepath: s.SessionContext.Filepath, Workdir: s.SessionContext.WorkDir, @@ -616,6 +622,8 @@ func (s *Session) Save() error { } func (s *Session) ToModel() *models.Session { + name := s.Name + s.sysInfoMu.RLock() sessModel := &models.Session{ SessionID: s.ID, RawID: s.RawID, @@ -630,7 +638,8 @@ func (s *Session) ToModel() *models.Session { LastCheckin: s.LastCheckinUnix(), DataString: s.Marshal(), } - artifact, err := sessionDBGetArtifact(s.Name) + s.sysInfoMu.RUnlock() + artifact, err := sessionDBGetArtifact(name) if err == nil && artifact.ProfileName != "" { if _, profileErr := sessionDBGetProfile(artifact.ProfileName); profileErr == nil { sessModel.ProfileName = artifact.ProfileName @@ -702,14 +711,17 @@ func (s *Session) Update(req *clientpb.RegisterSession) { s.Jitter = req.RegisterData.Timer.Jitter } s.SessionContext.Update(req) - // 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) + s.sysInfoMu.Lock() + shouldPublishInit = !s.Initialized + s.updateSysInfoLocked(req.RegisterData.Sysinfo) + s.sysInfoMu.Unlock() + } + if shouldPublishInit { + s.Publish(consts.CtrlSessionInit, fmt.Sprintf("session %s init", s.ID), true, true) } err := s.Save() @@ -722,6 +734,18 @@ func (s *Session) UpdateSysInfo(info *implantpb.SysInfo) { if info == nil { return } + s.sysInfoMu.Lock() + defer s.sysInfoMu.Unlock() + s.updateSysInfoLocked(info) +} + +func (s *Session) SetWorkDir(workDir string) { + s.sysInfoMu.Lock() + s.WorkDir = workDir + s.sysInfoMu.Unlock() +} + +func (s *Session) updateSysInfoLocked(info *implantpb.SysInfo) { s.Initialized = true osInfo := mergeSysInfoOS(s.Os, info.GetOs()) processInfo := mergeSysInfoProcess(s.Process, info.GetProcess()) @@ -882,15 +906,18 @@ func hasSignatureMetadata(value *implantpb.Process) bool { } func (s *Session) FillSysInfo() { - artifact, err := sessionDBGetArtifact(s.Name) + name := s.Name + 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.sysInfoMu.Lock() s.Os = &implantpb.Os{ Name: artifact.Os, Arch: artifact.Arch, } + s.sysInfoMu.Unlock() } func (s *Session) Publish(Op string, msg string, notify bool, important bool) { @@ -927,17 +954,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) diff --git a/server/internal/core/session_request_timeout_audit_test.go b/server/internal/core/session_request_timeout_audit_test.go deleted file mode 100644 index 7d48dd1d5..000000000 --- a/server/internal/core/session_request_timeout_audit_test.go +++ /dev/null @@ -1,84 +0,0 @@ -//go:build audit - -package core - -import ( - "errors" - "testing" - "time" - - "github.com/chainreactors/IoM-go/proto/client/clientpb" - implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb" -) - -func TestAuditSessionRequestAndWaitHonorsTimeout(t *testing.T) { - const taskID uint32 = 4242 - const requestTimeout = 25 * time.Millisecond - - sess := newTestSession("audit-request-timeout") - defer sess.Cancel() - - sent := make(chan struct{}) - stream := &testServerStream{ - sendMsg: func(interface{}) error { - close(sent) - return nil - }, - } - - type result struct { - spite *implantpb.Spite - err error - } - resultCh := make(chan result, 1) - exited := make(chan struct{}) - go func() { - defer close(exited) - spite, err := sess.RequestAndWait(&clientpb.SpiteRequest{ - Task: &clientpb.Task{TaskId: taskID}, - }, stream, requestTimeout) - resultCh <- result{spite: spite, err: err} - }() - defer func() { - deadline := time.NewTimer(time.Second) - defer deadline.Stop() - for { - select { - case <-exited: - sess.RemoveResp(taskID) - return - case <-deadline.C: - sess.RemoveResp(taskID) - t.Errorf("RequestAndWait goroutine did not exit during test cleanup") - return - default: - } - - if respCh, ok := sess.GetResp(taskID); ok { - select { - case respCh <- &implantpb.Spite{Name: "audit-cleanup"}: - default: - } - } - time.Sleep(time.Millisecond) - } - }() - - select { - case <-sent: - case <-time.After(time.Second): - t.Fatal("request was not sent") - } - - select { - case got := <-resultCh: - if !errors.Is(got.err, ErrImplantSendTimeout) { - t.Fatalf("RequestAndWait error = %v, want %v", got.err, ErrImplantSendTimeout) - } - if got.spite != nil { - t.Fatalf("RequestAndWait response = %#v, want nil", got.spite) - } - case <-time.After(10 * requestTimeout): - t.Fatalf("RequestAndWait did not return within %s; the timeout argument is not honored", 10*requestTimeout) - } -} diff --git a/server/internal/core/session_test.go b/server/internal/core/session_test.go index 1a794d2ae..fcda14890 100644 --- a/server/internal/core/session_test.go +++ b/server/internal/core/session_test.go @@ -454,6 +454,61 @@ 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() +} + // ---------- Task management ---------- func TestSession_NewTask_IncrementsSeq(t *testing.T) { 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-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_implant_session_handlers_test.go b/server/rpc/rpc_implant_session_handlers_test.go index 6762e2f8b..4bf8dd8ec 100644 --- a/server/rpc/rpc_implant_session_handlers_test.go +++ b/server/rpc/rpc_implant_session_handlers_test.go @@ -91,10 +91,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") } func TestKeepaliveEnablesSessionAfterResponse(t *testing.T) { @@ -129,7 +126,10 @@ func TestKeepaliveEnablesSessionAfterResponse(t *testing.T) { Name: consts.ModuleKeepalive, Body: &implantpb.Spite_Common{Common: &implantpb.CommonBody{Name: consts.ModuleKeepalive}}, }) - waitForCondition(t, 2*time.Second, sess.IsKeepaliveEnabled, "keepalive to become enabled") + waitForTaskDone(t, sess.Tasks.Get(task.TaskId), "keepalive task") + if !sess.IsKeepaliveEnabled() { + t.Fatal("keepalive stayed disabled after task completion") + } } func TestInfoUpdatesSessionSysinfoFromResponse(t *testing.T) { @@ -155,9 +155,11 @@ func TestInfoUpdatesSessionSysinfoFromResponse(t *testing.T) { Process: &implantpb.Process{Name: "agent.bin"}, }}, }) - waitForCondition(t, 2*time.Second, func() bool { - return sess.Os != nil && sess.Os.Name == "linux" && sess.Os.Arch == "x64" - }, "session sysinfo update") + waitForTaskDone(t, sess.Tasks.Get(task.TaskId), "info task") + 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) + } } func TestGetSessionReturnsDatabaseRecordWithoutRecovery(t *testing.T) { @@ -432,3 +434,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) +} From 4fc14271e1ef270d3a324968a8b8b6dc45042e1f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sat, 11 Jul 2026 05:26:44 +0800 Subject: [PATCH 05/29] ci: align workflows with Go 1.25 --- .github/workflows/ci.yaml | 20 ++++++++++++++++---- .github/workflows/nightly.yaml | 5 ++++- .github/workflows/realimplant.yaml | 5 ++++- .github/workflows/releaser.yaml | 7 ++++--- .github/workflows/update-search-index.yaml | 5 ++++- 5 files changed, 32 insertions(+), 10 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index b390882ad..1376228be 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -19,8 +19,11 @@ 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 @@ -50,8 +53,11 @@ 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 @@ -77,8 +83,11 @@ 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 @@ -107,8 +116,11 @@ 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 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..73bc7f3af 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: | 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/" From 055d82d7e272a2fe7cd4f5b27ad06df5f6c0a684 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sun, 12 Jul 2026 01:08:25 +0800 Subject: [PATCH 06/29] fix(rpc): wait for forward listener shutdown --- server/rpc/rpc-forward-listener.go | 2 +- server/rpc/rpc_forward_listener_test.go | 58 +++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/server/rpc/rpc-forward-listener.go b/server/rpc/rpc-forward-listener.go index f8275a23d..addbcfcca 100644 --- a/server/rpc/rpc-forward-listener.go +++ b/server/rpc/rpc-forward-listener.go @@ -207,7 +207,7 @@ func (r *forwardListenerRuntime) stop() { return } r.stopOnce.Do(func() { - forwardListenerRuntimes.Delete(r.listenerID) + defer forwardListenerRuntimes.Delete(r.listenerID) clearForwardTaskStreams(r.listenerID) if r.cancel != nil { r.cancel() diff --git a/server/rpc/rpc_forward_listener_test.go b/server/rpc/rpc_forward_listener_test.go index 4419bf39a..edb36498f 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 From d4ba4ef2f835ac370bd5b189c0a7a271d75f61f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sun, 12 Jul 2026 01:42:27 +0800 Subject: [PATCH 07/29] fix(runtime): close event and artifact consistency gaps --- client/command/build/artifact_command_test.go | 10 +- client/command/website/webcontent.go | 63 ++++++------ client/command/website/webcontent_cmd_test.go | 57 +++++++---- client/core/server.go | 13 ++- external/IoM-go | 2 +- server/internal/core/event.go | 31 ++++-- server/internal/core/event_runtime_test.go | 97 +++++++++++++++---- server/internal/core/safe_test.go | 2 +- .../internal/core/session_lifecycle_test.go | 18 ++-- server/listener/rem_runtime_test.go | 5 +- server/rpc/generic_runtime_test.go | 9 +- server/rpc/rpc-events.go | 13 ++- server/rpc/rpc_runtime_test.go | 7 +- 13 files changed, 231 insertions(+), 96 deletions(-) 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/website/webcontent.go b/client/command/website/webcontent.go index 9e878e797..fe414d660 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 04e11fe93..16511d796 100644 --- a/client/core/server.go +++ b/client/core/server.go @@ -130,10 +130,9 @@ func NewServerWithOptions(conn *grpc.ClientConn, config *mtls.ClientConfig, supp } for _, event := range events.GetEvents() { if suppressStartupOutput { - ser.ReconcileEvent(event) continue } - ser.HandlerEvent(event) + ser.handleEventOutput(event) } return ser, nil } @@ -435,6 +434,9 @@ func (s *Server) EventHandler() { if err != nil { return } + if _, err := eventStream.Header(); err != nil { + return + } s.Update() if s.Session != nil { s.UpdateSession(s.GetInteractive().SessionId) @@ -496,6 +498,13 @@ func (s *Server) HandlerEvent(event *clientpb.Event) { } // Reconcile state first — single entry point for all map updates s.ReconcileEvent(event) + s.handleEventOutput(event) +} + +func (s *Server) handleEventOutput(event *clientpb.Event) { + if s == nil || event == nil { + return + } // Quiet mode (non-index mux pane): suppress notification events but let // task events through so user-initiated commands still show results. diff --git a/external/IoM-go b/external/IoM-go index 03c1b18f6..41d5e8a4c 160000 --- a/external/IoM-go +++ b/external/IoM-go @@ -1 +1 @@ -Subproject commit 03c1b18f6bad962b1c76e5155a465adab877b322 +Subproject commit 41d5e8a4c4dcc5347156c7c0715f6244fb15b94d diff --git a/server/internal/core/event.go b/server/internal/core/event.go index eff3823c6..fbb51df25 100644 --- a/server/internal/core/event.go +++ b/server/internal/core/event.go @@ -220,10 +220,15 @@ func (event *Event) ToProtobuf() *clientpb.Event { } } +type eventSubscription struct { + events chan Event + ready chan struct{} +} + 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 @@ -254,8 +259,9 @@ func (broker *eventBroker) run() error { select { case <-broker.stop: return nil - case sub := <-broker.subscribe: - subscribers[sub] = struct{}{} + case subscription := <-broker.subscribe: + subscribers[subscription.events] = struct{}{} + close(subscription.ready) case sub := <-broker.unsubscribe: delete(subscribers, sub) case event := <-broker.publish: @@ -315,15 +321,24 @@ func (broker *eventBroker) Stop() { } // Subscribe - Generate a new subscription channel -func (broker *eventBroker) Subscribe() chan Event { +func (broker *eventBroker) Subscribe() (chan Event, error) { events := make(chan Event, eventBufSize) - broker.subscribe <- events - return events + ready := make(chan struct{}) + select { + case broker.subscribe <- eventSubscription{events: events, ready: ready}: + case <-broker.stop: + return nil, ErrEventBrokerUnavailable + } + <-ready + return events, nil } // 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) } @@ -382,7 +397,7 @@ 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(), diff --git a/server/internal/core/event_runtime_test.go b/server/internal/core/event_runtime_test.go index 2ed8f8459..db49be9df 100644 --- a/server/internal/core/event_runtime_test.go +++ b/server/internal/core/event_runtime_test.go @@ -53,22 +53,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) @@ -118,7 +108,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", @@ -127,7 +119,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) @@ -176,6 +171,72 @@ 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 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 TestEventToProtobufWebsiteWithoutTLSDoesNotPanic(t *testing.T) { event := Event{ EventType: consts.EventJob, diff --git a/server/internal/core/safe_test.go b/server/internal/core/safe_test.go index aed714365..a7ad1568f 100644 --- a/server/internal/core/safe_test.go +++ b/server/internal/core/safe_test.go @@ -18,7 +18,7 @@ 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(), diff --git a/server/internal/core/session_lifecycle_test.go b/server/internal/core/session_lifecycle_test.go index 940b38b9a..496362dcc 100644 --- a/server/internal/core/session_lifecycle_test.go +++ b/server/internal/core/session_lifecycle_test.go @@ -25,7 +25,7 @@ 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(), @@ -88,7 +88,7 @@ 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(), @@ -167,7 +167,7 @@ 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(), @@ -279,7 +279,7 @@ 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(), @@ -363,7 +363,7 @@ 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(), @@ -466,7 +466,7 @@ 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(), @@ -594,7 +594,7 @@ 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(), @@ -786,7 +786,7 @@ 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(), @@ -936,7 +936,7 @@ 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(), diff --git a/server/listener/rem_runtime_test.go b/server/listener/rem_runtime_test.go index 62d15db38..7655a76bf 100644 --- a/server/listener/rem_runtime_test.go +++ b/server/listener/rem_runtime_test.go @@ -63,7 +63,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) 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/rpc-events.go b/server/rpc/rpc-events.go index 85dfe75a0..eb1c5acd2 100644 --- a/server/rpc/rpc-events.go +++ b/server/rpc/rpc-events.go @@ -2,17 +2,26 @@ package rpc import ( "context" + "strconv" + "github.com/chainreactors/IoM-go/consts" "github.com/chainreactors/IoM-go/proto/client/clientpb" "github.com/chainreactors/IoM-go/proto/services/clientrpc" "github.com/chainreactors/logs" "github.com/chainreactors/malice-network/server/internal/core" "github.com/chainreactors/malice-network/server/internal/db" - "strconv" + "google.golang.org/grpc/metadata" ) func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsServer) error { - events := core.EventBroker.Subscribe() + events, err := core.EventBroker.Subscribe() + if err != nil { + return err + } + if err := stream.SendHeader(metadata.Pairs(consts.EventStreamReadyHeader, "true")); err != nil { + core.EventBroker.Unsubscribe(events) + return err + } clientID := core.GetCurrentID() defer func() { logs.Log.Infof("client: %d disconnected", clientID) diff --git a/server/rpc/rpc_runtime_test.go b/server/rpc/rpc_runtime_test.go index 9d2dfe9dd..5a186a855 100644 --- a/server/rpc/rpc_runtime_test.go +++ b/server/rpc/rpc_runtime_test.go @@ -114,7 +114,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 @@ -130,7 +133,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, From 05047feb6c85d28bc9659a3c5ea7549c9bc3ff42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Sun, 12 Jul 2026 14:43:08 +0800 Subject: [PATCH 08/29] fix(parser): cap snappy decoded payloads --- helper/utils/compress/compress.go | 5 +++++ server/internal/parser/malefic/parser.go | 11 ++++++++++- server/internal/parser/malefic/parser_test.go | 12 ++++++++++++ .../parser/malefic/untrusted_input_fuzz_test.go | 13 +++++++++++-- 4 files changed, 38 insertions(+), 3 deletions(-) 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/server/internal/parser/malefic/parser.go b/server/internal/parser/malefic/parser.go index 736768840..c83f0033a 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,8 +25,11 @@ const ( HeaderLength = 9 DefaultStartDelimiter = 0xd1 DefaultEndDelimiter = 0xd2 + maxDecodedPayloadSize = uint64(2 << 30) ) +var ErrDecodedPayloadTooLarge = errors.New("decoded payload exceeds 2 GiB limit") + func NewMaleficParser() *MaleficParser { return &MaleficParser{ StartDelimiter: DefaultStartDelimiter, @@ -158,7 +162,12 @@ func (parser *MaleficParser) Parse(buf []byte) (*implantpb.Spites, error) { } } - buf, err := compress.Decompress(buf) + decodedLength, err := compress.DecodedLen(buf) + if err == nil && uint64(decodedLength) > maxDecodedPayloadSize { + return nil, fmt.Errorf("%w: %d bytes", ErrDecodedPayloadTooLarge, decodedLength) + } + + buf, err = compress.Decompress(buf) if err != nil { logs.Log.Debugf("trying plaintext: %v", err) } diff --git a/server/internal/parser/malefic/parser_test.go b/server/internal/parser/malefic/parser_test.go index 1186d7984..dd3fa3f7f 100644 --- a/server/internal/parser/malefic/parser_test.go +++ b/server/internal/parser/malefic/parser_test.go @@ -177,6 +177,18 @@ func TestMaleficParser_Parse_EmptyPayload(t *testing.T) { } } +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_ZeroLengthSlice(t *testing.T) { p := NewMaleficParser() defer func() { diff --git a/server/internal/parser/malefic/untrusted_input_fuzz_test.go b/server/internal/parser/malefic/untrusted_input_fuzz_test.go index 97dfbb13a..db11dab35 100644 --- a/server/internal/parser/malefic/untrusted_input_fuzz_test.go +++ b/server/internal/parser/malefic/untrusted_input_fuzz_test.go @@ -80,8 +80,17 @@ func FuzzMaleficParserParse(f *testing.F) { // 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 && decodedLength > fuzzMaxPayloadLength { - return + 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 + } } } From f6224cd2d5e9177e9bdc0105d86c20b32481ca48 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:37:58 +0800 Subject: [PATCH 09/29] fix(client): synchronize shared runtime state --- client/command/build/commands.go | 4 ++-- client/command/implant.go | 7 ++++--- client/core/server.go | 12 ++++++++---- external/IoM-go | 2 +- 4 files changed, 15 insertions(+), 10 deletions(-) diff --git a/client/command/build/commands.go b/client/command/build/commands.go index 90488890f..cc5971159 100644 --- a/client/command/build/commands.go +++ b/client/command/build/commands.go @@ -597,13 +597,13 @@ artifact comment artifact_name "" 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 cf5b0d3f5..7890a067c 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/core/server.go b/client/core/server.go index 16511d796..4a388f3ac 100644 --- a/client/core/server.go +++ b/client/core/server.go @@ -437,9 +437,13 @@ func (s *Server) EventHandler() { if _, err := eventStream.Header(); err != nil { return } - s.Update() - if s.Session != nil { - s.UpdateSession(s.GetInteractive().SessionId) + 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() { @@ -452,7 +456,7 @@ func (s *Server) EventHandler() { } s.dispatchEventHooks(event) - if fn, ok := s.EventCallback[event.Op]; ok { + if fn, ok := s.GetEventCallback(event.Op); ok { fn(event) } s.HandlerEvent(event) diff --git a/external/IoM-go b/external/IoM-go index 41d5e8a4c..1f1bf8df3 160000 --- a/external/IoM-go +++ b/external/IoM-go @@ -1 +1 @@ -Subproject commit 41d5e8a4c4dcc5347156c7c0715f6244fb15b94d +Subproject commit 1f1bf8df33972dcaf46447da4261bdafc029f9b4 From 3bafe462930d4f3088d2e9971bd6448fd30b7818 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:43:36 +0800 Subject: [PATCH 10/29] fix(core): synchronize session state snapshots --- server/internal/core/session.go | 189 +++++++++++++++++++-------- server/internal/core/session_test.go | 68 ++++++++++ server/rpc/rpc-module.go | 15 +-- server/rpc/rpc-session.go | 4 +- 4 files changed, 203 insertions(+), 73 deletions(-) diff --git a/server/internal/core/session.go b/server/internal/core/session.go index c96cbc9c9..6bd09d2e7 100644 --- a/server/internal/core/session.go +++ b/server/internal/core/session.go @@ -271,7 +271,7 @@ type Session struct { keepaliveMu sync.Mutex keepaliveEnabled bool anyMu sync.RWMutex - sysInfoMu sync.RWMutex + stateMu sync.RWMutex Ctx context.Context Cancel context.CancelFunc @@ -354,8 +354,8 @@ func (w *SpiteStreamWriter) finish(err error) { } func (s *Session) Abstract() string { - s.sysInfoMu.RLock() - defer s.sysInfoMu.RUnlock() + s.stateMu.RLock() + defer s.stateMu.RUnlock() if s.Os == nil { return fmt.Sprintf("%s(%s)", s.Name, s.ID) } else { @@ -496,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 @@ -539,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 } @@ -554,41 +566,25 @@ func (s *Session) isAlived() bool { } func (s *Session) ToProtobuf() *clientpb.Session { - s.sysInfoMu.RLock() - pb := &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: 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: s.Modules, - Addons: s.Addons, - KeyPair: s.KeyPair, // 添加密钥对 - Data: s.Marshal(), - } - s.sysInfoMu.RUnlock() + pb := s.ToProtobufLite() pb.Tasks = s.Tasks.ToProtobuf() return pb } func (s *Session) ToProtobufLite() *clientpb.Session { - s.sysInfoMu.RLock() - defer s.sysInfoMu.RUnlock() + 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, @@ -596,7 +592,7 @@ func (s *Session) ToProtobufLite() *clientpb.Session { Note: s.Note, Name: s.Name, GroupName: s.Group, - IsAlive: s.isAlived(), + IsAlive: s.isAliveAtLocked(now), IsPrivilege: s.IsPrivilege, Target: s.Target, PipelineId: s.PipelineID, @@ -610,20 +606,33 @@ func (s *Session) ToProtobufLite() *clientpb.Session { 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, + Modules: modules, + Addons: addons, + KeyPair: keyPair, Data: s.Marshal(), } } +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 { return sessionDBSave(s.ToModel()) } func (s *Session) ToModel() *models.Session { + s.stateMu.RLock() name := s.Name - s.sysInfoMu.RLock() sessModel := &models.Session{ SessionID: s.ID, RawID: s.RawID, @@ -634,11 +643,11 @@ 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(), } - s.sysInfoMu.RUnlock() + s.stateMu.RUnlock() artifact, err := sessionDBGetArtifact(name) if err == nil && artifact.ProfileName != "" { if _, profileErr := sessionDBGetProfile(artifact.ProfileName); profileErr == nil { @@ -668,6 +677,50 @@ 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) 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 { @@ -692,16 +745,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 @@ -710,16 +774,16 @@ 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 { - s.sysInfoMu.Lock() shouldPublishInit = !s.Initialized s.updateSysInfoLocked(req.RegisterData.Sysinfo) - s.sysInfoMu.Unlock() } + s.stateMu.Unlock() if shouldPublishInit { s.Publish(consts.CtrlSessionInit, fmt.Sprintf("session %s init", s.ID), true, true) } @@ -734,15 +798,15 @@ func (s *Session) UpdateSysInfo(info *implantpb.SysInfo) { if info == nil { return } - s.sysInfoMu.Lock() - defer s.sysInfoMu.Unlock() + s.stateMu.Lock() + defer s.stateMu.Unlock() s.updateSysInfoLocked(info) } func (s *Session) SetWorkDir(workDir string) { - s.sysInfoMu.Lock() + s.stateMu.Lock() s.WorkDir = workDir - s.sysInfoMu.Unlock() + s.stateMu.Unlock() } func (s *Session) updateSysInfoLocked(info *implantpb.SysInfo) { @@ -906,18 +970,20 @@ func hasSignatureMetadata(value *implantpb.Process) bool { } func (s *Session) FillSysInfo() { + 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", name, err) return } - s.sysInfoMu.Lock() + s.stateMu.Lock() s.Os = &implantpb.Os{ Name: artifact.Os, Arch: artifact.Arch, } - s.sysInfoMu.Unlock() + s.stateMu.Unlock() } func (s *Session) Publish(Op string, msg string, notify bool, important bool) { @@ -1058,10 +1124,17 @@ func (s *Session) DeleteResp(taskId uint32) { // UpdateKeyPair 更新KeyPair并同步到SecureManager func (s *Session) UpdateKeyPair(keyPair *clientpb.KeyPair) { - s.SessionContext.KeyPair = keyPair + var stored *clientpb.KeyPair + if keyPair != nil { + stored = proto.Clone(keyPair).(*clientpb.KeyPair) + } + s.stateMu.Lock() + s.SessionContext.KeyPair = stored + manager := s.SecureManager + s.stateMu.Unlock() // 更新SecureManager中的KeyPair引用 - if s.SecureManager != nil { - s.SecureManager.UpdateKeyPair(keyPair) + if manager != nil { + manager.UpdateKeyPair(stored) } } @@ -1243,10 +1316,12 @@ func (s *Session) UpdatePrivateKey(key string) { func (s *Session) UpdateKeyPairFieldsAndPushCtrl(publicKey string, privateKey string) { next := &clientpb.KeyPair{} + s.stateMu.RLock() if s.KeyPair != nil { next.PublicKey = s.KeyPair.PublicKey next.PrivateKey = s.KeyPair.PrivateKey } + s.stateMu.RUnlock() if publicKey != "" { next.PublicKey = publicKey } diff --git a/server/internal/core/session_test.go b/server/internal/core/session_test.go index fcda14890..5cfcad34f 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" @@ -509,6 +510,73 @@ func TestSessionSysInfoSnapshotConcurrent(t *testing.T) { 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/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 } From 9512851cbf4339ccf972a365e54c8844641c81b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:47:24 +0800 Subject: [PATCH 11/29] fix(parser): reject unsafe encoded payloads --- server/internal/parser/malefic/parser.go | 12 ++++-- server/internal/parser/malefic/parser_test.go | 43 +++++++------------ .../malefic/untrusted_input_fuzz_test.go | 13 ++++-- server/internal/parser/parser.go | 28 ++++-------- server/internal/parser/parser_test.go | 39 +++++++++-------- .../internal/parser/untrusted_length_test.go | 32 +++++++++++--- 6 files changed, 91 insertions(+), 76 deletions(-) diff --git a/server/internal/parser/malefic/parser.go b/server/internal/parser/malefic/parser.go index c83f0033a..03d3cb4d8 100644 --- a/server/internal/parser/malefic/parser.go +++ b/server/internal/parser/malefic/parser.go @@ -28,7 +28,10 @@ const ( maxDecodedPayloadSize = uint64(2 << 30) ) -var ErrDecodedPayloadTooLarge = errors.New("decoded payload exceeds 2 GiB limit") +var ( + ErrDecodedPayloadTooLarge = errors.New("decoded payload exceeds 2 GiB limit") + ErrInvalidCompressedPayload = errors.New("invalid Snappy payload") +) func NewMaleficParser() *MaleficParser { return &MaleficParser{ @@ -163,13 +166,16 @@ func (parser *MaleficParser) Parse(buf []byte) (*implantpb.Spites, error) { } decodedLength, err := compress.DecodedLen(buf) - if err == nil && uint64(decodedLength) > maxDecodedPayloadSize { + 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) } 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 dd3fa3f7f..4aa4142bb 100644 --- a/server/internal/parser/malefic/parser_test.go +++ b/server/internal/parser/malefic/parser_test.go @@ -164,16 +164,9 @@ 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) } } @@ -191,13 +184,17 @@ func TestMaleficParser_Parse_RejectsDecodedPayloadAboveLimit(t *testing.T) { 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) { @@ -271,8 +268,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 @@ -286,15 +281,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 index db11dab35..1b62f6d46 100644 --- a/server/internal/parser/malefic/untrusted_input_fuzz_test.go +++ b/server/internal/parser/malefic/untrusted_input_fuzz_test.go @@ -49,14 +49,19 @@ func FuzzMaleficParserReadHeader(f *testing.F) { } return } - // A complete header may be rejected by current or future boundary checks. - // If it is accepted, its decoded fields must still be internally consistent. - if err != nil { + 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 := binary.LittleEndian.Uint32(data[MsgSessionEnd:HeaderLength]) + 1 + wantLength := declaredLength + 1 if sid != wantSID || length != wantLength { t.Fatalf("ReadHeader() = (%d, %d), want (%d, %d)", sid, length, wantSID, wantLength) } diff --git a/server/internal/parser/parser.go b/server/internal/parser/parser.go index 2f53544c3..3261090c3 100644 --- a/server/internal/parser/parser.go +++ b/server/internal/parser/parser.go @@ -57,7 +57,9 @@ type MessageParser struct { PacketParser } -const payloadReadChunkSize = 64 * 1024 +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) { @@ -105,27 +107,15 @@ func (parser *MessageParser) ReadPacket(conn io.ReadWriteCloser) (uint32, *impla } 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 } - - chunkLength := payloadReadChunkSize - if uint32(chunkLength) > length { - chunkLength = int(length) - } - chunk := make([]byte, chunkLength) - payload := make([]byte, 0, chunkLength) - remaining := uint64(length) - for remaining > 0 { - readLength := len(chunk) - if uint64(readLength) > remaining { - readLength = int(remaining) - } - if _, err := io.ReadFull(conn, chunk[:readLength]); err != nil { - return nil, err - } - payload = append(payload, chunk[:readLength]...) - remaining -= uint64(readLength) + payload := make([]byte, int(length)) + if _, err := io.ReadFull(conn, payload); err != nil { + return nil, err } return payload, 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 index 4f41da43c..1a912be9b 100644 --- a/server/internal/parser/untrusted_length_test.go +++ b/server/internal/parser/untrusted_length_test.go @@ -33,9 +33,9 @@ func (c *headerThenEOFConn) Read(p []byte) (int, error) { func (c *headerThenEOFConn) Write(p []byte) (int, error) { return len(p), nil } func (c *headerThenEOFConn) Close() error { return nil } -func TestReadPacket_DeclaredLengthUsesBoundedReadBuffer(t *testing.T) { +func TestReadPacket_DeclaredLengthBelowLimitReadsDeclaredPayload(t *testing.T) { const configuredLimit = uint32(1024) - declaredLength := uint32(payloadReadChunkSize * 4) + declaredLength := uint32(256 * 1024) header := make([]byte, malefic.HeaderLength) header[malefic.MsgStart] = malefic.DefaultStartDelimiter @@ -53,8 +53,30 @@ func TestReadPacket_DeclaredLengthUsesBoundedReadBuffer(t *testing.T) { if !errors.Is(err, io.EOF) { t.Fatalf("ReadPacket error = %v, want EOF after the header", err) } - if conn.maxReadRequest > payloadReadChunkSize { - t.Fatalf("ReadPacket requested a %d-byte read buffer, want at most %d", conn.maxReadRequest, payloadReadChunkSize) + 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) } } @@ -66,7 +88,7 @@ func TestReadPacket_AcceptsValidFrameLargerThanPacketLength(t *testing.T) { } p.WithMaxPacketLength(configuredPacketLength) - stdout := make([]byte, payloadReadChunkSize*2) + stdout := make([]byte, 128*1024) x := uint32(1) for i := range stdout { x ^= x << 13 From a206fc8824410fbed15d6f1952077263964fc81f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:51:22 +0800 Subject: [PATCH 12/29] fix(events): make stream shutdown explicit --- client/core/server.go | 8 +++- server/internal/core/event.go | 42 ++++++++++++------- server/internal/core/event_runtime_test.go | 18 ++++++++ server/internal/core/safe_test.go | 1 - .../internal/core/session_lifecycle_test.go | 9 ---- server/rpc/rpc-events.go | 5 ++- 6 files changed, 55 insertions(+), 28 deletions(-) diff --git a/client/core/server.go b/client/core/server.go index 4a388f3ac..c6ef22d55 100644 --- a/client/core/server.go +++ b/client/core/server.go @@ -451,7 +451,13 @@ 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) diff --git a/server/internal/core/event.go b/server/internal/core/event.go index fbb51df25..8c60ccb4e 100644 --- a/server/internal/core/event.go +++ b/server/internal/core/event.go @@ -26,6 +26,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 ) @@ -233,12 +234,12 @@ type eventBroker struct { send chan Event notifier inotify.Notifier - lock *sync.Mutex cache *RingCache alive atomic.Bool managed atomic.Bool startOnce sync.Once + stopOnce sync.Once } func (broker *eventBroker) run() error { @@ -247,12 +248,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 { @@ -270,18 +266,24 @@ func (broker *eventBroker) run() error { } 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() } } } +func closeEventSubscriber(sub chan Event) { + defer func() { + _ = recover() + }() + close(sub) +} + func (broker *eventBroker) dispatch(sub chan Event, event Event) (err error) { defer func() { if recovered := recover(); recovered != nil { @@ -290,10 +292,10 @@ 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 + return ErrEventSubscriberSlow } - return nil } func (broker *eventBroker) Start() { @@ -317,7 +319,12 @@ func (broker *eventBroker) Start() { // Stop - Close the broker channel func (broker *eventBroker) Stop() { - close(broker.stop) + if broker == nil { + return + } + broker.stopOnce.Do(func() { + close(broker.stop) + }) } // Subscribe - Generate a new subscription channel @@ -329,8 +336,12 @@ func (broker *eventBroker) Subscribe() (chan Event, error) { case <-broker.stop: return nil, ErrEventBrokerUnavailable } - <-ready - return events, nil + select { + case <-ready: + return events, nil + case <-broker.stop: + return nil, ErrEventBrokerUnavailable + } } // Unsubscribe - Remove a subscription channel @@ -402,7 +413,6 @@ func NewBroker() *eventBroker { 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 db49be9df..8aa036fb4 100644 --- a/server/internal/core/event_runtime_test.go +++ b/server/internal/core/event_runtime_test.go @@ -77,6 +77,24 @@ func TestEventBrokerRunStopClosesSubscribers(t *testing.T) { } } +func TestEventBrokerStopIsIdempotent(t *testing.T) { + broker := newTestBroker() + broker.Stop() + broker.Stop() +} + +func TestEventBrokerRejectsSlowSubscriber(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 !errors.Is(err, ErrEventSubscriberSlow) { + t.Fatalf("dispatch error = %v, want ErrEventSubscriberSlow", err) + } +} + func TestEventBrokerTryPublishReturnsUnavailableWhenStopped(t *testing.T) { broker := newTestBroker() broker.managed.Store(true) diff --git a/server/internal/core/safe_test.go b/server/internal/core/safe_test.go index a7ad1568f..177d2f544 100644 --- a/server/internal/core/safe_test.go +++ b/server/internal/core/safe_test.go @@ -23,7 +23,6 @@ func newTestBroker() *eventBroker { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } } diff --git a/server/internal/core/session_lifecycle_test.go b/server/internal/core/session_lifecycle_test.go index 496362dcc..71c84382c 100644 --- a/server/internal/core/session_lifecycle_test.go +++ b/server/internal/core/session_lifecycle_test.go @@ -30,7 +30,6 @@ func TestLifecycle_TickerMarksDead(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -93,7 +92,6 @@ func TestLifecycle_TickerKeepsAlive(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -172,7 +170,6 @@ func TestLifecycle_RegisterCheckinDeadReborn(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -284,7 +281,6 @@ func TestLifecycle_ConcurrentCheckinAndTicker(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -368,7 +364,6 @@ func TestLifecycle_MixedAliveDeadSessions(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -471,7 +466,6 @@ func TestEdge_RegisterLongSilenceThenReborn(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -599,7 +593,6 @@ func TestEdge_ConcurrentRemove(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -791,7 +784,6 @@ func TestEdge_RapidFlapping(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker @@ -941,7 +933,6 @@ func TestEdge_MultipleDeathsInSameTicker(t *testing.T) { send: make(chan Event, eventBufSize), notifier: inotify.NewNotifier(), cache: NewMessageCache(eventBufSize), - lock: &sync.Mutex{}, } oldBroker := EventBroker EventBroker = broker diff --git a/server/rpc/rpc-events.go b/server/rpc/rpc-events.go index eb1c5acd2..1601fc15d 100644 --- a/server/rpc/rpc-events.go +++ b/server/rpc/rpc-events.go @@ -33,7 +33,10 @@ func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsSe 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 { From b86a685e2ec75826eff33588b5f29845be525b99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:29:50 +0800 Subject: [PATCH 13/29] fix(listener): harden forward lifecycle cleanup --- server/internal/core/forward.go | 50 ++++- server/internal/core/forward_runtime_test.go | 101 +++++++++ server/listener/forward.go | 194 ++++++++++++++++-- .../forward_backpressure_audit_test.go | 2 +- .../forward_closed_stream_audit_test.go | 38 ++++ server/listener/forward_rollback_test.go | 136 +++++++++++- server/listener/forward_runtime_test.go | 134 +++++++++++- server/listener/http.go | 15 +- server/listener/listener.go | 10 +- server/listener/tcp.go | 27 ++- 10 files changed, 644 insertions(+), 63 deletions(-) diff --git a/server/internal/core/forward.go b/server/internal/core/forward.go index bc5e59df3..57377e13e 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,18 @@ 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) +} + +func (f *forwarders) removeIfSame(id string, fw *Forward) error { + if fw == nil { + return 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) { @@ -213,16 +223,40 @@ func (f *Forward) shutdown() { // 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 { - f.abortErr = stream.CloseSend() + 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 { - <-f.handlerDone + select { + case <-f.handlerDone: + case <-ctx.Done(): + f.abortErr = errors.Join(f.abortErr, fmt.Errorf("%w: %v", ErrForwardShutdownTimeout, ctx.Err())) + } } }) return f.abortErr diff --git a/server/internal/core/forward_runtime_test.go b/server/internal/core/forward_runtime_test.go index 933ab9f70..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" @@ -41,6 +42,23 @@ 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") @@ -108,6 +126,54 @@ func TestForwardAbortClosesStreamCancelsContextAndStopsHandler(t *testing.T) { } } +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{ @@ -153,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/listener/forward.go b/server/listener/forward.go index 5a466879c..6c3aea6ea 100644 --- a/server/listener/forward.go +++ b/server/listener/forward.go @@ -23,6 +23,14 @@ 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") +) + type ForwardListener struct { lns *listener server *grpc.Server @@ -129,7 +137,10 @@ func (s *forwardListenerService) TaskStream(stream forwardrpc.ForwardListener_Ta if err != nil { return err } - local := s.registry.get(listenerID, pipelineID) + local, err := s.registry.open(listenerID, pipelineID) + if err != nil { + return status.Error(codes.Unavailable, err.Error()) + } local.attach(stream) return local.serve(stream) } @@ -140,19 +151,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, @@ -170,14 +190,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, @@ -195,49 +221,131 @@ 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) + } + if r.poisoned[key] != nil { + return nil, fmt.Errorf("%w: %s", ErrForwardStreamPoisoned, key) } - stream := &forwardLocalStream{ + 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 @@ -249,6 +357,8 @@ type forwardLocalStream struct { slotsOnce sync.Once eventSlots chan struct{} remoteSend sync.WaitGroup + closeErr error + drained chan struct{} } func (s *forwardLocalStream) Send(resp *clientpb.SpiteResponse) error { @@ -267,6 +377,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: @@ -337,7 +454,16 @@ 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) @@ -345,23 +471,47 @@ func (s *forwardLocalStream) close() { close(s.done) } s.eventMu.Unlock() - s.remoteSend.Wait() + + 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 <-slots: s.eventMu.Lock() - if s.closed.Load() { + 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 @@ -369,6 +519,8 @@ func (s *forwardLocalStream) sendEvent(event *clientpb.SpiteRequest) error { return nil case <-s.done: return io.ErrClosedPipe + case <-ctx.Done(): + return ctx.Err() case <-timer.C: return fmt.Errorf("forward stream %s:%s event queue full", s.listenerID, s.pipelineID) } diff --git a/server/listener/forward_backpressure_audit_test.go b/server/listener/forward_backpressure_audit_test.go index 8ca109069..be370ff41 100644 --- a/server/listener/forward_backpressure_audit_test.go +++ b/server/listener/forward_backpressure_audit_test.go @@ -46,7 +46,7 @@ func (s *auditRequestTaskStream) SendMsg(interface{}) error { return nil } func (s *auditRequestTaskStream) RecvMsg(interface{}) error { return nil } func TestAuditForwardRegistryUsesBoundedQueues(t *testing.T) { - stream := newForwardStreamRegistry().get("audit-listener", "audit-pipeline") + stream := mustOpenForwardLocalStream(t, newForwardStreamRegistry(), "audit-listener", "audit-pipeline") if got := cap(stream.requests); got != 255 { t.Fatalf("request queue capacity = %d, want 255", got) } diff --git a/server/listener/forward_closed_stream_audit_test.go b/server/listener/forward_closed_stream_audit_test.go index f3509e7ce..7c4ca8fb8 100644 --- a/server/listener/forward_closed_stream_audit_test.go +++ b/server/listener/forward_closed_stream_audit_test.go @@ -146,3 +146,41 @@ func TestForwardCloseWaitsForInFlightRemoteSend(t *testing.T) { 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_rollback_test.go b/server/listener/forward_rollback_test.go index 17f8a6667..701acf468 100644 --- a/server/listener/forward_rollback_test.go +++ b/server/listener/forward_rollback_test.go @@ -21,6 +21,7 @@ var errListenerInitialization = errors.New("listener initialization failed") type rollbackForwardStream struct { closeCalls atomic.Int32 + closeErr error } func (*rollbackForwardStream) Send(*clientpb.SpiteResponse) error { return nil } @@ -29,7 +30,7 @@ func (*rollbackForwardStream) Recv() (*clientpb.SpiteRequest, error) { } func (s *rollbackForwardStream) CloseSend() error { s.closeCalls.Add(1) - return nil + return s.closeErr } type rollbackPipelineRPC struct { @@ -40,6 +41,7 @@ type rollbackListener struct { closed chan struct{} closeOnce sync.Once closeCalls atomic.Int32 + closeErr error } func newRollbackListener() *rollbackListener { @@ -55,7 +57,7 @@ func (l *rollbackListener) Close() error { l.closeCalls.Add(1) close(l.closed) }) - return nil + return l.closeErr } func (*rollbackListener) Addr() net.Addr { return rollbackAddr("rollback") } @@ -117,6 +119,63 @@ func TestTCPPipelineRollsBackForwardWhenListenFails(t *testing.T) { 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 } @@ -134,9 +193,68 @@ func TestHTTPPipelineRollsBackForwardWhenListenFails(t *testing.T) { 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 @@ -151,7 +269,11 @@ func TestTCPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { if err != nil { t.Fatalf("NewTcpPipeline failed: %v", err) } - assertForwardRolledBack(t, pipeline, stream, pipeline.Start) + 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) } @@ -159,7 +281,9 @@ func TestTCPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { 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 @@ -174,7 +298,11 @@ func TestHTTPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { if err != nil { t.Fatalf("NewHttpPipeline failed: %v", err) } - assertForwardRolledBack(t, pipeline, stream, pipeline.Start) + 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) } diff --git a/server/listener/forward_runtime_test.go b/server/listener/forward_runtime_test.go index 70d35b591..3f03491f4 100644 --- a/server/listener/forward_runtime_test.go +++ b/server/listener/forward_runtime_test.go @@ -22,6 +22,15 @@ import ( "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 +199,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 +207,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 +219,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 +233,121 @@ 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 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 +373,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 203f91891..6b521a241 100644 --- a/server/listener/http.go +++ b/server/listener/http.go @@ -126,13 +126,10 @@ func (pipeline *HTTPPipeline) Close() error { if ln == nil { return nil } - if err := ln.Close(); err != nil && !errors.Is(err, net.ErrClosed) { - return err - } - return nil + return closePipelineListener(ln) } -func (pipeline *HTTPPipeline) Start() error { +func (pipeline *HTTPPipeline) Start() (err error) { if !pipeline.beginStart() { return nil } @@ -146,7 +143,7 @@ func (pipeline *HTTPPipeline) Start() error { committed := false defer func() { if !committed { - _ = forward.Abort() + err = errors.Join(err, forward.Abort()) pipeline.abortStart() } }() @@ -162,8 +159,7 @@ func (pipeline *HTTPPipeline) Start() error { if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable && pipeline.TLSConfig.Cert != nil { err := httpStartWithCmux(pipeline, ln, mux) if err != nil { - _ = ln.Close() - return err + return errors.Join(err, closePipelineListener(ln)) } } else { // 非 TLS 模式，使用原有逻辑 @@ -176,8 +172,7 @@ func (pipeline *HTTPPipeline) Start() error { }, pipeline.runtimeErrorHandler("serve loop")) } if !pipeline.commitStart(ln, forward) { - _ = ln.Close() - return nil + return closePipelineListener(ln) } committed = true diff --git a/server/listener/listener.go b/server/listener/listener.go index 9ae6ef982..ade9ca6d0 100644 --- a/server/listener/listener.go +++ b/server/listener/listener.go @@ -598,10 +598,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 { @@ -770,7 +774,7 @@ 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 diff --git a/server/listener/tcp.go b/server/listener/tcp.go index 7915e9273..43a4089fd 100644 --- a/server/listener/tcp.go +++ b/server/listener/tcp.go @@ -102,14 +102,10 @@ func (pipeline *TCPPipeline) Close() error { if ln == nil { return nil } - err := ln.Close() - if err != nil && !errors.Is(err, net.ErrClosed) { - return err - } - return nil + return closePipelineListener(ln) } -func (pipeline *TCPPipeline) Start() error { +func (pipeline *TCPPipeline) Start() (err error) { if !pipeline.beginStart() { return nil } @@ -123,7 +119,7 @@ func (pipeline *TCPPipeline) Start() error { committed := false defer func() { if !committed { - _ = forward.Abort() + err = errors.Join(err, forward.Abort()) pipeline.abortStart() } }() @@ -133,8 +129,7 @@ func (pipeline *TCPPipeline) Start() error { return err } if !pipeline.commitStart(ln, forward) { - _ = ln.Close() - return nil + return closePipelineListener(ln) } committed = true logs.Log.Infof("pipeline.tcp - start host=%s port=%d parser=%s tls=%t", @@ -218,8 +213,7 @@ func (pipeline *TCPPipeline) handler() (net.Listener, error) { if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable { cmuxListener, err := pipeline.handleWithCmux(ln) if err != nil { - _ = ln.Close() - return nil, err + return nil, errors.Join(err, closePipelineListener(ln)) } return cmuxListener, nil } @@ -231,6 +225,17 @@ func (pipeline *TCPPipeline) handler() (net.Listener, error) { return ln, nil } +func closePipelineListener(ln net.Listener) error { + if ln == nil { + return nil + } + err := ln.Close() + if errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + // handleWithCmux 使用 cmux 实现 TLS 和非 TLS 的端口复用 func (pipeline *TCPPipeline) handleWithCmux(ln net.Listener) (net.Listener, error) { var tlsConfig *tls.Config From c1cd467197320496773306300b49d0d3208ffb45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:43:03 +0800 Subject: [PATCH 14/29] fix(events): make startup replay atomic --- client/core/server.go | 118 +++++++-- client/core/server_state_test.go | 11 + external/IoM-go | 2 +- server/internal/core/event.go | 71 +++++- server/internal/core/event_runtime_test.go | 86 +++++++ server/rpc/rpc-events.go | 31 ++- server/rpc/rpc_events_runtime_test.go | 282 +++++++++++++++++++++ 7 files changed, 567 insertions(+), 34 deletions(-) create mode 100644 server/rpc/rpc_events_runtime_test.go diff --git a/client/core/server.go b/client/core/server.go index c6ef22d55..c993dffaf 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,16 +142,11 @@ 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 - } - for _, event := range events.GetEvents() { - if suppressStartupOutput { - continue - } - ser.handleEventOutput(event) + ser := &Server{ + ServerState: s, + taskMessages: make(map[string]string), + suppressStartupOutput: suppressStartupOutput, + startupHistoryPending: true, } return ser, nil } @@ -430,13 +444,79 @@ 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 } - 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 } + 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 + } + if !s.suppressStartupOutput { + for _, event := range history.GetEvents() { + s.handleEventOutput(event) + } + } + } + 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 { + s.handleEventOutput(event) + } + } + s.markStartupHistoryComplete() + } if err := s.Update(); err != nil { client.Log.Warnf("failed to refresh server state: %v\n", err) } diff --git a/client/core/server_state_test.go b/client/core/server_state_test.go index 055e0c17e..d816d9e11 100644 --- a/client/core/server_state_test.go +++ b/client/core/server_state_test.go @@ -85,6 +85,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/external/IoM-go b/external/IoM-go index 1f1bf8df3..dbb6ae8e3 160000 --- a/external/IoM-go +++ b/external/IoM-go @@ -1 +1 @@ -Subproject commit 1f1bf8df33972dcaf46447da4261bdafc029f9b4 +Subproject commit dbb6ae8e3ffac1d3f1f8a1c11088fbe64ff125ce diff --git a/server/internal/core/event.go b/server/internal/core/event.go index 8c60ccb4e..15757df79 100644 --- a/server/internal/core/event.go +++ b/server/internal/core/event.go @@ -222,8 +222,9 @@ func (event *Event) ToProtobuf() *clientpb.Event { } type eventSubscription struct { - events chan Event - ready chan struct{} + events chan Event + ready chan struct{} + history chan []Event } type eventBroker struct { @@ -240,6 +241,8 @@ type eventBroker struct { managed atomic.Bool startOnce sync.Once stopOnce sync.Once + publishMu sync.Mutex + stopped bool } func (broker *eventBroker) run() error { @@ -254,13 +257,27 @@ func (broker *eventBroker) run() error { for { select { case <-broker.stop: + broker.cacheAcceptedEvents() return nil 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 event := <-broker.publish: + if event.Important { + eventCopy := event + broker.cache.Add(&eventCopy) + } if event.Important { logs.Log.Infof("event.%s - %s", event.EventType, event.String()) } else if event.EventType != consts.EventHeartbeat { @@ -277,6 +294,20 @@ func (broker *eventBroker) run() error { } } +func (broker *eventBroker) cacheAcceptedEvents() { + for { + select { + case event := <-broker.publish: + if event.Important { + eventCopy := event + broker.cache.Add(&eventCopy) + } + default: + return + } + } +} + func closeEventSubscriber(sub chan Event) { defer func() { _ = recover() @@ -323,24 +354,43 @@ func (broker *eventBroker) Stop() { 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, 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{}) + var history chan []Event + if includeHistory { + history = make(chan []Event, 1) + } select { - case broker.subscribe <- eventSubscription{events: events, ready: ready}: + case broker.subscribe <- eventSubscription{events: events, ready: ready, history: history}: case <-broker.stop: - return nil, ErrEventBrokerUnavailable + return nil, nil, ErrEventBrokerUnavailable } select { case <-ready: - return events, nil + if history != nil { + return events, <-history, nil + } + return events, nil, nil case <-broker.stop: - return nil, ErrEventBrokerUnavailable + return nil, nil, ErrEventBrokerUnavailable } } @@ -367,17 +417,22 @@ 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: + broker.publishMu.Unlock() return ErrEventBrokerQueueFull } + broker.publishMu.Unlock() if event.IsNotify { broker.Notify(event) } diff --git a/server/internal/core/event_runtime_test.go b/server/internal/core/event_runtime_test.go index 8aa036fb4..8d8e38ae6 100644 --- a/server/internal/core/event_runtime_test.go +++ b/server/internal/core/event_runtime_test.go @@ -83,6 +83,33 @@ func TestEventBrokerStopIsIdempotent(t *testing.T) { 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 TestEventBrokerRejectsSlowSubscriber(t *testing.T) { broker := newTestBroker() sub := make(chan Event, eventBufSize) @@ -229,6 +256,65 @@ func TestEventBrokerSubscribeWaitsForRegistration(t *testing.T) { } } +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() diff --git a/server/rpc/rpc-events.go b/server/rpc/rpc-events.go index 1601fc15d..d71506497 100644 --- a/server/rpc/rpc-events.go +++ b/server/rpc/rpc-events.go @@ -10,16 +10,23 @@ import ( "github.com/chainreactors/logs" "github.com/chainreactors/malice-network/server/internal/core" "github.com/chainreactors/malice-network/server/internal/db" + "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" ) func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsServer) error { - events, err := core.EventBroker.Subscribe() - if err != nil { - return err + 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 = core.EventBroker.SubscribeWithHistory() + } else { + events, err = core.EventBroker.Subscribe() } - if err := stream.SendHeader(metadata.Pairs(consts.EventStreamReadyHeader, "true")); err != nil { - core.EventBroker.Unsubscribe(events) + if err != nil { return err } clientID := core.GetCurrentID() @@ -28,6 +35,18 @@ func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsSe core.Clients.Remove(int(clientID)) core.EventBroker.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 { @@ -35,7 +54,7 @@ func (rpc *Server) Events(_ *clientpb.Empty, stream clientrpc.MaliceRPC_EventsSe return nil case event, ok := <-events: if !ok { - return nil + return status.Error(codes.ResourceExhausted, "event subscriber fell behind; reconnect and resynchronize") } pb := event.ToProtobuf() err := stream.Send(pb) diff --git a/server/rpc/rpc_events_runtime_test.go b/server/rpc/rpc_events_runtime_test.go new file mode 100644 index 000000000..9c8142986 --- /dev/null +++ b/server/rpc/rpc_events_runtime_test.go @@ -0,0 +1,282 @@ +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/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +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 TestEventsReportsSubscriberClosureAsResyncError(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 status.Code(err) != codes.ResourceExhausted { + t.Fatalf("Events status = %v, want ResourceExhausted: %v", status.Code(err), err) + } + case <-time.After(2 * time.Second): + t.Fatal("Events did not report subscriber closure") + } +} + +func TestEventsReportsRealSubscriberOverflow(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 = 128 + 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) + select { + case err := <-result: + if status.Code(err) != codes.ResourceExhausted { + t.Fatalf("Events status = %v, want ResourceExhausted: %v", status.Code(err), err) + } + case <-time.After(2 * time.Second): + t.Fatal("Events did not report real subscriber overflow") + } + if len(stream.events) == 0 { + t.Fatal("overflow test did not deliver any queued events") + } +} From a6b64132ad8cbc3f4160757c1381a856ae15f815 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:50:41 +0800 Subject: [PATCH 15/29] fix(transport): bound payload read duration --- server/internal/core/connection.go | 30 +++- .../internal/core/connection_runtime_test.go | 136 ++++++++++++++++++ server/internal/stream/conn.go | 33 +++++ server/internal/stream/conn_test.go | 47 ++++++ server/listener/http.go | 5 + server/listener/http_pipeline_e2e_test.go | 48 ++++++- 6 files changed, 294 insertions(+), 5 deletions(-) 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/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/listener/http.go b/server/listener/http.go index 6b521a241..e51f402df 100644 --- a/server/listener/http.go +++ b/server/listener/http.go @@ -15,6 +15,7 @@ import ( "net/http" "strconv" "sync" + "time" "github.com/chainreactors/logs" "github.com/chainreactors/malice-network/server/internal/certutils" @@ -393,6 +394,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 } 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() From 3815162d6ade1a7d949ec21bcc21bf96b5a9504c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:56:18 +0800 Subject: [PATCH 16/29] test(ci): expand concurrency and fuzz coverage --- .github/workflows/ci.yaml | 42 +++++++++++++++++++++--------- .github/workflows/realimplant.yaml | 3 +++ 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 1376228be..f282b4d6e 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -25,8 +25,10 @@ jobs: - 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 @@ -59,18 +61,28 @@ jobs: - 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 @@ -89,8 +101,10 @@ jobs: - 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 @@ -122,8 +136,10 @@ jobs: - 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/realimplant.yaml b/.github/workflows/realimplant.yaml index 73bc7f3af..dc9ecf649 100644 --- a/.github/workflows/realimplant.yaml +++ b/.github/workflows/realimplant.yaml @@ -41,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: | From 47c837f9c42ddf0c591f746df441c28400d80322 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:40:53 +0800 Subject: [PATCH 17/29] fix(session): synchronize mutable runtime state --- server/internal/core/session.go | 119 ++++++++++++++++++ .../core/session_state_accessors_test.go | 80 ++++++++++++ server/rpc/generic.go | 9 +- server/rpc/rpc-addon.go | 42 ++----- server/rpc/rpc-forward-listener.go | 3 +- server/rpc/rpc-implant.go | 17 +-- server/rpc/rpc-listener.go | 3 +- server/rpc/rpc_addon_runtime_test.go | 24 ++-- .../rpc/rpc_implant_session_handlers_test.go | 5 +- 9 files changed, 241 insertions(+), 61 deletions(-) create mode 100644 server/internal/core/session_state_accessors_test.go diff --git a/server/internal/core/session.go b/server/internal/core/session.go index d6bff6081..210a4527c 100644 --- a/server/internal/core/session.go +++ b/server/internal/core/session.go @@ -691,6 +691,125 @@ func (s *Session) SetGroup(group string) { 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 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/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/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-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..629028a7e 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()}, 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_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_implant_session_handlers_test.go b/server/rpc/rpc_implant_session_handlers_test.go index c985d2f7b..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 { From 1d0b9f9dde2d39ebea80b17451e74378db0b6239 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:43:24 +0800 Subject: [PATCH 18/29] fix(transport): bound initial protocol reads --- server/internal/stream/peekconn.go | 36 +++++++-- .../internal/stream/peekconn_deadline_test.go | 80 +++++++++++++++++++ server/listener/cmux.go | 14 +++- server/listener/cmux_server_test.go | 73 +++++++++++++++++ 4 files changed, 197 insertions(+), 6 deletions(-) create mode 100644 server/internal/stream/peekconn_deadline_test.go create mode 100644 server/listener/cmux_server_test.go 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") + } +} From 257e7ed84a7d624063ce201dce569b01ad78bc8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:43:58 +0800 Subject: [PATCH 19/29] fix(forward): enforce single active task stream --- server/listener/forward.go | 39 +++++++++++-- server/listener/forward_runtime_test.go | 76 +++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 5 deletions(-) diff --git a/server/listener/forward.go b/server/listener/forward.go index 6c3aea6ea..b4bcd9433 100644 --- a/server/listener/forward.go +++ b/server/listener/forward.go @@ -26,9 +26,10 @@ import ( 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") + 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 { @@ -141,7 +142,14 @@ func (s *forwardListenerService) TaskStream(stream forwardrpc.ForwardListener_Ta if err != nil { return status.Error(codes.Unavailable, err.Error()) } - local.attach(stream) + 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) } @@ -359,6 +367,8 @@ type forwardLocalStream struct { remoteSend sync.WaitGroup closeErr error drained chan struct{} + attachMu sync.Mutex + attached bool } func (s *forwardLocalStream) Send(resp *clientpb.SpiteResponse) error { @@ -396,7 +406,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()) diff --git a/server/listener/forward_runtime_test.go b/server/listener/forward_runtime_test.go index 3f03491f4..953fe2acf 100644 --- a/server/listener/forward_runtime_test.go +++ b/server/listener/forward_runtime_test.go @@ -19,6 +19,9 @@ 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" ) @@ -326,6 +329,79 @@ func TestForwardStreamRegistryBlocksReuseUntilTimedOutSendDrains(t *testing.T) { } } +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} From 97b7fe2dcbc250b8bac82b8b904b84d519cd52d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:52:41 +0800 Subject: [PATCH 20/29] fix(rem): harden runtime cancellation and cleanup --- server/listener/rem.go | 174 +++++++++++++++++-------- server/listener/rem_runtime_test.go | 188 +++++++++++++++++++++++++++- 2 files changed, 309 insertions(+), 53 deletions(-) diff --git a/server/listener/rem.go b/server/listener/rem.go index 90f6da03c..176afe317 100644 --- a/server/listener/rem.go +++ b/server/listener/rem.go @@ -34,6 +34,25 @@ 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 +) + +type remStartState struct { + done chan struct{} + cleanupErr error +} + func NewRem(rpc listenerrpc.ListenerRPCClient, pipeline *clientpb.Pipeline) (*REM, error) { remConfig := pipeline.GetRem() var conURL string @@ -62,8 +81,9 @@ func NewRem(rpc listenerrpc.ListenerRPCClient, pipeline *clientpb.Pipeline) (*RE type REM struct { stateMu sync.RWMutex starting bool - startDone chan struct{} - stopCh chan struct{} + startState *remStartState + runCtx context.Context + runCancel context.CancelFunc healthInterval time.Duration con *rem.RemConsole rpc listenerrpc.ListenerRPCClient @@ -87,19 +107,19 @@ func (rem *REM) Start() error { err := remConsoleListen(rem.con) if err != nil { - rem.abortStart() + rem.abortStart(nil) return err } if !rem.enabled() { - _ = remConsoleClose(rem.con) - rem.abortStart() - return nil + cleanupErr := remConsoleClose(rem.con) + rem.abortStart(cleanupErr) + return cleanupErr } logs.Log.Important(rem.con.Link()) if !rem.commitStart() { - _ = remConsoleClose(rem.con) - rem.abortStart() - return nil + cleanupErr := remConsoleClose(rem.con) + rem.abortStart(cleanupErr) + return cleanupErr } return nil } @@ -204,20 +224,21 @@ func (rem *REM) Close() error { rem.stateMu.Lock() wasActive := rem.Enable rem.Enable = false - stopCh := rem.stopCh - rem.stopCh = nil + cancel := rem.runCancel + rem.runCancel = nil + rem.runCtx = nil if rem.starting { - done := rem.startDone + startState := rem.startState rem.stateMu.Unlock() - if stopCh != nil { - close(stopCh) + if cancel != nil { + cancel() } - <-done - return nil + <-startState.done + return startState.cleanupErr } rem.stateMu.Unlock() - if stopCh != nil { - close(stopCh) + if cancel != nil { + cancel() } if !wasActive || rem.con == nil { return nil @@ -235,7 +256,7 @@ func (rem *REM) beginStart() bool { for { rem.stateMu.Lock() if rem.starting { - done := rem.startDone + done := rem.startState.done rem.stateMu.Unlock() <-done continue @@ -246,27 +267,29 @@ func (rem *REM) beginStart() bool { } rem.Enable = true rem.starting = true - rem.startDone = make(chan struct{}) - rem.stopCh = make(chan struct{}) + rem.startState = &remStartState{done: make(chan struct{})} + rem.runCtx, rem.runCancel = context.WithCancel(context.Background()) rem.stateMu.Unlock() return true } } -func (rem *REM) abortStart() { +func (rem *REM) abortStart(cleanupErr error) { rem.stateMu.Lock() rem.starting = false rem.Enable = false - done := rem.startDone - rem.startDone = nil - stopCh := rem.stopCh - rem.stopCh = nil - if done != nil { - close(done) + 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 stopCh != nil { - close(stopCh) + if cancel != nil { + cancel() } } @@ -276,58 +299,84 @@ func (rem *REM) commitStart() bool { if !rem.Enable { return false } - core.GoGuarded("rem-accept:"+rem.Name, rem.acceptLoop, rem.runtimeErrorHandler("accept loop")) - core.GoGuarded("rem-health:"+rem.Name, rem.healthLoop, rem.runtimeErrorHandler("health loop")) + 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 - done := rem.startDone - rem.startDone = nil - close(done) + startState := rem.startState + rem.startState = nil + close(startState.done) return true } -func (rem *REM) acceptLoop() error { +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 := rem.con.Accept() + ag, err := remConsoleAccept(rem.con) if err != nil { - if !rem.enabled() { + 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 - stopCh := rem.stopSignal() for rem.enabled() { - if err := remHealthCheck(rem.rpc, context.Background(), rem.ToProtobuf()); err != nil { + 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 { @@ -363,7 +412,7 @@ func (rem *REM) healthLoop() error { timer := time.NewTimer(interval) select { case <-timer.C: - case <-stopCh: + case <-runCtx.Done(): if !timer.Stop() { select { case <-timer.C: @@ -376,13 +425,36 @@ func (rem *REM) healthLoop() error { return nil } -func (rem *REM) stopSignal() <-chan struct{} { - rem.stateMu.Lock() - defer rem.stateMu.Unlock() - if rem.stopCh == nil { - rem.stopCh = make(chan struct{}) +func (rem *REM) healthCheck(runCtx context.Context) error { + if runCtx == nil { + runCtx = context.Background() } - return rem.stopCh + 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 { diff --git a/server/listener/rem_runtime_test.go b/server/listener/rem_runtime_test.go index 7655a76bf..722dfcdaf 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,186 @@ 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 + defer func() { remConsoleAccept = oldAccept }() + + 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) }() + + 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(100 * time.Millisecond): + t.Fatal("second Accept did not occur within bounded backoff") + } + select { + case <-calls: + case <-time.After(100 * time.Millisecond): + 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: + if err != nil { + t.Fatalf("acceptLoop error = %v", err) + } + if elapsed := time.Since(startedClose); elapsed > 100*time.Millisecond { + t.Fatalf("acceptLoop cancellation took %v", elapsed) + } + case <-time.After(250 * time.Millisecond): + 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", @@ -39,8 +218,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) @@ -103,8 +286,9 @@ func TestREMHealthLoopPublishesDegradedAndRecoveredEvents(t *testing.T) { } done := make(chan error, 1) + runCtx := setREMTestRunContext(rem) go func() { - done <- rem.healthLoop() + done <- rem.healthLoopContext(runCtx) }() deadline := time.After(2 * time.Second) From 6beacf2af15a29024f41be3ba4e83867cec71b27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:00:30 +0800 Subject: [PATCH 21/29] fix(listener): linearize pipeline serving startup --- server/internal/core/forward.go | 18 +- server/listener/forward_rollback_test.go | 77 ++++++++ server/listener/http.go | 51 ++++-- .../listener/lifecycle_linearization_test.go | 168 ++++++++++++++++++ server/listener/tcp.go | 92 ++++++---- 5 files changed, 357 insertions(+), 49 deletions(-) diff --git a/server/internal/core/forward.go b/server/internal/core/forward.go index 57377e13e..1d5868031 100644 --- a/server/internal/core/forward.go +++ b/server/internal/core/forward.go @@ -117,6 +117,18 @@ func (f *forwarders) Remove(id string) error { 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 + } + f.forwarders.CompareAndDelete(id, fw) + return fw.Abort() +} + func (f *forwarders) removeIfSame(id string, fw *Forward) error { if fw == nil { return nil @@ -174,7 +186,6 @@ func NewForward(rpc ForwardClient, pipeline Pipeline) (*Forward, error) { type Forward struct { ctx context.Context cancel context.CancelFunc - count int Pipeline ListenerId string Stream ForwardStream @@ -204,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) } @@ -262,10 +272,6 @@ func (f *Forward) AbortContext(ctx context.Context) error { return f.abortErr } -func (f *Forward) Count() int { - return f.count -} - func (f *Forward) Context(sid string) context.Context { return metadata.NewOutgoingContext(f.ctx, metadata.Pairs( "session_id", sid, diff --git a/server/listener/forward_rollback_test.go b/server/listener/forward_rollback_test.go index 701acf468..73c648e29 100644 --- a/server/listener/forward_rollback_test.go +++ b/server/listener/forward_rollback_test.go @@ -102,6 +102,17 @@ func assertRollbackState(t *testing.T, pipeline core.Pipeline, stream *rollbackF } } +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 } @@ -307,3 +318,69 @@ func TestHTTPPipelineRollsBackForwardWhenCmuxFails(t *testing.T) { 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/http.go b/server/listener/http.go index e51f402df..b02c0ac20 100644 --- a/server/listener/http.go +++ b/server/listener/http.go @@ -142,9 +142,14 @@ func (pipeline *HTTPPipeline) Start() (err error) { } forward.ListenerId = pipeline.ListenerID committed := false + registered := false defer func() { if !committed { - err = errors.Join(err, forward.Abort()) + if registered { + err = errors.Join(err, pipeline.rollbackCommittedStart(forward)) + } else { + err = errors.Join(err, forward.Abort()) + } pipeline.abortStart() } }() @@ -157,13 +162,16 @@ func (pipeline *HTTPPipeline) Start() (err 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 := httpStartWithCmux(pipeline, ln, mux) if err != nil { - return errors.Join(err, closePipelineListener(ln)) + return err } } else { - // 非 TLS 模式，使用原有逻辑 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) { @@ -172,10 +180,11 @@ func (pipeline *HTTPPipeline) Start() (err error) { return nil }, pipeline.runtimeErrorHandler("serve loop")) } - if !pipeline.commitStart(ln, forward) { - return closePipelineListener(ln) + 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) @@ -229,6 +238,33 @@ func (pipeline *HTTPPipeline) commitStart(ln net.Listener, forward *core.Forward } 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() @@ -241,11 +277,6 @@ func (pipeline *HTTPPipeline) commitStart(ln net.Listener, forward *core.Forward dispatchForwardTaskRequest("http", pipeline.Name, msg) } }, pipeline.runtimeErrorHandler("forward recv loop")) - pipeline.starting = false - done := pipeline.startDone - pipeline.startDone = nil - close(done) - return true } // startWithCmux 使用 cmux 实现 HTTP TLS 和非 TLS 的端口复用 diff --git a/server/listener/lifecycle_linearization_test.go b/server/listener/lifecycle_linearization_test.go index ffefc730f..c6ee4f0d4 100644 --- a/server/listener/lifecycle_linearization_test.go +++ b/server/listener/lifecycle_linearization_test.go @@ -2,7 +2,11 @@ package listener import ( "context" + "crypto/tls" "errors" + "net" + "net/http" + "sync" "sync/atomic" "testing" "time" @@ -25,6 +29,42 @@ type linearizedStartRPC 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) @@ -109,6 +149,134 @@ func TestHTTPPipelineCloseLinearizesWithInFlightStart(t *testing.T) { 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 diff --git a/server/listener/tcp.go b/server/listener/tcp.go index 43a4089fd..420cbb0d1 100644 --- a/server/listener/tcp.go +++ b/server/listener/tcp.go @@ -117,21 +117,34 @@ func (pipeline *TCPPipeline) Start() (err error) { } forward.ListenerId = pipeline.ListenerID committed := false + registered := false defer func() { if !committed { - err = errors.Join(err, forward.Abort()) + if registered { + err = errors.Join(err, pipeline.rollbackCommittedStart(forward)) + } else { + err = errors.Join(err, forward.Abort()) + } pipeline.abortStart() } }() - 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) return nil @@ -184,6 +197,23 @@ func (pipeline *TCPPipeline) commitStart(ln net.Listener, forward *core.Forward) } pipeline.ln = ln core.Forwarders.Add(forward) + return true +} + +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() @@ -196,33 +226,38 @@ func (pipeline *TCPPipeline) commitStart(ln net.Listener, forward *core.Forward) dispatchForwardTaskRequest("tcp", pipeline.Name, msg) } }, pipeline.runtimeErrorHandler("forward recv loop")) - pipeline.starting = false - done := pipeline.startDone - pipeline.startDone = nil - close(done) - return true } -func (pipeline *TCPPipeline) handler() (net.Listener, error) { +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, err + return nil, nil, err } - // 如果启用了 TLS，使用 cmux 实现 TLS 和非 TLS 的端口复用 - if pipeline.TLSConfig != nil && pipeline.TLSConfig.Enable { - cmuxListener, err := pipeline.handleWithCmux(ln) + var tlsConfig *tls.Config + 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) + } else { + tlsConfig, err = certutils.GetTlsConfig(pipeline.TLSConfig.Cert) + } if err != nil { - return nil, errors.Join(err, closePipelineListener(ln)) + return nil, nil, errors.Join(err, closePipelineListener(ln)) } - return cmuxListener, nil } + return ln, tlsConfig, nil +} - // 非 TLS 模式，使用原有逻辑 +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 ln, nil + return nil } func closePipelineListener(ln net.Listener) error { @@ -236,23 +271,14 @@ func closePipelineListener(ln net.Listener) error { return err } -// handleWithCmux 使用 cmux 实现 TLS 和非 TLS 的端口复用 -func (pipeline *TCPPipeline) handleWithCmux(ln net.Listener) (net.Listener, error) { - var tlsConfig *tls.Config - if pipeline.TLSConfig.Cert != nil { - var err error - 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) - } else { - tlsConfig, err = certutils.GetTlsConfig(pipeline.TLSConfig.Cert) - } - if err != nil { - return nil, err - } - } - - return tcpStartCmux(ln, tlsConfig, pipeline.HandleConnection, pipeline.runtimeErrorHandler("cmux")) +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 模式) From 42e094cfaf35844ad377fe371f557679d8400b79 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:00:38 +0800 Subject: [PATCH 22/29] fix(listener): replace failed pipeline runtimes --- server/listener/listener.go | 21 +++- .../listener_runtime_recovery_test.go | 112 ++++++++++++++++++ 2 files changed, 128 insertions(+), 5 deletions(-) create mode 100644 server/listener/listener_runtime_recovery_test.go diff --git a/server/listener/listener.go b/server/listener/listener.go index ade9ca6d0..006393fe3 100644 --- a/server/listener/listener.go +++ b/server/listener/listener.go @@ -549,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 @@ -781,6 +789,9 @@ func (lns *listener) handleStartRem(job *clientpb.Job) error { } // 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()) } @@ -794,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) + } +} From 5f1bc109bab74f0c34dfe234d604bc85ef3cae99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:05:54 +0800 Subject: [PATCH 23/29] fix(parser): validate snappy blocks before allocation --- helper/utils/compress/compress_test.go | 54 ++++++++++++ helper/utils/compress/snappy_validate.go | 83 +++++++++++++++++++ server/internal/parser/malefic/parser.go | 3 + server/internal/parser/malefic/parser_test.go | 11 +++ 4 files changed, 151 insertions(+) create mode 100644 helper/utils/compress/snappy_validate.go 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/internal/parser/malefic/parser.go b/server/internal/parser/malefic/parser.go index 03d3cb4d8..c9c756390 100644 --- a/server/internal/parser/malefic/parser.go +++ b/server/internal/parser/malefic/parser.go @@ -172,6 +172,9 @@ func (parser *MaleficParser) Parse(buf []byte) (*implantpb.Spites, error) { 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 { diff --git a/server/internal/parser/malefic/parser_test.go b/server/internal/parser/malefic/parser_test.go index 4aa4142bb..5daba1390 100644 --- a/server/internal/parser/malefic/parser_test.go +++ b/server/internal/parser/malefic/parser_test.go @@ -182,6 +182,17 @@ func TestMaleficParser_Parse_RejectsDecodedPayloadAboveLimit(t *testing.T) { } } +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() if _, err := p.Parse(nil); err == nil { From d2e36d749c18606e20043ce0609df4d5c55bbf0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:14:15 -0700 Subject: [PATCH 24/29] fix(core): make session keypair update atomic Merge the read-modify-write of the session key pair under a single write lock so a re-register and a key rotation cannot read the same base snapshot and clobber each other. Rename UpdateKeyPairFieldsAndPushCtrl to UpdateKeyPair (the old full-replacement method is now unused). Co-Authored-By: Claude Opus 4.8 --- server/internal/core/session.go | 40 +++++++++++++++------------------ server/rpc/rpc-implant.go | 2 +- 2 files changed, 19 insertions(+), 23 deletions(-) diff --git a/server/internal/core/session.go b/server/internal/core/session.go index 210a4527c..ebf5854bc 100644 --- a/server/internal/core/session.go +++ b/server/internal/core/session.go @@ -1243,22 +1243,6 @@ func (s *Session) DeleteResp(taskId uint32) { } } -// UpdateKeyPair 更新KeyPair并同步到SecureManager -func (s *Session) UpdateKeyPair(keyPair *clientpb.KeyPair) { - var stored *clientpb.KeyPair - if keyPair != nil { - stored = proto.Clone(keyPair).(*clientpb.KeyPair) - } - s.stateMu.Lock() - s.SessionContext.KeyPair = stored - manager := s.SecureManager - s.stateMu.Unlock() - // 更新SecureManager中的KeyPair引用 - if manager != nil { - manager.UpdateKeyPair(stored) - } -} - // SetKeepalive updates the keepalive state. Returns the previous state. func (s *Session) SetKeepalive(enabled bool) bool { s.keepaliveMu.Lock() @@ -1428,28 +1412,40 @@ 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{} - s.stateMu.RLock() if s.KeyPair != nil { next.PublicKey = s.KeyPair.PublicKey next.PrivateKey = s.KeyPair.PrivateKey } - s.stateMu.RUnlock() if publicKey != "" { next.PublicKey = publicKey } 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 + manager := s.SecureManager + s.stateMu.Unlock() + if manager != nil { + manager.UpdateKeyPair(next) + } s.PushCtrl() } diff --git a/server/rpc/rpc-implant.go b/server/rpc/rpc-implant.go index 629028a7e..f6dedc648 100644 --- a/server/rpc/rpc-implant.go +++ b/server/rpc/rpc-implant.go @@ -550,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 } From 2ede218cb2de706c32204d64c6d3472f42aafaf0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:14:16 -0700 Subject: [PATCH 25/29] fix(listener): bound rem close wait for in-flight start Close waited indefinitely on startState.done while con.Listen ignores cancellation, so a stuck console link could wedge listener shutdown. Cap the wait at remStartCloseTimeout. Co-Authored-By: Claude Opus 4.8 --- server/listener/rem.go | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/server/listener/rem.go b/server/listener/rem.go index 176afe317..65a26a4df 100644 --- a/server/listener/rem.go +++ b/server/listener/rem.go @@ -46,6 +46,12 @@ 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 { @@ -233,8 +239,12 @@ func (rem *REM) Close() error { if cancel != nil { cancel() } - <-startState.done - return startState.cleanupErr + 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 { From 94a3cca5d8e53023a84d6e45e9bb2f7dababde8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:14:16 -0700 Subject: [PATCH 26/29] test(listener): stabilize flaky rem accept-backoff test Stop the accept loop before restoring the seam so a late retry cannot call the default Accept on a nil console and crash the whole package; widen the timing windows to remove flakiness. Co-Authored-By: Claude Opus 4.8 --- server/listener/rem_runtime_test.go | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/server/listener/rem_runtime_test.go b/server/listener/rem_runtime_test.go index 722dfcdaf..6b8e66089 100644 --- a/server/listener/rem_runtime_test.go +++ b/server/listener/rem_runtime_test.go @@ -124,7 +124,6 @@ func TestREMCloseCancelsHealthRPCContextWithDeadline(t *testing.T) { func TestREMAcceptErrorsUseInterruptibleBoundedBackoff(t *testing.T) { oldAccept := remConsoleAccept - defer func() { remConsoleAccept = oldAccept }() calls := make(chan time.Time, 8) remConsoleAccept = func(*remhelper.RemConsole) (*agent.Agent, error) { @@ -139,18 +138,33 @@ func TestREMAcceptErrorsUseInterruptibleBoundedBackoff(t *testing.T) { 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(100 * time.Millisecond): + case <-time.After(2 * time.Second): t.Fatal("second Accept did not occur within bounded backoff") } select { case <-calls: - case <-time.After(100 * time.Millisecond): + case <-time.After(2 * time.Second): t.Fatal("third Accept did not occur within maximum backoff") } @@ -160,13 +174,14 @@ func TestREMAcceptErrorsUseInterruptibleBoundedBackoff(t *testing.T) { } select { case err := <-done: + drained = true if err != nil { t.Fatalf("acceptLoop error = %v", err) } - if elapsed := time.Since(startedClose); elapsed > 100*time.Millisecond { + if elapsed := time.Since(startedClose); elapsed > time.Second { t.Fatalf("acceptLoop cancellation took %v", elapsed) } - case <-time.After(250 * time.Millisecond): + case <-time.After(2 * time.Second): t.Fatal("Close did not interrupt Accept retry backoff") } } From 844875c18f2b244b9c2979b8a54b5113a8d6ae52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:14:16 -0700 Subject: [PATCH 27/29] docs: document event queue-full and forward slot trade-offs Explain that an Important event dropped on a full broker queue is also absent from replay history, and document the eventSlots one-slot-per-in-flight-entry invariant. Co-Authored-By: Claude Opus 4.8 --- server/internal/core/event.go | 4 ++++ server/listener/forward.go | 7 +++++++ 2 files changed, 11 insertions(+) diff --git a/server/internal/core/event.go b/server/internal/core/event.go index ad3709f76..b05317f35 100644 --- a/server/internal/core/event.go +++ b/server/internal/core/event.go @@ -733,6 +733,10 @@ func (broker *eventBroker) TryPublish(event Event) error { 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 } diff --git a/server/listener/forward.go b/server/listener/forward.go index b4bcd9433..e1bcc4bdd 100644 --- a/server/listener/forward.go +++ b/server/listener/forward.go @@ -363,6 +363,13 @@ type forwardLocalStream struct { 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 From 1a7448b03f41135b3fc801fa5ee0b53874f13281 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Tue, 14 Jul 2026 01:40:35 +0800 Subject: [PATCH 28/29] fix(core): keep session keypairs synchronized --- server/internal/core/secure_test.go | 43 +++++++++++++++++++++++++++++ server/internal/core/session.go | 10 ++++--- 2 files changed, 49 insertions(+), 4 deletions(-) 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 ebf5854bc..6804d1194 100644 --- a/server/internal/core/session.go +++ b/server/internal/core/session.go @@ -1441,11 +1441,13 @@ func (s *Session) UpdateKeyPair(publicKey string, privateKey string) { // next is freshly allocated here and shared with no caller, so no defensive // clone is needed before storing it. s.SessionContext.KeyPair = next - manager := s.SecureManager - s.stateMu.Unlock() - if manager != nil { - manager.UpdateKeyPair(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() } From 1f8e5e4c4d1ffd2e4b9c969cee27377ba333725e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=95=E6=AD=A2?= <68958533+h3zh1@users.noreply.github.com> Date: Tue, 14 Jul 2026 01:40:40 +0800 Subject: [PATCH 29/29] test(integration): expect stopped pipelines to remain cached --- client/command/listener/pipeline_integration_test.go | 6 +++--- server/client_server_integration_test.go | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) 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/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") }