fix(rpc): fallback to DB when session is offline for task and context queries

Session sweep removes dead sessions from memory, causing all task/context
RPC handlers to fail with "Session ID not found". Add DB+disk fallback
for read-only task handlers, context save operations, and session recovery
in SpiteStream so that offline sessions no longer block data retrieval.

Author: h3zh1

server/rpc/rpc-context.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/chainreactors/IoM-go/client"
 	"github.com/chainreactors/IoM-go/consts"
 	"github.com/chainreactors/IoM-go/proto/client/clientpb"
 	errs "github.com/chainreactors/IoM-go/types"
@@ -69,7 +70,8 @@ func getTaskFromContext(req *clientpb.Context) (*core.Task, error) {
 
 	sess, err := core.Sessions.Get(sessionID)
 	if err != nil {
-		return nil, err
+		// Session not in memory (dead/offline), fall back to DB.
+		return getTaskFromDB(sessionID, req.Task.TaskId)
 	}
 
 	task := sess.Tasks.Get(req.Task.TaskId)
@@ -80,6 +82,37 @@ func getTaskFromContext(req *clientpb.Context) (*core.Task, error) {
 	return task, nil
 }
 
+// getTaskFromDB constructs a minimal core.Task from DB when session is not in memory.
+func getTaskFromDB(sessionID string, taskID uint32) (*core.Task, error) {
+	dbTask, err := db.GetTaskBySessionAndSeq(sessionID, taskID)
+	if err != nil {
+		return nil, errs.ErrNotFoundTask
+	}
+	task := core.FromTaskProtobuf(dbTask.ToProtobuf())
+	// Build a minimal Session so that context handlers (HandleScreenshot, etc.)
+	// can resolve the file path and bindResolvedTask can produce session metadata.
+	task.Session = &core.Session{
+		ID:             sessionID,
+		SessionContext: &client.SessionContext{},
+	}
+	// Try to enrich from DB; non-fatal if it fails.
+	if dbSess, err := db.FindSession(sessionID); err == nil && dbSess != nil {
+		task.Session.Target = dbSess.Target
+		task.Session.PipelineID = dbSess.PipelineID
+		task.Session.ListenerID = dbSess.ListenerID
+		task.Session.Name = dbSess.ProfileName
+		task.Session.Note = dbSess.Note
+		task.Session.Group = dbSess.GroupName
+		task.Session.Type = dbSess.Type
+		task.Session.RawID = dbSess.RawID
+		task.Session.CreatedAt = dbSess.CreatedAt
+		if dbSess.Data != nil {
+			task.Session.SessionContext = dbSess.Data
+		}
+	}
+	return task, nil
+}
+
 func bindResolvedTask(req *clientpb.Context, task *core.Task) *clientpb.Context {
 	if req == nil || task == nil || task.Session == nil {
 		return req

server/rpc/rpc-listener.go

@@ -10,6 +10,7 @@ import (
 	"github.com/chainreactors/IoM-go/types"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/malice-network/server/internal/core"
+	"github.com/chainreactors/malice-network/server/internal/db"
 	"google.golang.org/protobuf/proto"
 	"io"
 	"time"
@@ -67,8 +68,20 @@ func (rpc *Server) SpiteStream(stream listenerrpc.ListenerRPC_SpiteStreamServer)
 
 		sess, err := core.Sessions.Get(msg.SessionId)
 		if err != nil {
-			logs.Log.Warnf("session %s not found", msg.SessionId)
-			continue
+			// Session not in memory — try to recover from DB since the implant
+			// is actively sending data, meaning it's still alive.
+			dbSess, dbErr := db.FindSession(msg.SessionId)
+			if dbErr != nil || dbSess == nil {
+				logs.Log.Warnf("session %s not found in memory or DB", msg.SessionId)
+				continue
+			}
+			sess, err = core.RecoverSession(dbSess)
+			if err != nil {
+				logs.Log.Warnf("session %s recovery failed: %v", msg.SessionId, err)
+				continue
+			}
+			core.Sessions.Add(sess)
+			logs.Log.Importantf("session %s recovered from DB via SpiteStream", msg.SessionId)
 		}
 		sess.SetLastCheckin(time.Now().Unix())
 		if sess.MarkAlive() {

server/rpc/rpc-task.go

@@ -139,6 +139,66 @@ func getAllTaskContextsFromDisk(sess *core.Session, task *core.Task) (*clientpb.
 	}, nil
 }
 
+// getTaskContextFromDB fetches task context from DB + disk when session is not in memory.
+func getTaskContextFromDB(sessionID string, taskID uint32, index int32) (*clientpb.TaskContext, error) {
+	dbSess, err := db.FindSession(sessionID)
+	if err != nil || dbSess == nil {
+		return nil, types.ErrNotFoundSession
+	}
+	dbTask, err := db.GetTaskBySessionAndSeq(sessionID, taskID)
+	if err != nil {
+		return nil, types.ErrNotFoundTask
+	}
+	entries, err := readTaskSpitesFromDisk(sessionID, dbTask.Seq)
+	if err != nil || len(entries) == 0 {
+		return nil, types.ErrNotFoundTaskContent
+	}
+	var spite *implantpb.Spite
+	if index == -1 {
+		spite = entries[len(entries)-1].Spite
+	} else {
+		for _, e := range entries {
+			if e.Index == int(index) {
+				spite = e.Spite
+				break
+			}
+		}
+	}
+	if spite == nil {
+		return nil, types.ErrNotFoundTaskContent
+	}
+	return &clientpb.TaskContext{
+		Task:    dbTask.ToProtobuf(),
+		Session: dbSess.ToProtobuf(),
+		Spite:   spite,
+	}, nil
+}
+
+// getAllTaskContextFromDB fetches all task contexts from DB + disk when session is not in memory.
+func getAllTaskContextFromDB(sessionID string, taskID uint32) (*clientpb.TaskContexts, error) {
+	dbSess, err := db.FindSession(sessionID)
+	if err != nil || dbSess == nil {
+		return nil, types.ErrNotFoundSession
+	}
+	dbTask, err := db.GetTaskBySessionAndSeq(sessionID, taskID)
+	if err != nil {
+		return nil, types.ErrNotFoundTask
+	}
+	entries, err := readTaskSpitesFromDisk(sessionID, dbTask.Seq)
+	if err != nil || len(entries) == 0 {
+		return nil, types.ErrNotFoundTaskContent
+	}
+	spites := make([]*implantpb.Spite, 0, len(entries))
+	for _, e := range entries {
+		spites = append(spites, e.Spite)
+	}
+	return &clientpb.TaskContexts{
+		Task:    dbTask.ToProtobuf(),
+		Session: dbSess.ToProtobuf(),
+		Spites:  spites,
+	}, nil
+}
+
 func getTaskContext(sess *core.Session, task *core.Task, index int32) (*clientpb.TaskContext, error) {
 	var msg *implantpb.Spite
 	var ok bool
@@ -174,7 +234,12 @@ func (rpc *Server) GetTasks(ctx context.Context, req *clientpb.TaskRequest) (*cl
 	} else {
 		sess, err := core.Sessions.Get(req.SessionId)
 		if err != nil {
-			return nil, types.ErrNotFoundSession
+			// Fallback to DB when session is not in memory (e.g., dead session)
+			modelTasks, dbErr := db.ListTasksBySession(req.SessionId)
+			if dbErr != nil {
+				return nil, types.ErrNotFoundSession
+			}
+			return modelTasks.ToProtobuf(), nil
 		}
 		return sess.Tasks.ToProtobuf(), nil
 	}
@@ -198,7 +263,7 @@ func (rpc *Server) GetTaskContent(ctx context.Context, req *clientpb.Task) (*cli
 	}
 	sess, err := core.Sessions.Get(req.SessionId)
 	if err != nil {
-		return nil, types.ErrNotFoundSession
+		return getTaskContextFromDB(req.SessionId, req.TaskId, req.Need)
 	}
 	task := sess.Tasks.GetOrRecover(sess, req.TaskId)
 	if task == nil {
@@ -217,7 +282,8 @@ func (rpc *Server) WaitTaskContent(ctx context.Context, req *clientpb.Task) (*cl
 	}
 	sess, err := core.Sessions.Get(req.SessionId)
 	if err != nil {
-		return nil, types.ErrNotFoundSession
+		// Session not in memory (dead), try DB+disk directly
+		return getTaskContextFromDB(req.SessionId, req.TaskId, req.Need)
 	}
 	task := sess.Tasks.GetOrRecover(sess, req.TaskId)
 	if task == nil {
@@ -269,7 +335,8 @@ func (rpc *Server) WaitTaskFinish(ctx context.Context, req *clientpb.Task) (*cli
 	}
 	sess, err := core.Sessions.Get(req.SessionId)
 	if err != nil {
-		return nil, types.ErrNotFoundSession
+		// Session not in memory (dead), try DB+disk directly
+		return getTaskContextFromDB(req.SessionId, req.TaskId, -1)
 	}
 	task := sess.Tasks.GetOrRecover(sess, req.TaskId)
 	if task == nil {
@@ -309,7 +376,7 @@ func (rpc *Server) GetAllTaskContent(ctx context.Context, req *clientpb.Task) (*
 	}
 	sess, err := core.Sessions.Get(req.SessionId)
 	if err != nil {
-		return nil, types.ErrNotFoundSession
+		return getAllTaskContextFromDB(req.SessionId, req.TaskId)
 	}
 	task := sess.Tasks.GetOrRecover(sess, req.TaskId)
 	if task == nil {