feat: add archive checksum

Author: christophwitzko

internal/batch/download.go

@@ -56,28 +56,30 @@ func downloadFileAndVerifyChecksum(ctx context.Context, tarWriter *tar.Writer, f
 	return nil
 }
 
-func DownloadFilesAndTarGz(ctx context.Context, batchResponse *registry.BatchResponse) (string, error) {
-	tarFile, err := os.CreateTemp("", "plugin-archive-*.tar.gz")
+func DownloadFilesAndTarGz(ctx context.Context, batchResponse *registry.BatchResponse) (string, string, error) {
+	tgzFile, err := os.CreateTemp("", "plugin-archive-*.tar.gz")
 	if err != nil {
-		return "", fmt.Errorf("failed to create temp file: %w", err)
+		return "", "", fmt.Errorf("failed to create temp file: %w", err)
 	}
-	defer tarFile.Close()
-	gzipWriter := gzip.NewWriter(tarFile)
+	defer tgzFile.Close()
+
+	tgzHash := sha256.New()
+	gzipWriter := gzip.NewWriter(io.MultiWriter(tgzFile, tgzHash))
 	tarWriter := tar.NewWriter(gzipWriter)
 	for _, plugin := range batchResponse.Plugins {
 		fileName := fmt.Sprintf("%s_%s/%s/%s/%s", batchResponse.OS, batchResponse.Arch, plugin.FullName, plugin.Version, plugin.FileName)
 		err = downloadFileAndVerifyChecksum(ctx, tarWriter, fileName, plugin.URL, plugin.Checksum)
 		if err != nil {
-			return "", fmt.Errorf("failed to add file to tar archive: %w", err)
+			return "", "", fmt.Errorf("failed to add file to tar archive: %w", err)
 		}
 	}
 	err = tarWriter.Close()
 	if err != nil {
-		return "", fmt.Errorf("failed to close tar writer: %w", err)
+		return "", "", fmt.Errorf("failed to close tar writer: %w", err)
 	}
 	err = gzipWriter.Close()
 	if err != nil {
-		return "", fmt.Errorf("failed to close gzip writer: %w", err)
+		return "", "", fmt.Errorf("failed to close gzip writer: %w", err)
 	}
-	return tarFile.Name(), nil
+	return tgzFile.Name(), hex.EncodeToString(tgzHash.Sum(nil)), nil
 }

internal/batch/download_test.go

@@ -5,6 +5,8 @@ import (
 	"bytes"
 	"compress/gzip"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"io"
 	"net/http"
@@ -74,14 +76,18 @@ func TestDownloadFilesAndTarGz(t *testing.T) {
 		Plugins: plugins,
 	}
 
-	tgzFileName, err := DownloadFilesAndTarGz(context.Background(), batchResponse)
+	tgzFileName, tgzChecksum, err := DownloadFilesAndTarGz(context.Background(), batchResponse)
 	require.NoError(t, err)
 	require.NotEmpty(t, tgzFileName)
 	defer os.Remove(tgzFileName)
 
 	tgzFile, err := os.ReadFile(tgzFileName)
 	require.NoError(t, err)
 
+	tarHash := sha256.New()
+	tarHash.Write(tgzFile)
+	require.Equal(t, tgzChecksum, hex.EncodeToString(tarHash.Sum(nil)))
+
 	gzipReader, err := gzip.NewReader(bytes.NewReader(tgzFile))
 	require.NoError(t, err)
 	defer gzipReader.Close()

internal/server/handler_batch.go

@@ -107,13 +107,14 @@ func (s *Server) batchGetPlugins(w http.ResponseWriter, r *http.Request) {
 	// the download url is deterministic, so we can set it here
 	batchResponse.DownloadURL = s.config.GetPublicPluginCacheDownloadURL(archiveKey)
 
-	_, err = s.storage.HeadObject(r.Context(), &s3.HeadObjectInput{
+	headRes, err := s.storage.HeadObject(r.Context(), &s3.HeadObjectInput{
 		Bucket: s.config.GetBucket(),
 		Key:    &archiveKey,
 	})
 	if err == nil {
 		// the archive already exists, return the response
 		s.log.Infof("found cached archive %s", archiveKey)
+		batchResponse.DownloadChecksum = headRes.Metadata["checksum"]
 		s.setInCache(batchRequestCacheKey, batchResponse)
 		s.writeJSON(w, batchResponse)
 		return
@@ -126,13 +127,14 @@ func (s *Server) batchGetPlugins(w http.ResponseWriter, r *http.Request) {
 	}
 
 	s.log.Infof("plugin archive %s not found, creating...", archiveKey)
-	tarFileName, err := batch.DownloadFilesAndTarGz(r.Context(), batchResponse)
+	tgzFileName, tgzChecksum, err := batch.DownloadFilesAndTarGz(r.Context(), batchResponse)
 	if err != nil {
 		s.writeJSONError(w, r, http.StatusInternalServerError, err, "could not create plugin archive")
 		return
 	}
-	s.log.Infof("created plugin archive %s, uploading...", tarFileName)
-	tarFile, err := os.Open(tarFileName)
+	batchResponse.DownloadChecksum = tgzChecksum
+	s.log.Infof("created plugin archive %s, uploading...", tgzFileName)
+	tarFile, err := os.Open(tgzFileName)
 	if err != nil {
 		s.writeJSONError(w, r, http.StatusInternalServerError, err, "could not open plugin archive")
 		return
@@ -143,6 +145,9 @@ func (s *Server) batchGetPlugins(w http.ResponseWriter, r *http.Request) {
 		Key:         &archiveKey,
 		Body:        tarFile,
 		ContentType: aws.String("application/gzip"),
+		Metadata: map[string]string{
+			"checksum": tgzChecksum,
+		},
 	})
 	if closeErr := tarFile.Close(); closeErr != nil {
 		s.log.Errorf("could not close plugin archive file: %v", closeErr)
@@ -153,7 +158,7 @@ func (s *Server) batchGetPlugins(w http.ResponseWriter, r *http.Request) {
 	}
 
 	s.log.Infof("uploaded plugin archive.")
-	if rmErr := os.Remove(tarFileName); rmErr != nil {
+	if rmErr := os.Remove(tgzFileName); rmErr != nil {
 		s.log.Errorf("could not remove plugin archive file: %v", rmErr)
 	}
 

pkg/registry/registry.go

@@ -125,11 +125,12 @@ func (b BatchResponsePlugins) Has(fullName string) bool {
 }
 
 type BatchResponse struct {
-	OS           string
-	Arch         string
-	Plugins      BatchResponsePlugins
-	DownloadHash string
-	DownloadURL  string
+	OS               string
+	Arch             string
+	Plugins          BatchResponsePlugins
+	DownloadHash     string
+	DownloadURL      string
+	DownloadChecksum string
 }
 
 func NewBatchResponse(req *BatchRequest, plugins BatchResponsePlugins) *BatchResponse {