From f148b5dded456df6513e6d81bf97d2ec36b0f585 Mon Sep 17 00:00:00 2001
From: Cody Soyland <codysoyland@github.com>
Date: Wed, 18 Dec 2024 15:00:13 -0500
Subject: [PATCH] Add multihasher for computing multiple hashes at once

Signed-off-by: Cody Soyland <codysoyland@github.com>
---
 pkg/verify/signature.go               | 34 +++++++++++++++++++++++
 pkg/verify/signature_internal_test.go | 40 +++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 pkg/verify/signature_internal_test.go

diff --git a/pkg/verify/signature.go b/pkg/verify/signature.go
index c7056c4..760b139 100644
--- a/pkg/verify/signature.go
+++ b/pkg/verify/signature.go
@@ -265,3 +265,37 @@ func verifyMessageSignatureWithArtifactDigest(verifier signature.Verifier, msg M
 
 	return nil
 }
+
+type multihasher struct {
+	hashfuncs []crypto.Hash
+	hashes    []hash.Hash
+}
+
+func newMultihasher(hashfuncs []crypto.Hash) *multihasher {
+	hashes := make([]hash.Hash, len(hashfuncs))
+	for i := range hashfuncs {
+		hashes[i] = hashfuncs[i].New()
+	}
+	return &multihasher{
+		hashfuncs: hashfuncs,
+		hashes:    hashes,
+	}
+}
+
+func (m *multihasher) Write(p []byte) (n int, err error) {
+	for i := range m.hashes {
+		n, err = m.hashes[i].Write(p)
+		if err != nil {
+			return
+		}
+	}
+	return
+}
+
+func (m *multihasher) Sum(b []byte) map[crypto.Hash][]byte {
+	sums := make(map[crypto.Hash][]byte, len(m.hashes))
+	for i := range m.hashes {
+		sums[m.hashfuncs[i]] = m.hashes[i].Sum(b)
+	}
+	return sums
+}
diff --git a/pkg/verify/signature_internal_test.go b/pkg/verify/signature_internal_test.go
new file mode 100644
index 0000000..8388483
--- /dev/null
+++ b/pkg/verify/signature_internal_test.go
@@ -0,0 +1,40 @@
+// Copyright 2024 The Sigstore Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package verify
+
+import (
+	"crypto"
+	"crypto/sha256"
+	"crypto/sha512"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMultiHasher(t *testing.T) {
+	testBytes := []byte("Hello, world!")
+	hash256 := sha256.Sum256(testBytes)
+	hash512 := sha512.Sum512(testBytes)
+
+	hasher := newMultihasher([]crypto.Hash{crypto.SHA256, crypto.SHA512})
+	_, err := hasher.Write(testBytes)
+	assert.NoError(t, err)
+
+	hashes := hasher.Sum(nil)
+
+	assert.Equal(t, 2, len(hashes))
+	assert.EqualValues(t, hash256[:], hashes[crypto.SHA256])
+	assert.EqualValues(t, hash512[:], hashes[crypto.SHA512])
+}