Skip to content

Commit 645019e

Browse files
committed
fix(download): add per-chunk timeout and retry to prevent deadlock
When a chunk response is lost in transit (e.g., REM layer drops the connection mid-transfer), recvSpite blocked indefinitely waiting for the chunk that would never arrive. The implant was also stuck waiting for the next chunk request, creating a mutual deadlock that only resolved after session_dead fired (~4 minutes later). Add a 30s per-chunk timeout using a select on time.After(MinTimeout). On timeout, resend the current chunk request to the implant. After maxChunkRetries (3) consecutive timeouts for the same chunk, return an error so the task fails fast rather than hanging silently. Retries counter resets to 0 on each successful chunk receipt, so transient single-packet loss does not consume retry budget.
1 parent 1cb86eb commit 645019e

1 file changed

Lines changed: 23 additions & 1 deletion

File tree

server/rpc/rpc-file.go

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,11 @@ import (
1616
"github.com/chainreactors/malice-network/server/internal/parser"
1717
"os"
1818
"path/filepath"
19+
"time"
1920
)
2021

22+
const maxChunkRetries = 3
23+
2124
var rpcFileSaveContext = db.SaveContext
2225

2326
func downloadChunkCount(size int, chunkSize int) int {
@@ -345,8 +348,27 @@ func (rpc *Server) Download(ctx context.Context, req *implantpb.DownloadRequest)
345348
return err
346349
}
347350

351+
retries := 0
348352
for {
349-
resp, ok := recvSpite(greq.Task.Ctx, out)
353+
var resp *implantpb.Spite
354+
var ok bool
355+
select {
356+
case resp, ok = <-out:
357+
retries = 0
358+
case <-greq.Task.Ctx.Done():
359+
return ErrTaskContextCancelled
360+
case <-time.After(consts.MinTimeout):
361+
retries++
362+
if retries >= maxChunkRetries {
363+
return fmt.Errorf("chunk %d timed out after %d retries", current_cur, maxChunkRetries)
364+
}
365+
logs.Log.Debugf("[download] chunk %d timeout, retrying (%d/%d)", current_cur, retries, maxChunkRetries)
366+
if err := in.Send(curRequest); err != nil {
367+
return err
368+
}
369+
continue
370+
}
371+
350372
if !ok {
351373
return ErrTaskContextCancelled
352374
}

0 commit comments

Comments
 (0)