refactoring

Signed-off-by: Patrick Zheng <[email protected]>
notaryproject · Nov 25, 2024 · 45c0fa2 · 45c0fa2
1 parent 3ef5c27
commit 45c0fa2
Show file tree

Hide file tree

Showing 4 changed files with 195 additions and 188 deletions.
diff --git a/internal/algorithm/algorithm.go b/internal/algorithm/algorithm.go
@@ -11,16 +11,46 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// package algorithm includes signature algorithms accepted by Notary Project
 package algorithm
 
 import (
+	"crypto"
 	"crypto/ecdsa"
 	"crypto/rsa"
 	"crypto/x509"
 	"errors"
 	"fmt"
 )
 
+// Algorithm defines the signature algorithm.
+type Algorithm int
+
+// Signature algorithms supported by this library.
+//
+// Reference: https://github.com/notaryproject/notaryproject/blob/main/specs/signature-specification.md#algorithm-selection
+const (
+	AlgorithmPS256 Algorithm = 1 + iota // RSASSA-PSS with SHA-256
+	AlgorithmPS384                      // RSASSA-PSS with SHA-384
+	AlgorithmPS512                      // RSASSA-PSS with SHA-512
+	AlgorithmES256                      // ECDSA on secp256r1 with SHA-256
+	AlgorithmES384                      // ECDSA on secp384r1 with SHA-384
+	AlgorithmES512                      // ECDSA on secp521r1 with SHA-512
+)
+
+// Hash returns the hash function of the algorithm.
+func (alg Algorithm) Hash() crypto.Hash {
+	switch alg {
+	case AlgorithmPS256, AlgorithmES256:
+		return crypto.SHA256
+	case AlgorithmPS384, AlgorithmES384:
+		return crypto.SHA384
+	case AlgorithmPS512, AlgorithmES512:
+		return crypto.SHA512
+	}
+	return 0
+}
+
 // KeyType defines the key type.
 type KeyType int
 
@@ -38,6 +68,31 @@ type KeySpec struct {
 	Size int
 }
 
+// SignatureAlgorithm returns the signing algorithm associated with the KeySpec.
+func (k KeySpec) SignatureAlgorithm() Algorithm {
+	switch k.Type {
+	case KeyTypeEC:
+		switch k.Size {
+		case 256:
+			return AlgorithmES256
+		case 384:
+			return AlgorithmES384
+		case 521:
+			return AlgorithmES512
+		}
+	case KeyTypeRSA:
+		switch k.Size {
+		case 2048:
+			return AlgorithmPS256
+		case 3072:
+			return AlgorithmPS384
+		case 4096:
+			return AlgorithmPS512
+		}
+	}
+	return 0
+}
+
 // ExtractKeySpec extracts KeySpec from the signing certificate.
 func ExtractKeySpec(signingCert *x509.Certificate) (KeySpec, error) {
 	switch key := signingCert.PublicKey.(type) {

diff --git a/internal/algorithm/algorithm_test.go b/internal/algorithm/algorithm_test.go
@@ -14,6 +14,7 @@
 package algorithm
 
 import (
+	"crypto"
 	"crypto/ecdsa"
 	"crypto/ed25519"
 	"crypto/elliptic"
@@ -27,6 +28,59 @@ import (
 	"github.com/notaryproject/notation-core-go/testhelper"
 )
 
+func TestHash(t *testing.T) {
+	tests := []struct {
+		name   string
+		alg    Algorithm
+		expect crypto.Hash
+	}{
+		{
+			name:   "PS256",
+			alg:    AlgorithmPS256,
+			expect: crypto.SHA256,
+		},
+		{
+			name:   "ES256",
+			alg:    AlgorithmES256,
+			expect: crypto.SHA256,
+		},
+		{
+			name:   "PS384",
+			alg:    AlgorithmPS384,
+			expect: crypto.SHA384,
+		},
+		{
+			name:   "ES384",
+			alg:    AlgorithmES384,
+			expect: crypto.SHA384,
+		},
+		{
+			name:   "PS512",
+			alg:    AlgorithmPS512,
+			expect: crypto.SHA512,
+		},
+		{
+			name:   "ES512",
+			alg:    AlgorithmES512,
+			expect: crypto.SHA512,
+		},
+		{
+			name:   "UnsupportedAlgorithm",
+			alg:    0,
+			expect: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			hash := tt.alg.Hash()
+			if hash != tt.expect {
+				t.Fatalf("Expected %v, got %v", tt.expect, hash)
+			}
+		})
+	}
+}
+
 func TestExtractKeySpec(t *testing.T) {
 	type testCase struct {
 		name      string
@@ -114,3 +168,77 @@ func TestExtractKeySpec(t *testing.T) {
 		})
 	}
 }
+
+func TestSignatureAlgorithm(t *testing.T) {
+	tests := []struct {
+		name    string
+		keySpec KeySpec
+		expect  Algorithm
+	}{
+		{
+			name: "EC 256",
+			keySpec: KeySpec{
+				Type: KeyTypeEC,
+				Size: 256,
+			},
+			expect: AlgorithmES256,
+		},
+		{
+			name: "EC 384",
+			keySpec: KeySpec{
+				Type: KeyTypeEC,
+				Size: 384,
+			},
+			expect: AlgorithmES384,
+		},
+		{
+			name: "EC 521",
+			keySpec: KeySpec{
+				Type: KeyTypeEC,
+				Size: 521,
+			},
+			expect: AlgorithmES512,
+		},
+		{
+			name: "RSA 2048",
+			keySpec: KeySpec{
+				Type: KeyTypeRSA,
+				Size: 2048,
+			},
+			expect: AlgorithmPS256,
+		},
+		{
+			name: "RSA 3072",
+			keySpec: KeySpec{
+				Type: KeyTypeRSA,
+				Size: 3072,
+			},
+			expect: AlgorithmPS384,
+		},
+		{
+			name: "RSA 4096",
+			keySpec: KeySpec{
+				Type: KeyTypeRSA,
+				Size: 4096,
+			},
+			expect: AlgorithmPS512,
+		},
+		{
+			name: "Unsupported key spec",
+			keySpec: KeySpec{
+				Type: 0,
+				Size: 0,
+			},
+			expect: 0,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			alg := tt.keySpec.SignatureAlgorithm()
+			if alg != tt.expect {
+				t.Errorf("unexpected signature algorithm: %v, expect: %v", alg, tt.expect)
+			}
+		})
+	}
+}
diff --git a/signature/algorithm.go b/signature/algorithm.go
@@ -14,56 +14,36 @@
 package signature
 
 import (
-	"crypto"
 	"crypto/x509"
 
 	"github.com/notaryproject/notation-core-go/internal/algorithm"
 )
 
 // Algorithm defines the signature algorithm.
-type Algorithm int
+type Algorithm = algorithm.Algorithm
 
 // Signature algorithms supported by this library.
 //
 // Reference: https://github.com/notaryproject/notaryproject/blob/main/specs/signature-specification.md#algorithm-selection
 const (
-	AlgorithmPS256 Algorithm = 1 + iota // RSASSA-PSS with SHA-256
-	AlgorithmPS384                      // RSASSA-PSS with SHA-384
-	AlgorithmPS512                      // RSASSA-PSS with SHA-512
-	AlgorithmES256                      // ECDSA on secp256r1 with SHA-256
-	AlgorithmES384                      // ECDSA on secp384r1 with SHA-384
-	AlgorithmES512                      // ECDSA on secp521r1 with SHA-512
+	AlgorithmPS256 = algorithm.AlgorithmPS256 // RSASSA-PSS with SHA-256
+	AlgorithmPS384 = algorithm.AlgorithmPS384 // RSASSA-PSS with SHA-384
+	AlgorithmPS512 = algorithm.AlgorithmPS512 // RSASSA-PSS with SHA-512
+	AlgorithmES256 = algorithm.AlgorithmES256 // ECDSA on secp256r1 with SHA-256
+	AlgorithmES384 = algorithm.AlgorithmES384 // ECDSA on secp384r1 with SHA-384
+	AlgorithmES512 = algorithm.AlgorithmES512 // ECDSA on secp521r1 with SHA-512
 )
 
 // KeyType defines the key type.
-type KeyType int
+type KeyType = algorithm.KeyType
 
 const (
-	KeyTypeRSA KeyType = 1 + iota // KeyType RSA
-	KeyTypeEC                     // KeyType EC
+	KeyTypeRSA = algorithm.KeyTypeRSA // KeyType RSA
+	KeyTypeEC  = algorithm.KeyTypeEC  // KeyType EC
 )
 
 // KeySpec defines a key type and size.
-type KeySpec struct {
-	// KeyType is the type of the key.
-	Type KeyType
-
-	// KeySize is the size of the key in bits.
-	Size int
-}
-
-// Hash returns the hash function of the algorithm.
-func (alg Algorithm) Hash() crypto.Hash {
-	switch alg {
-	case AlgorithmPS256, AlgorithmES256:
-		return crypto.SHA256
-	case AlgorithmPS384, AlgorithmES384:
-		return crypto.SHA384
-	case AlgorithmPS512, AlgorithmES512:
-		return crypto.SHA512
-	}
-	return 0
-}
+type KeySpec = algorithm.KeySpec
 
 // ExtractKeySpec extracts KeySpec from the signing certificate.
 func ExtractKeySpec(signingCert *x509.Certificate) (KeySpec, error) {
@@ -73,33 +53,5 @@ func ExtractKeySpec(signingCert *x509.Certificate) (KeySpec, error) {
 			Msg: err.Error(),
 		}
 	}
-	return KeySpec{
-		Type: KeyType(ks.Type),
-		Size: ks.Size,
-	}, nil
-}
-
-// SignatureAlgorithm returns the signing algorithm associated with the KeySpec.
-func (k KeySpec) SignatureAlgorithm() Algorithm {
-	switch k.Type {
-	case KeyTypeEC:
-		switch k.Size {
-		case 256:
-			return AlgorithmES256
-		case 384:
-			return AlgorithmES384
-		case 521:
-			return AlgorithmES512
-		}
-	case KeyTypeRSA:
-		switch k.Size {
-		case 2048:
-			return AlgorithmPS256
-		case 3072:
-			return AlgorithmPS384
-		case 4096:
-			return AlgorithmPS512
-		}
-	}
-	return 0
+	return ks, nil
 }