cmd/cosign/cli: stop loading the entire blob into memory (#2500)

sign-blob always loads the entire blob into memory. This can lead to
problems when signing large blob files and increased memory usage.

Signed-off-by: Anish Shah <[email protected]>

Signed-off-by: Anish Shah <[email protected]>

cmd/cosign/cli/policy_init.go

@@ -18,6 +18,7 @@ package cli
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -286,7 +287,11 @@ func signPolicy() *cobra.Command {
 				if err != nil {
 					return err
 				}
-				entry, err := cosign.TLogUpload(ctx, rekorClient, sig, signed.Signed, rekorBytes)
+				checkSum := sha256.New()
+				if _, err := checkSum.Write(signed.Signed); err != nil {
+					return err
+				}
+				entry, err := cosign.TLogUpload(ctx, rekorClient, sig, checkSum, rekorBytes)
 				if err != nil {
 					return err
 				}

cmd/cosign/cli/sign/sign_blob.go

@@ -16,12 +16,11 @@
 package sign
 
 import (
-	"bytes"
 	"context"
+	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"io"
 	"os"
 	"path/filepath"
 
@@ -31,21 +30,26 @@ import (
 
 	"github.com/sigstore/cosign/cmd/cosign/cli/options"
 	"github.com/sigstore/cosign/cmd/cosign/cli/rekor"
+	internal "github.com/sigstore/cosign/internal/pkg/cosign"
 	"github.com/sigstore/cosign/pkg/cosign"
 	signatureoptions "github.com/sigstore/sigstore/pkg/signature/options"
 )
 
 // nolint
 func SignBlobCmd(ro *options.RootOptions, ko options.KeyOpts, regOpts options.RegistryOptions, payloadPath string, b64 bool, outputSignature string, outputCertificate string, tlogUpload bool) ([]byte, error) {
-	var payload []byte
+	var payload internal.HashReader
 	var err error
 	var rekorBytes []byte
 
 	if payloadPath == "-" {
-		payload, err = io.ReadAll(os.Stdin)
+		payload = internal.NewHashReader(os.Stdin, sha256.New())
 	} else {
 		fmt.Fprintln(os.Stderr, "Using payload from:", payloadPath)
-		payload, err = os.ReadFile(filepath.Clean(payloadPath))
+		f, err := os.Open(filepath.Clean(payloadPath))
+		if err != nil {
+			return nil, err
+		}
+		payload = internal.NewHashReader(f, sha256.New())
 	}
 	if err != nil {
 		return nil, err
@@ -60,7 +64,7 @@ func SignBlobCmd(ro *options.RootOptions, ko options.KeyOpts, regOpts options.Re
 	}
 	defer sv.Close()
 
-	sig, err := sv.SignMessage(bytes.NewReader(payload), signatureoptions.WithContext(ctx))
+	sig, err := sv.SignMessage(&payload, signatureoptions.WithContext(ctx))
 	if err != nil {
 		return nil, fmt.Errorf("signing blob: %w", err)
 	}
@@ -90,7 +94,7 @@ func SignBlobCmd(ro *options.RootOptions, ko options.KeyOpts, regOpts options.Re
 		if err != nil {
 			return nil, err
 		}
-		entry, err := cosign.TLogUpload(ctx, rekorClient, sig, payload, rekorBytes)
+		entry, err := cosign.TLogUpload(ctx, rekorClient, sig, &payload, rekorBytes)
 		if err != nil {
 			return nil, err
 		}

internal/pkg/cosign/common.go

@@ -14,7 +14,12 @@
 
 package cosign
 
-import "os"
+import (
+	"errors"
+	"hash"
+	"io"
+	"os"
+)
 
 func FileExists(filename string) (bool, error) {
 	info, err := os.Stat(filename)
@@ -26,3 +31,34 @@ func FileExists(filename string) (bool, error) {
 	}
 	return !info.IsDir(), nil
 }
+
+// HashReader hashes while it reads.
+type HashReader struct {
+	r io.Reader
+	h hash.Hash
+}
+
+func NewHashReader(r io.Reader, h hash.Hash) HashReader {
+	return HashReader{
+		r: io.TeeReader(r, h),
+		h: h,
+	}
+}
+
+// Read implements io.Reader.
+func (h *HashReader) Read(p []byte) (n int, err error) { return h.r.Read(p) }
+
+// Sum implements hash.Hash.
+func (h *HashReader) Sum(p []byte) []byte { return h.h.Sum(p) }
+
+// Reset implements hash.Hash.
+func (h *HashReader) Reset() { h.h.Reset() }
+
+// Size implements hash.Hash.
+func (h *HashReader) Size() int { return h.h.Size() }
+
+// BlockSize implements hash.Hash.
+func (h *HashReader) BlockSize() int { return h.h.BlockSize() }
+
+// Write implements hash.Hash
+func (h *HashReader) Write(p []byte) (int, error) { return 0, errors.New("not implemented") }

internal/pkg/cosign/common_test.go

@@ -16,6 +16,9 @@
 package cosign
 
 import (
+	"bytes"
+	"crypto/sha256"
+	"io"
 	"os"
 	"testing"
 )
@@ -49,3 +52,22 @@ func Test_FileExists(t *testing.T) {
 		})
 	}
 }
+
+func Test_HashReader(t *testing.T) {
+	input := []byte("hello world")
+	r := NewHashReader(bytes.NewReader(input), sha256.New())
+
+	got, err := io.ReadAll(&r)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(got, input) {
+		t.Errorf("io.ReadAll returned %s, want %s", got, input)
+	}
+
+	gotHash := r.Sum(nil)
+	if hash := sha256.Sum256(input); !bytes.Equal(gotHash, hash[:]) {
+		t.Errorf("Sum returned %s, want %s", gotHash, hash)
+	}
+}

internal/pkg/cosign/rekor/signer.go

@@ -17,6 +17,7 @@ package rekor
 import (
 	"context"
 	"crypto"
+	"crypto/sha256"
 	"encoding/base64"
 	"fmt"
 	"io"
@@ -90,7 +91,11 @@ func (rs *signerWrapper) Sign(ctx context.Context, payload io.Reader) (oci.Signa
 	}
 
 	bundle, err := uploadToTlog(rekorBytes, rs.rClient, func(r *client.Rekor, b []byte) (*models.LogEntryAnon, error) {
-		return cosignv1.TLogUpload(ctx, r, sigBytes, payloadBytes, b)
+		checkSum := sha256.New()
+		if _, err := checkSum.Write(payloadBytes); err != nil {
+			return nil, err
+		}
+		return cosignv1.TLogUpload(ctx, r, sigBytes, checkSum, b)
 	})
 	if err != nil {
 		return nil, nil, err

pkg/cosign/tlog.go

@@ -25,6 +25,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"hash"
 	"os"
 	"strconv"
 	"strings"
@@ -138,8 +139,8 @@ func GetRekorPubs(ctx context.Context, _ *client.Rekor) (map[string]RekorPubKey,
 }
 
 // TLogUpload will upload the signature, public key and payload to the transparency log.
-func TLogUpload(ctx context.Context, rekorClient *client.Rekor, signature, payload []byte, pemBytes []byte) (*models.LogEntryAnon, error) {
-	re := rekorEntry(payload, signature, pemBytes)
+func TLogUpload(ctx context.Context, rekorClient *client.Rekor, signature []byte, sha256CheckSum hash.Hash, pemBytes []byte) (*models.LogEntryAnon, error) {
+	re := rekorEntry(sha256CheckSum, signature, pemBytes)
 	returnVal := models.Hashedrekord{
 		APIVersion: swag.String(re.APIVersion()),
 		Spec:       re.HashedRekordObj,
@@ -184,17 +185,16 @@ func doUpload(ctx context.Context, rekorClient *client.Rekor, pe models.Proposed
 	return nil, errors.New("bad response from server")
 }
 
-func rekorEntry(payload, signature, pubKey []byte) hashedrekord_v001.V001Entry {
+func rekorEntry(sha256CheckSum hash.Hash, signature, pubKey []byte) hashedrekord_v001.V001Entry {
 	// TODO: Signatures created on a digest using a hash algorithm other than SHA256 will fail
 	// upload right now. Plumb information on the hash algorithm used when signing from the
 	// SignerVerifier to use for the HashedRekordObj.Data.Hash.Algorithm.
-	h := sha256.Sum256(payload)
 	return hashedrekord_v001.V001Entry{
 		HashedRekordObj: models.HashedrekordV001Schema{
 			Data: &models.HashedrekordV001SchemaData{
 				Hash: &models.HashedrekordV001SchemaDataHash{
 					Algorithm: swag.String(models.HashedrekordV001SchemaDataHashAlgorithmSha256),
-					Value:     swag.String(hex.EncodeToString(h[:])),
+					Value:     swag.String(hex.EncodeToString(sha256CheckSum.Sum(nil))),
 				},
 			},
 			Signature: &models.HashedrekordV001SchemaSignature{
@@ -341,7 +341,11 @@ func proposedEntry(b64Sig string, payload, pubKey []byte) ([]models.ProposedEntr
 		}
 		proposedEntry = []models.ProposedEntry{e}
 	} else {
-		re := rekorEntry(payload, signature, pubKey)
+		sha256CheckSum := sha256.New()
+		if _, err := sha256CheckSum.Write(payload); err != nil {
+			return nil, err
+		}
+		re := rekorEntry(sha256CheckSum, signature, pubKey)
 		entry := &models.Hashedrekord{
 			APIVersion: swag.String(re.APIVersion()),
 			Spec:       re.HashedRekordObj,