add type BlobCompression to reduce code duplication

Author: majewsky

internal/models/blob.go

@@ -4,9 +4,16 @@
 package models
 
 import (
+	"compress/gzip"
+	"errors"
+	"fmt"
+	"io"
 	"time"
 
+	"github.com/klauspost/compress/zstd"
 	"github.com/opencontainers/go-digest"
+	v1 "github.com/opencontainers/image-spec/specs-go/v1"
+	"go.podman.io/image/v5/manifest"
 	. "go.xyrillian.de/gg/option"
 )
 
@@ -42,6 +49,55 @@ func (b Blob) SafeMediaType() string {
 	return b.MediaType
 }
 
+// BlobCompression is an enum indicating the compression format used within the blob.
+// This information is derived from the blob's MediaType, and only reported when the MediaType is understood by Keppel.
+type BlobCompression string
+
+const (
+	// BlobCompressionUnknown indicates that we do not recognize the blob's MediaType, and thus do not know whether it is uncompressed.
+	BlobCompressionUnknown BlobCompression = "unknown"
+	// BlobCompressionNone indicates that we recognize the blob's MediaType and thus know that the payload is uncompressed.
+	BlobCompressionNone BlobCompression = "none"
+	BlobCompressionGzip BlobCompression = "gzip"
+	BlobCompressionZstd BlobCompression = "zstd"
+)
+
+// Compression reports the compression format used within the blob,
+// or None if the blob's MediaType is not understood by Keppel.
+func (b Blob) Compression() BlobCompression {
+	switch b.MediaType {
+	case manifest.DockerV2SchemaLayerMediaTypeUncompressed, v1.MediaTypeImageLayer:
+		return BlobCompressionNone
+	case manifest.DockerV2Schema2LayerMediaType, v1.MediaTypeImageLayerGzip:
+		return BlobCompressionGzip
+	case manifest.DockerV2SchemaLayerMediaTypeZstd, v1.MediaTypeImageLayerZstd:
+		return BlobCompressionZstd
+	default:
+		return BlobCompressionUnknown
+	}
+}
+
+// Reader wraps a reader for compressed data in this format with the respective
+// decompressor, returning a reader that yields uncompressed data.
+func (c BlobCompression) Reader(r io.Reader) (io.ReadCloser, error) {
+	switch c {
+	case BlobCompressionNone:
+		return io.NopCloser(r), nil
+	case BlobCompressionGzip:
+		return gzip.NewReader(r)
+	case BlobCompressionZstd:
+		zr, err := zstd.NewReader(r)
+		if err != nil {
+			return nil, err
+		}
+		return zr.IOReadCloser(), nil
+	case BlobCompressionUnknown:
+		return nil, errors.New("do not know how to handle read data with unknown compression format")
+	default:
+		panic(fmt.Sprintf("unexpected BlobCompression value: %q", c))
+	}
+}
+
 const (
 	// BlobValidationInterval is how often each blob will be validated by BlobValidationJob.
 	// This is here instead of near the job because package processor also needs to know it.

internal/tasks/manifests.go

@@ -4,7 +4,6 @@
 package tasks
 
 import (
-	"compress/gzip"
 	"context"
 	"database/sql"
 	"errors"
@@ -17,7 +16,6 @@ import (
 	"time"
 
 	"github.com/go-gorp/gorp/v3"
-	"github.com/klauspost/compress/zstd"
 	"github.com/opencontainers/go-digest"
 	imagespecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/prometheus/client_golang/prometheus"
@@ -815,15 +813,12 @@ func (j *Janitor) checkPreConditionsForTrivy(ctx context.Context, account models
 
 	// filter media types that trivy is known to support
 	for _, blob := range layerBlobs {
-		if blob.MediaType == imageManifest.DockerV2SchemaLayerMediaTypeUncompressed || blob.MediaType == imageManifest.DockerV2Schema2LayerMediaType || blob.MediaType == imageManifest.DockerV2SchemaLayerMediaTypeZstd ||
-			blob.MediaType == imagespecs.MediaTypeImageLayer || blob.MediaType == imagespecs.MediaTypeImageLayerGzip || blob.MediaType == imagespecs.MediaTypeImageLayerZstd {
-			continue
+		if blob.Compression() == models.BlobCompressionUnknown { // None = unknown compression method because we don't recognize `blob.MediaType`
+			securityInfo.VulnerabilityStatus = models.UnsupportedVulnerabilityStatus
+			securityInfo.Message = fmt.Sprintf("vulnerability scanning is not supported for blob layers with media type %q", blob.MediaType)
+			securityInfo.NextCheckAt = Some(j.timeNow().Add(j.addJitter(24 * time.Hour)))
+			return false, layerBlobs, nil
 		}
-
-		securityInfo.VulnerabilityStatus = models.UnsupportedVulnerabilityStatus
-		securityInfo.Message = fmt.Sprintf("vulnerability scanning is not supported for blob layers with media type %q", blob.MediaType)
-		securityInfo.NextCheckAt = Some(j.timeNow().Add(j.addJitter(24 * time.Hour)))
-		return false, layerBlobs, nil
 	}
 
 	// can only validate when all blobs are present in the storage
@@ -845,47 +840,32 @@ func (j *Janitor) checkPreConditionsForTrivy(ctx context.Context, account models
 		}
 
 		if blob.BlocksVulnScanning.IsNone() {
-			isUncompressed := blob.MediaType == imageManifest.DockerV2SchemaLayerMediaTypeUncompressed || blob.MediaType == imagespecs.MediaTypeImageLayer
-			isGzip := blob.MediaType == imageManifest.DockerV2Schema2LayerMediaType || blob.MediaType == imagespecs.MediaTypeImageLayerGzip
-			isZstd := blob.MediaType == imageManifest.DockerV2SchemaLayerMediaTypeZstd || blob.MediaType == imagespecs.MediaTypeImageLayerZstd
+			compression := blob.Compression()
 
 			// when measuring uncompressed size, use LimitReader as a simple but
 			// effective guard against zip bombs
 			limitBytes := int64(1 << 30 * blobUncompressedSizeTooBigGiB)
 			var numberBytes int64
 
-			if isUncompressed {
+			if compression == models.BlobCompressionNone {
 				numberBytes = int64(blob.SizeBytes) //nolint:gosec
-			} else if isGzip || isZstd {
+			} else {
 				// uncompress the blob to check if it's too large for Trivy to handle within its allotted timeout
 				reader, _, err := j.sd.ReadBlob(ctx, account, blob.StorageID)
 				if err != nil {
 					return false, layerBlobs, fmt.Errorf("cannot read blob %s: %w", blob.Digest, err)
 				}
 				defer reader.Close()
 
-				if isGzip {
-					compressedReader, err := gzip.NewReader(reader)
-					if err != nil {
-						return false, layerBlobs, fmt.Errorf("cannot create gzip reader for blob %s: %w", blob.Digest, err)
-					}
-					defer compressedReader.Close()
-
-					numberBytes, err = io.Copy(io.Discard, io.LimitReader(compressedReader, limitBytes+1))
-					if err != nil {
-						return false, layerBlobs, fmt.Errorf("cannot decompress gzip blob %s: %w", blob.Digest, err)
-					}
-				} else if isZstd {
-					compressedReader, err := zstd.NewReader(reader)
-					if err != nil {
-						return false, layerBlobs, fmt.Errorf("cannot create zstd reader for blob %s: %w", blob.Digest, err)
-					}
-					defer compressedReader.Close()
-
-					numberBytes, err = io.Copy(io.Discard, io.LimitReader(compressedReader, limitBytes+1))
-					if err != nil {
-						return false, layerBlobs, fmt.Errorf("cannot decompress zstd blob %s: %w", blob.Digest, err)
-					}
+				compressedReader, err := compression.Reader(reader)
+				if err != nil {
+					return false, layerBlobs, fmt.Errorf("cannot create %s reader for blob %s: %w", compression, blob.Digest, err)
+				}
+				defer compressedReader.Close()
+
+				numberBytes, err = io.Copy(io.Discard, io.LimitReader(compressedReader, limitBytes+1))
+				if err != nil {
+					return false, layerBlobs, fmt.Errorf("cannot decompress %s blob %s: %w", compression, blob.Digest, err)
 				}
 			}