diff --git a/pkg/handlers/apk.go b/pkg/handlers/apk.go index 858161e53a95..e23caa738786 100644 --- a/pkg/handlers/apk.go +++ b/pkg/handlers/apk.go @@ -17,6 +17,7 @@ import ( "github.com/avast/apkparser" dextk "github.com/csnewman/dextk" + "github.com/trufflesecurity/trufflehog/v3/pkg/common" logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context" "github.com/trufflesecurity/trufflehog/v3/pkg/engine/defaults" "github.com/trufflesecurity/trufflehog/v3/pkg/iobuf" @@ -205,6 +206,9 @@ func (h *apkHandler) processAPK(ctx logContext.Context, input fileReader, apkCha // Process all files for secrets for _, file := range zipReader.File { + if common.IsDone(ctx) { + return ctx.Err() + } if err := h.processFile(ctx, file, resTable, apkChan); err != nil { ctx.Logger().V(2).Info(fmt.Sprintf("failed to process file: %s", file.Name), "error", err) } diff --git a/pkg/handlers/apk_test.go b/pkg/handlers/apk_test.go index b8ca00530a0e..94240050b7e1 100644 --- a/pkg/handlers/apk_test.go +++ b/pkg/handlers/apk_test.go @@ -1,15 +1,21 @@ package handlers import ( + "archive/zip" + "bytes" + stdCtx "context" + "fmt" "io" "net/http" "regexp" "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/trufflesecurity/trufflehog/v3/pkg/context" + "github.com/trufflesecurity/trufflehog/v3/pkg/iobuf" ) func TestAPKHandler(t *testing.T) { @@ -108,3 +114,79 @@ func TestOpenValidZipInvalidAPK(t *testing.T) { err = handler.processAPK(ctx, newReader, archiveChan) assert.Contains(t, err.Error(), "resources.arsc file not found") } + +// buildMinimalAPK creates a zip archive containing a minimal valid resources.arsc +// (empty resource table with 0 packages) and the specified number of dummy XML files. +func buildMinimalAPK(t *testing.T, fileCount int) []byte { + t.Helper() + + var buf bytes.Buffer + zw := zip.NewWriter(&buf) + + // Minimal resources.arsc: RES_TABLE_TYPE header with 0 packages. + w, err := zw.Create("resources.arsc") + if err != nil { + t.Fatal(err) + } + resArsc := []byte{ + 0x02, 0x00, // id = RES_TABLE_TYPE + 0x0c, 0x00, // headerSize = 12 + 0x0c, 0x00, 0x00, 0x00, // totalSize = 12 + 0x00, 0x00, 0x00, 0x00, // packageCount = 0 + } + if _, err := w.Write(resArsc); err != nil { + t.Fatal(err) + } + + for i := 0; i < fileCount; i++ { + fw, err := zw.Create(fmt.Sprintf("res/layout/file_%04d.xml", i)) + if err != nil { + t.Fatal(err) + } + if _, err := fw.Write([]byte("data")); err != nil { + t.Fatal(err) + } + } + + if err := zw.Close(); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +func TestProcessAPK_ExitsOnContextCancellation(t *testing.T) { + apkData := buildMinimalAPK(t, 500) + + rdr := iobuf.NewBufferedReaderSeeker(bytes.NewReader(apkData)) + defer rdr.Close() + + handler := newAPKHandler() + apkChan := make(chan DataOrErr, defaultBufferSize) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := handler.processAPK(context.AddLogger(ctx), fileReader{BufferedReadSeeker: rdr}, apkChan) + if err != stdCtx.Canceled { + t.Fatalf("expected context.Canceled, got: %v", err) + } +} + +func TestProcessAPK_ExitsOnDeadlineExceeded(t *testing.T) { + apkData := buildMinimalAPK(t, 500) + + rdr := iobuf.NewBufferedReaderSeeker(bytes.NewReader(apkData)) + defer rdr.Close() + + handler := newAPKHandler() + apkChan := make(chan DataOrErr, defaultBufferSize) + + ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond) + defer cancel() + time.Sleep(time.Millisecond) + + err := handler.processAPK(context.AddLogger(ctx), fileReader{BufferedReadSeeker: rdr}, apkChan) + if err != stdCtx.DeadlineExceeded { + t.Fatalf("expected context.DeadlineExceeded, got: %v", err) + } +}