feat(secure): support cold-start key exchange with empty pre-shared keys

Allow key exchange to proceed when both pipeline and implant enable
secure mode but neither provides pre-shared Age keys. The first
communication happens in plaintext; once key exchange completes the
session switches to encrypted.

Changes:
- GetKeyPairForSession returns empty KeyPair instead of nil when secure
  is enabled but keys are empty, so the parser stays in secure-aware
  plaintext mode and picks up new keys after PushCtrl.
- initializeSecureManager safely handles nil ImplantKeypair/ServerKeypair
  to prevent panics on pipelines with no pre-shared keys.
- Move ResetCounters into the successful key-exchange callback so a
  failed attempt triggers an immediate retry on the next checkin instead
  of waiting 100 messages.

Author: h3zh1

server/internal/core/connection.go

@@ -118,7 +118,11 @@ func GetKeyPairForSession(sid uint32, secureConfig *implanttypes.SecureConfig) *
 	}
 
 	if publicKey == "" && len(privateCandidates) == 0 {
-		return nil
+		// secure is enabled but no keys established yet (cold start scenario).
+		// Return an empty KeyPair so the parser enters "secure but plaintext" mode;
+		// once key exchange completes and PushCtrl syncs the new keys, the parser
+		// will pick them up on the next GetConnection / WithSecure call.
+		return &clientpb.KeyPair{}
 	}
 
 	return &clientpb.KeyPair{

server/internal/core/connection_keypair_test.go

@@ -0,0 +1,106 @@
+package core
+
+import (
+	"testing"
+
+	"github.com/chainreactors/IoM-go/proto/client/clientpb"
+	"github.com/chainreactors/malice-network/helper/implanttypes"
+)
+
+// TestGetKeyPairForSession_Disabled verifies that nil is returned when
+// secure is not enabled, regardless of key content.
+func TestGetKeyPairForSession_Disabled(t *testing.T) {
+	t.Parallel()
+	cfg := &implanttypes.SecureConfig{Enable: false}
+	kp := GetKeyPairForSession(1, cfg)
+	if kp != nil {
+		t.Fatal("expected nil KeyPair when secure is disabled")
+	}
+}
+
+// TestGetKeyPairForSession_NilConfig verifies nil config returns nil KeyPair.
+func TestGetKeyPairForSession_NilConfig(t *testing.T) {
+	t.Parallel()
+	kp := GetKeyPairForSession(1, nil)
+	if kp != nil {
+		t.Fatal("expected nil KeyPair for nil config")
+	}
+}
+
+// TestGetKeyPairForSession_ColdStart verifies that an empty (non-nil) KeyPair
+// is returned when secure is enabled but no pre-shared keys exist.
+// This is the cold-start scenario where encryption starts after key exchange.
+func TestGetKeyPairForSession_ColdStart(t *testing.T) {
+	t.Parallel()
+	cfg := &implanttypes.SecureConfig{
+		Enable:           true,
+		ServerPublicKey:  "",
+		ServerPrivateKey: "",
+		ImplantPublicKey: "",
+	}
+	kp := GetKeyPairForSession(42, cfg)
+	if kp == nil {
+		t.Fatal("expected non-nil KeyPair for cold-start secure config")
+	}
+	if kp.PublicKey != "" {
+		t.Errorf("PublicKey = %q, want empty", kp.PublicKey)
+	}
+	if kp.PrivateKey != "" {
+		t.Errorf("PrivateKey = %q, want empty", kp.PrivateKey)
+	}
+}
+
+// TestGetKeyPairForSession_WithPreSharedKeys verifies that pre-shared keys
+// are correctly assembled into the KeyPair.
+func TestGetKeyPairForSession_WithPreSharedKeys(t *testing.T) {
+	t.Parallel()
+	cfg := &implanttypes.SecureConfig{
+		Enable:           true,
+		ServerPrivateKey: "server-priv",
+		ImplantPublicKey: "implant-pub",
+	}
+	kp := GetKeyPairForSession(99, cfg)
+	if kp == nil {
+		t.Fatal("expected non-nil KeyPair")
+	}
+	if kp.PublicKey != "implant-pub" {
+		t.Errorf("PublicKey = %q, want %q", kp.PublicKey, "implant-pub")
+	}
+	if kp.PrivateKey != "server-priv" {
+		t.Errorf("PrivateKey = %q, want %q", kp.PrivateKey, "server-priv")
+	}
+}
+
+// TestGetKeyPairForSession_AfterKeyExchange verifies that session-level keys
+// (from key exchange) take priority over pipeline-level pre-shared keys.
+func TestGetKeyPairForSession_AfterKeyExchange(t *testing.T) {
+	t.Parallel()
+	var sid uint32 = 200
+	cfg := &implanttypes.SecureConfig{
+		Enable:           true,
+		ServerPrivateKey: "pipeline-server-priv",
+		ImplantPublicKey: "pipeline-implant-pub",
+	}
+
+	// Simulate key exchange result stored in ListenerSessions
+	ListenerSessions.Add(&clientpb.Session{
+		RawId: sid,
+		KeyPair: &clientpb.KeyPair{
+			PublicKey:  "exchanged-pub",
+			PrivateKey: "exchanged-priv",
+		},
+	})
+	t.Cleanup(func() { ListenerSessions.Remove(sid) })
+
+	kp := GetKeyPairForSession(sid, cfg)
+	if kp == nil {
+		t.Fatal("expected non-nil KeyPair")
+	}
+	if kp.PublicKey != "exchanged-pub" {
+		t.Errorf("PublicKey = %q, want %q", kp.PublicKey, "exchanged-pub")
+	}
+	// PrivateKey should contain both exchanged and pipeline keys
+	if kp.PrivateKey == "" {
+		t.Fatal("PrivateKey should not be empty")
+	}
+}

server/internal/core/secure_test.go

@@ -290,6 +290,44 @@ func TestSecureManager_CounterOverflow(t *testing.T) {
 	}
 }
 
+// TestSecureManager_EmptyKeyPairSession verifies that a SecureManager created
+// with an empty (non-nil, zero-value) KeyPair still works correctly.
+// This is the cold-start scenario where keys will be populated after exchange.
+func TestSecureManager_EmptyKeyPairSession(t *testing.T) {
+	t.Parallel()
+	sess := &Session{
+		ID:   "empty-keypair-session",
+		Type: "test",
+		SessionContext: &client.SessionContext{
+			SessionInfo: &client.SessionInfo{},
+			KeyPair:     &clientpb.KeyPair{PublicKey: "", PrivateKey: ""},
+		},
+	}
+
+	sm := NewSecureSpiteManager(sess)
+	if sm == nil {
+		t.Fatal("NewSecureSpiteManager returned nil")
+	}
+
+	// Counter operations must work normally
+	for i := 0; i < 100; i++ {
+		sm.IncrementCounter()
+	}
+	if !sm.ShouldRotateKey() {
+		t.Error("ShouldRotateKey should be true at counter=100 even with empty keys")
+	}
+
+	// UpdateKeyPair should transition to real keys
+	sm.UpdateKeyPair(&clientpb.KeyPair{
+		PublicKey:  "exchanged-pub",
+		PrivateKey: "exchanged-priv",
+	})
+	sm.ResetCounters()
+	if sm.ShouldRotateKey() {
+		t.Error("ShouldRotateKey should be false after reset")
+	}
+}
+
 // TestSecureManager_NilKeyPairSession verifies creating a SecureManager with
 // a session that has a nil KeyPair does not panic.
 func TestSecureManager_NilKeyPairSession(t *testing.T) {

server/internal/core/session.go

@@ -972,9 +972,19 @@ func (s *Session) initializeSecureManager(req *clientpb.RegisterSession) error {
 	}
 
 	if s.KeyPair == nil {
+		// Populate from pipeline pre-shared keys when available (may be empty
+		// in cold-start scenarios where no keys are distributed ahead of time).
+		pubKey := ""
+		privKey := ""
+		if pipeline.Secure.ImplantKeypair != nil {
+			pubKey = pipeline.Secure.ImplantKeypair.PublicKey
+		}
+		if pipeline.Secure.ServerKeypair != nil {
+			privKey = pipeline.Secure.ServerKeypair.PrivateKey
+		}
 		s.KeyPair = &clientpb.KeyPair{
-			PublicKey:  pipeline.Secure.ImplantKeypair.PublicKey,
-			PrivateKey: pipeline.Secure.ServerKeypair.PrivateKey,
+			PublicKey:  pubKey,
+			PrivateKey: privKey,
 		}
 	}
 

server/real_implant_e2e_test.go

@@ -470,6 +470,76 @@ func TestRealImplantSecureKeyExchangeE2E(t *testing.T) {
 	t.Log("secure key exchange E2E test passed: cold start → key exchange → encrypted commands")
 }
 
+// TestRealImplantFullColdStartE2E verifies key exchange works when the pipeline
+// has NO pre-shared keys at all (neither server nor implant keypair). Both sides
+// start with empty keys; the server triggers key exchange on first registration
+// and the implant generates its keypair during the exchange.
+func TestRealImplantFullColdStartE2E(t *testing.T) {
+	testsupport.RequireRealImplantEnv(t)
+
+	h := testsupport.NewControlPlaneHarness(t)
+	listenerName := fmt.Sprintf("real-fullcold-listener-%d", time.Now().UnixNano())
+	pipelineName := fmt.Sprintf("real-fullcold-pipe-%d", time.Now().UnixNano())
+
+	// Create a fully cold-start pipeline (no keys at all)
+	pipeline := testsupport.NewRealFullColdStartTCPPipeline(t, listenerName, pipelineName)
+	implant := testsupport.NewRealImplant(t, h, pipeline)
+	if err := implant.Start(t); err != nil {
+		t.Fatalf("real full-cold-start implant start failed: %v", err)
+	}
+
+	// Wait for session registration + initial key exchange
+	waitRealActiveConnection(t, implant.SessionID)
+	waitRealPostRegisterCheckin(t, implant.SessionID)
+
+	// Verify session was registered with secure mode
+	runtimeSession := mustRealRuntimeSession(t, implant.SessionID)
+	if runtimeSession.SecureManager == nil {
+		t.Fatal("session should have a SecureManager after full cold-start registration")
+	}
+
+	// Connect as admin client
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	conn, err := h.Connect(ctx)
+	if err != nil {
+		t.Fatalf("Connect failed: %v", err)
+	}
+	t.Cleanup(func() { _ = conn.Close() })
+
+	rpc := clientrpc.NewMaliceRPCClient(conn)
+	sessionCtx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs(
+		"session_id", implant.SessionID,
+		"callee", consts.CalleeCMD,
+	))
+
+	// Enable keepalive for faster command execution
+	enableRealKeepalive(t, &realRPCFixture{
+		h: h, implant: implant, rpc: rpc, session: sessionCtx,
+	}, true)
+
+	waitRealActiveConnection(t, implant.SessionID)
+
+	// Execute pwd command — this goes through age-encrypted channel
+	pwdTask, err := rpc.Pwd(sessionCtx, &implantpb.Request{Name: consts.ModulePwd})
+	if err != nil {
+		t.Fatalf("Pwd (post full cold-start key exchange) failed: %v", err)
+	}
+	pwdContent := waitRealTaskFinish(t, rpc, implant.SessionID, pwdTask.TaskId)
+	pwdOutput := strings.TrimSpace(pwdContent.GetSpite().GetResponse().GetOutput())
+	if pwdOutput == "" {
+		t.Fatal("pwd output should not be empty after full cold-start key exchange")
+	}
+	t.Logf("pwd after full cold-start key exchange: %s", pwdOutput)
+
+	// Disable keepalive
+	enableRealKeepalive(t, &realRPCFixture{
+		h: h, implant: implant, rpc: rpc, session: sessionCtx,
+	}, false)
+
+	t.Log("full cold-start key exchange E2E test passed: no pre-shared keys → key exchange → encrypted commands")
+}
+
 // TestRealImplantTLSPipelineE2E verifies the implant can connect to a TLS-enabled
 // pipeline (self-signed cert with CA verification) and execute commands normally.
 // This tests the full certificate chain: pipeline generates CA+cert → profile

server/rpc/rpc-implant.go

@@ -309,12 +309,16 @@ func (rpc *Server) triggerKeyExchange(ctx context.Context, sess *core.Session) e
 		Nonce:     nonce,
 		Signature: signature,
 	}
-	sess.SecureManager.ResetCounters()
+	// Reset counter only on successful exchange — if the callback fires, the
+	// implant accepted the request and sent back its new public key. Resetting
+	// before the response arrives would swallow a failed attempt and delay the
+	// next retry by a full rotation budget (100 checkins).
 	_, err = rpc.GenericInternal(ctx, req, consts.ModuleKeyExchange, func(spite *implantpb.Spite) {
 		resp := spite.GetKeyExchangeResponse()
 		if resp == nil {
 			return
 		}
+		sess.SecureManager.ResetCounters()
 		sess.UpdateKeyPairFieldsAndPushCtrl(resp.PublicKey, keyPair.Private)
 	})
 	return err

server/testsupport/real_implant.go

@@ -163,6 +163,20 @@ func NewRealSecureTCPPipeline(t testing.TB, listenerName, pipelineName string) *
 	return pipeline
 }
 
+// NewRealFullColdStartTCPPipeline creates a TCP pipeline with age key exchange
+// enabled but NO pre-shared keys at all (neither server nor implant keypair).
+// Both sides start from scratch; the first key exchange establishes encryption.
+func NewRealFullColdStartTCPPipeline(t testing.TB, listenerName, pipelineName string) *clientpb.Pipeline {
+	t.Helper()
+
+	pipeline := NewRealTCPPipeline(t, listenerName, pipelineName)
+	pipeline.Secure = &clientpb.Secure{
+		Enable: true,
+		// Both keypairs intentionally left empty for full cold-start scenario
+	}
+	return pipeline
+}
+
 // NewRealTLSTCPPipeline creates a TCP pipeline with TLS enabled (self-signed cert).
 // The pipeline generates its own CA + server certificate, so the implant can
 // verify the server using the embedded CA cert.