Merge pull request 'Corpus downloader' (#12) from quiet-afl into main

Reviewed-on: https://git.ts.mattnite.net/mattnite/cairn/pulls/12

actions/cairn-zig-fuzz-afl/action.yml

@@ -190,6 +190,7 @@ runs:
           mkdir -p "${FINDINGS}"
 
           AFL_EXIT=0
+          {
             AFL_NO_UI=1 \
             AFL_SKIP_CPUFREQ=1 \
             AFL_I_DONT_CARE_ABOUT_MISSING_CRASHES=1 \
@@ -199,8 +200,8 @@ runs:
               -i "${SEEDS}" \
               -o "${FINDINGS}" \
               ${EXTRA_AFL_ARGS} \
-            -- "${FUZZ_BIN}" \
-          || AFL_EXIT=$?
+              -- "${FUZZ_BIN}"
+          } >/dev/null 2>&1 || AFL_EXIT=$?
 
           if [ "${AFL_EXIT}" -eq 0 ]; then
             echo "AFL++ exited normally (completed run)"

cmd/cairn/main.go

@@ -1,7 +1,9 @@
 package main
 
 import (
+	"archive/tar"
 	"bytes"
+	"compress/gzip"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -544,45 +546,44 @@ func cmdCorpus(subcmd string, args []string) error {
 			return fmt.Errorf("creating output dir: %w", err)
 		}
 
-		// List all corpus entries.
-		resp, err := http.Get(serverURL + "/api/v1/targets/" + targetID + "/corpus?limit=10000")
+		resp, err := http.Get(serverURL + "/api/v1/targets/" + targetID + "/corpus/download")
 		if err != nil {
-			return fmt.Errorf("listing corpus: %w", err)
+			return fmt.Errorf("downloading corpus: %w", err)
 		}
 		defer resp.Body.Close()
-		body, _ := io.ReadAll(resp.Body)
+
+		if resp.StatusCode == http.StatusNoContent {
+			fmt.Printf("Downloaded 0 corpus entries to %s\n", dir)
+			return nil
+		}
 		if resp.StatusCode != http.StatusOK {
+			body, _ := io.ReadAll(resp.Body)
 			return fmt.Errorf("server returned %d: %s", resp.StatusCode, body)
 		}
 
-		var listResp struct {
-			Entries []struct {
-				ID      float64 `json:"id"`
-				BlobKey string  `json:"blob_key"`
-			} `json:"entries"`
-		}
-		if err := json.Unmarshal(body, &listResp); err != nil {
-			return fmt.Errorf("parsing response: %w", err)
-		}
-
-		var downloaded int
-		for _, entry := range listResp.Entries {
-			entryID := fmt.Sprintf("%d", int(entry.ID))
-			dlURL := serverURL + "/api/v1/targets/" + targetID + "/corpus/" + entryID + "/download"
-			dlResp, err := http.Get(dlURL)
+		gr, err := gzip.NewReader(resp.Body)
 		if err != nil {
-				return fmt.Errorf("downloading entry %s: %w", entryID, err)
+			return fmt.Errorf("decompressing corpus: %w", err)
+		}
+		defer gr.Close()
+
+		tr := tar.NewReader(gr)
+		var downloaded int
+		for {
+			hdr, err := tr.Next()
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				return fmt.Errorf("reading tar: %w", err)
 			}
 
-			filename := filepath.Base(entry.BlobKey)
-			outPath := filepath.Join(dir, filename)
+			outPath := filepath.Join(dir, filepath.Base(hdr.Name))
 			out, err := os.Create(outPath)
 			if err != nil {
-				dlResp.Body.Close()
 				return fmt.Errorf("creating file %s: %w", outPath, err)
 			}
-			_, err = io.Copy(out, dlResp.Body)
-			dlResp.Body.Close()
+			_, err = io.Copy(out, tr)
 			out.Close()
 			if err != nil {
 				return fmt.Errorf("writing file %s: %w", outPath, err)

internal/handler/targets.go

@@ -1,9 +1,12 @@
 package handler
 
 import (
+	"archive/tar"
+	"compress/gzip"
 	"fmt"
 	"io"
 	"net/http"
+	"path/filepath"
 	"strconv"
 
 	"github.com/gin-gonic/gin"
@@ -317,3 +320,57 @@ func (h *CorpusHandler) Download(c *gin.Context) {
 	c.Header("Content-Type", "application/octet-stream")
 	_, _ = io.Copy(c.Writer, reader)
 }
+
+func (h *CorpusHandler) DownloadAll(c *gin.Context) {
+	targetID, err := parseUintID(c.Param("id"), "target id")
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	ctx := c.Request.Context()
+
+	var entries []models.CorpusEntry
+	if err := h.DB.WithContext(ctx).Where("target_id = ?", targetID).Find(&entries).Error; err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
+		return
+	}
+
+	if len(entries) == 0 {
+		c.Status(http.StatusNoContent)
+		return
+	}
+
+	c.Header("Content-Type", "application/gzip")
+	c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=\"corpus-%d.tar.gz\"", targetID))
+
+	gw := gzip.NewWriter(c.Writer)
+	defer gw.Close()
+	tw := tar.NewWriter(gw)
+	defer tw.Close()
+
+	for _, entry := range entries {
+		reader, err := h.Store.Get(ctx, entry.BlobKey)
+		if err != nil {
+			continue
+		}
+
+		data, err := io.ReadAll(reader)
+		reader.Close()
+		if err != nil {
+			continue
+		}
+
+		hdr := &tar.Header{
+			Name: filepath.Base(entry.BlobKey),
+			Mode: 0o644,
+			Size: int64(len(data)),
+		}
+		if err := tw.WriteHeader(hdr); err != nil {
+			return
+		}
+		if _, err := tw.Write(data); err != nil {
+			return
+		}
+	}
+}

internal/web/routes.go

@@ -80,6 +80,7 @@ func NewRouter(cfg RouterConfig) (*gin.Engine, error) {
 	api.GET("/targets/:id", targetAPI.Detail)
 	api.POST("/targets/:id/corpus", corpusAPI.Upload)
 	api.GET("/targets/:id/corpus", corpusAPI.List)
+	api.GET("/targets/:id/corpus/download", corpusAPI.DownloadAll)
 	api.GET("/targets/:id/corpus/:entry_id/download", corpusAPI.Download)
 	api.POST("/runs", runAPI.Start)
 	api.GET("/runs", runAPI.List)