refactor(ec2): use x/crypto/ssh for key parsing and fingerprints (#567)

2026-10-15 17:11:30 +00:00 · 2026-07-22 15:14:08 +10:00
parent 5fd74a8212
commit 2119981924
2 changed files with 203 additions and 153 deletions
@@ -3,10 +3,6 @@ package handlers_ec2_key
 import (
 	"bytes"
 	"context"
-	"crypto/md5" //#nosec G501 - need md5 for AWS compatibility
-	"crypto/sha256"
-	"encoding/base64"
-	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -25,6 +21,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
+	"golang.org/x/crypto/ssh"
 )

 // Ensure KeyServiceImpl implements KeyService.
@@ -145,12 +142,14 @@ func (s *KeyServiceImpl) CreateKeyPair(ctx context.Context, input *ec2.CreateKey
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}

-	// Calculate fingerprint based on key type
-	fingerprint, err := s.calculateFingerprint(publicKeyData, keyType)
+	// Fingerprint the key we just generated; the digest algorithm follows the
+	// key algorithm, so it is derived from the parsed key rather than keyType.
+	publicKey, _, _, _, err := ssh.ParseAuthorizedKey(publicKeyData)
 	if err != nil {
-		slog.ErrorContext(ctx, "Failed to calculate fingerprint", "err", err)
+		slog.ErrorContext(ctx, "Failed to parse generated public key", "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
+	fingerprint := keyFingerprint(publicKey)

 	// Upload public key to S3
 	_, err = s.store.PutObject(ctx, &s3.PutObjectInput{
@@ -198,43 +197,38 @@ func (s *KeyServiceImpl) CreateKeyPair(ctx context.Context, input *ec2.CreateKey
 	return output, nil
 }

-// calculateFingerprint computes the SSH key fingerprint
-// - For RSA: SHA-1 hash of public key (MD5 for older format)
-// - For ED25519: SHA-256 hash of public key.
-func (s *KeyServiceImpl) calculateFingerprint(publicKeyData []byte, keyType string) (string, error) {
-	// Parse the public key to extract the key data
-	// Format: "ssh-ed25519 AAAAC3Nza... comment"
-	parts := strings.Fields(string(publicKeyData))
-	if len(parts) < 2 {
-		return "", fmt.Errorf("invalid public key format")
-	}
-
-	// Decode base64 key data
-	keyData, err := base64.StdEncoding.DecodeString(parts[1])
-	if err != nil {
-		return "", fmt.Errorf("failed to decode public key: %w", err)
-	}
-
-	if keyType == "ed25519" {
-		// ED25519 uses SHA-256 fingerprint
-		hash := sha256.Sum256(keyData)
-		return formatFingerprint(hash[:], "SHA256"), nil
-	} else {
-		// RSA uses SHA-1 or MD5 fingerprint
-		// AWS uses MD5 for RSA keys for backward compatibility
-		hash := md5.Sum(keyData) //#nosec G401 - need md5 for AWS compatibility
-		return formatFingerprint(hash[:], "MD5"), nil
+// keyFingerprint returns the OpenSSH fingerprint of a public key: the base64
+// SHA256 digest (SHA256:<base64>) for ED25519 keys, and the colon-separated MD5
+// digest of the SSH wire blob for RSA.
+//
+// ED25519 matches AWS. RSA does not, on either path. AWS hashes the DER
+// SubjectPublicKeyInfo encoding for an imported key, and the SHA-1 of the DER
+// PKCS#8 private key for one EC2 generated itself, so both differ from the value
+// returned here. AWS's own docs disagree on this -- the ImportKeyPair reference
+// cites RFC 4716 section 4, which is what this implements, while the fingerprint
+// verification guide gives the DER encoding -- and the verification guide is the
+// one that matches observed behaviour. Correcting it is deliberately out of
+// scope for the parse refactor because it changes a persisted, client-visible
+// value; do not "fix" the rendering here in isolation.
+func keyFingerprint(publicKey ssh.PublicKey) string {
+	if publicKey.Type() == ssh.KeyAlgoED25519 {
+		return ssh.FingerprintSHA256(publicKey)
 	}
+	return ssh.FingerprintLegacyMD5(publicKey)
 }

-// formatFingerprint formats the hash as a colon-separated hex string.
-func formatFingerprint(hash []byte, algorithm string) string {
-	if algorithm == "MD5" {
-		// MD5 format: aa:bb:cc:dd:...
-		return strings.ToLower(hex.EncodeToString(hash))
-	} else {
-		// SHA256 format: SHA256:base64encodedstring
-		return fmt.Sprintf("SHA256:%s", base64.RawStdEncoding.EncodeToString(hash))
+// keyPairType maps an SSH public key algorithm to the EC2 key type string. EC2
+// key pairs are RSA or ED25519 only, so every other algorithm -- ECDSA and
+// ssh-dss included -- is rejected rather than stored under a type the API has no
+// way to report back.
+func keyPairType(publicKey ssh.PublicKey) (string, error) {
+	switch publicKey.Type() {
+	case ssh.KeyAlgoED25519:
+		return "ed25519", nil
+	case ssh.KeyAlgoRSA:
+		return "rsa", nil
+	default:
+		return "", fmt.Errorf("unsupported key algorithm %q", publicKey.Type())
 	}
 }

@@ -710,46 +704,48 @@ func (s *KeyServiceImpl) ImportKeyPair(ctx context.Context, input *ec2.ImportKey
 		return nil, errors.New(awserrors.ErrorInvalidKeyPairDuplicate)
 	}

-	// Parse the public key material to extract key data and determine type
-	publicKeyData := input.PublicKeyMaterial
-	publicKeyString := string(publicKeyData)
-
-	// Parse the public key format: "ssh-rsa AAAAB..." or "ssh-ed25519 AAAAC..."
-	parts := strings.Fields(publicKeyString)
-	if len(parts) < 2 {
-		slog.ErrorContext(ctx, "Invalid public key format", "keyName", keyName)
+	// The material is stored verbatim and later served to instances as their
+	// authorized_keys, so it must hold exactly the one key the returned
+	// fingerprint describes. ParseAuthorizedKey skips leading comment and junk
+	// lines and stops at the first key, so a multi-line blob would otherwise be
+	// stored -- and trusted by the guest -- in full while only its first key was
+	// validated. Requiring a single line is how that rule is enforced; RFC 4716
+	// material is multi-line and would need normalising to an OpenSSH line
+	// before it reached here.
+	//
+	// This binds new imports only. Records written before the check exists are
+	// never revalidated, so material already in the object store may still be
+	// multi-line or option-prefixed and is still served to guests as-is.
+	publicKeyData := bytes.TrimSpace(input.PublicKeyMaterial)
+	if bytes.ContainsAny(publicKeyData, "\r\n") {
+		slog.ErrorContext(ctx, "Public key material is not a single key", "keyName", keyName)
 		return nil, errors.New(awserrors.ErrorInvalidKeyFormat)
 	}

-	// Determine key type from algorithm prefix
-	var keyType string
-	algorithmPrefix := parts[0]
-	switch {
-	case strings.HasPrefix(algorithmPrefix, "ssh-ed25519"):
-		keyType = "ed25519"
-	case strings.HasPrefix(algorithmPrefix, "ssh-rsa"):
-		keyType = "rsa"
-	case strings.HasPrefix(algorithmPrefix, "ecdsa-sha2-"):
-		// ECDSA keys are also supported but less common
-		keyType = "ecdsa"
-	default:
-		slog.ErrorContext(ctx, "Unsupported key type", "algorithm", algorithmPrefix, "keyName", keyName)
-		return nil, errors.New(awserrors.ErrorInvalidKeyFormat)
-	}
-
-	// Validate that the key data is valid base64
-	if _, err := base64.StdEncoding.DecodeString(parts[1]); err != nil {
-		slog.ErrorContext(ctx, "Invalid base64 in public key material", "keyName", keyName, "err", err)
-		return nil, errors.New(awserrors.ErrorInvalidKeyFormat)
-	}
-
-	// Calculate fingerprint from the imported public key
-	fingerprint, err := s.calculateFingerprint(publicKeyData, keyType)
+	// Parse the authorized-key line ("ssh-rsa AAAAB... comment"), which also
+	// validates the base64 body against the algorithm's wire encoding.
+	publicKey, _, options, _, err := ssh.ParseAuthorizedKey(publicKeyData)
 	if err != nil {
-		slog.ErrorContext(ctx, "Failed to calculate fingerprint", "err", err)
-		return nil, errors.New(awserrors.ErrorServerInternal)
+		slog.ErrorContext(ctx, "Invalid public key format", "keyName", keyName, "err", err)
+		return nil, errors.New(awserrors.ErrorInvalidKeyFormat)
 	}

+	// An option prefix ("command=...", "from=...") is not covered by the
+	// fingerprint, yet sshd would apply it to every login on every instance
+	// launched with this key pair. Refuse to import access the API cannot report.
+	if len(options) > 0 {
+		slog.ErrorContext(ctx, "Public key material carries authorized_keys options", "keyName", keyName, "options", options)
+		return nil, errors.New(awserrors.ErrorInvalidKeyFormat)
+	}
+
+	keyType, err := keyPairType(publicKey)
+	if err != nil {
+		slog.ErrorContext(ctx, "Unsupported key type", "algorithm", publicKey.Type(), "keyName", keyName, "err", err)
+		return nil, errors.New(awserrors.ErrorInvalidKeyFormat)
+	}
+
+	fingerprint := keyFingerprint(publicKey)
+
 	// Upload public key to S3
 	_, err = s.store.PutObject(ctx, &s3.PutObjectInput{
 		Bucket: aws.String(s.bucketName),
@@ -14,6 +14,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
 )

 const (
@@ -27,6 +28,20 @@ const testED25519PubKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOMqqnkVzrm0SdG6U
 // Valid RSA public key for import tests (generated locally, 2048-bit).
 const testRSAPubKey = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDP9LrByKWpgbX+prxBwnlf7lrmI5AfDwuiCofuvOAzt9/PwIDMMIAhlvlpm09jjOuuH/MRQApJB5A714Auv+hBKVK0lCq9KcTrnKZOpRj2aGgIZgaoO6P/POoZc+kBf9Y/GP18DCKc4y/XyBsp69dPP6XRdYBlEdeIeVZdgCPYrM7s5FjT7aML2ba2Y2EvH116hPxv+nJZGwM6yqWxWRyTOoNMMTAfNYmoKkNy2zC1FARUyqDwumJ2z5Fvo4ZdN1qoFPOsfPc3iB0NUtSZbLU1awScvHb0rwR5PRnelTZ3Nbkw8I8A0IAhBTE5veW9D38hDIJhRd4nW73BUhgmzDL7"

+// Well-formed ECDSA (nistp256) public key: parseable by x/crypto/ssh, but EC2
+// key pairs are RSA or ED25519 only, so import must reject it.
+const testECDSAPubKey = "ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBMZBEqs7mTlkOKGPSYp+tc5lZadVp9C2vCzIWZFTdnO1e3a8X59SdBWXiccQXjK1jxj+KLuQAEJY38kKqIUe/no="
+
+// Well-formed DSA public key: parseable by x/crypto/ssh but not a key type EC2 accepts.
+const testDSAPubKey = "ssh-dss AAAAB3NzaC1kc3MAAACBAMfooAGeGkAp6syFsxNueaASpwMr5t+BiLTxAmf9grlH/7MPhXHxxtrrzEdq4bK0mheLl4irAtardgR5ghOVxKXqz4dLlcEyv3H3tOY8Hq+JT/6w9j4FLjen4obXcilh7vfRPdL/A7Dk3Th9NSkrp43FmXUL4EyuMbqi7LcQYpMpAAAAFQDpwLJZvztWN3pEeT/MMLsVSmJedwAAAIBZjhEyHHk1iomvueP6GkmdrXt4V9+6BHjG/rHzQRlO79muU5ImX/BFALCc0RjaPNAoo0lF6ptaPf2HPeu3dtEAWM9iXH8SLqcAVX7B5FUYKFb7zsyQmlT3pKo21V3mCakKHDma8kbHSC2sysl1NOD4IkGTQalP4MuzIvCXNKbCdAAAAIEAl7OdP7hBngX0CuM0+cJXonZvnvIo1NOWGVu+dCn93mvoGjKFyZmLSEMIfFbmckQF2J4F9cM9aU6ht76k73DFnnA/F7WiJ+hIKLhL7Y8F0eDtWkawswvwxvHB+C7drrqezD6t5INX4CYlNQD4zqhgWKBSKVn3sQzQI95gFE51rKE="
+
+// Fingerprints as reported by `ssh-keygen -lf` (SHA256) and `ssh-keygen -l -E md5 -f`
+// (MD5, minus its "MD5:" prefix) for the keys above.
+const (
+	testED25519Fingerprint = "SHA256:+DiY3wvvV6TuJJhbpZisF/zLDA0zPMSvHdkr4UvCOqU"
+	testRSAFingerprint     = "35:76:17:80:c5:ca:87:56:ff:db:4b:0e:9c:76:63:6b"
+)
+
 func newTestKeyService() (*KeyServiceImpl, *objectstore.MemoryObjectStore) {
 	store := objectstore.NewMemoryObjectStore()
 	svc := NewKeyServiceImplWithStore(store, testBucket)
@@ -106,9 +121,8 @@ func TestCreateKeyPair_RSA(t *testing.T) {
 	require.NotNil(t, out)

 	assert.Equal(t, "my-rsa-key", *out.KeyName)
-	assert.NotEmpty(t, *out.KeyFingerprint)
-	// RSA fingerprint is MD5 hex (no "SHA256:" prefix, no colons in our format)
-	assert.False(t, strings.HasPrefix(*out.KeyFingerprint, "SHA256:"), "RSA fingerprint should not be SHA256 format")
+	// RSA fingerprint is the legacy colon-separated MD5 digest: 16 hex pairs.
+	assert.Regexp(t, `^([0-9a-f]{2}:){15}[0-9a-f]{2}$`, *out.KeyFingerprint)
 }

 func TestCreateKeyPair_NilInput(t *testing.T) {
@@ -186,8 +200,7 @@ func TestImportKeyPair_Success_ED25519(t *testing.T) {
 	assert.Equal(t, "imported-ed25519", *out.KeyName)
 	assert.NotEmpty(t, *out.KeyPairId)
 	assert.True(t, strings.HasPrefix(*out.KeyPairId, "key-"))
-	assert.NotEmpty(t, *out.KeyFingerprint)
-	assert.True(t, strings.HasPrefix(*out.KeyFingerprint, "SHA256:"))
+	assert.Equal(t, testED25519Fingerprint, *out.KeyFingerprint)

 	// Verify public key stored in S3
 	keyPath := "keys/" + testAccountID + "/imported-ed25519"
@@ -219,8 +232,25 @@ func TestImportKeyPair_Success_RSA(t *testing.T) {
 	require.NotNil(t, out)

 	assert.Equal(t, "imported-rsa", *out.KeyName)
-	assert.NotEmpty(t, *out.KeyFingerprint)
-	assert.False(t, strings.HasPrefix(*out.KeyFingerprint, "SHA256:"), "RSA fingerprint should be MD5 format")
+	assert.Equal(t, testRSAFingerprint, *out.KeyFingerprint)
+}
+
+// A trailing comment is accepted and preserved, and surrounding whitespace is
+// stripped, so what the guest is served is exactly the key the fingerprint
+// covers.
+func TestImportKeyPair_KeyWithComment(t *testing.T) {
+	svc, _ := newTestKeyService()
+
+	out, err := svc.ImportKeyPair(context.Background(), &ec2.ImportKeyPairInput{
+		KeyName:           aws.String("imported-commented"),
+		PublicKeyMaterial: []byte(testED25519PubKey + " user@laptop\n"),
+	}, testAccountID)
+	require.NoError(t, err)
+	assert.Equal(t, testED25519Fingerprint, *out.KeyFingerprint)
+
+	material, err := svc.GetPublicKeyMaterial(testAccountID, "imported-commented")
+	require.NoError(t, err)
+	assert.Equal(t, testED25519PubKey+" user@laptop", material)
 }

 func TestImportKeyPair_Duplicate(t *testing.T) {
@@ -267,10 +297,22 @@ func TestImportKeyPairInvalidKeyFormat(t *testing.T) {
 			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
 		},
 		{
-			name:           "UnsupportedAlgorithm",
+			name:           "TruncatedAlgorithmBody",
 			publicKey:      "ssh-dss AAAAB3NzaC1kc3MAAACB",
 			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
 		},
+		{
+			// Parses cleanly, but DSA is not a key type EC2 accepts.
+			name:           "UnsupportedAlgorithmDSA",
+			publicKey:      testDSAPubKey,
+			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
+		},
+		{
+			// Likewise ECDSA: EC2 key pairs are RSA or ED25519 only.
+			name:           "UnsupportedAlgorithmECDSA",
+			publicKey:      testECDSAPubKey,
+			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
+		},
 		{
 			name:           "InvalidBase64",
 			publicKey:      "ssh-rsa not-valid-base64!!!",
@@ -281,16 +323,47 @@ func TestImportKeyPairInvalidKeyFormat(t *testing.T) {
 			publicKey:      "ssh-ed25519 ",
 			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
 		},
+		{
+			// An option prefix parses, but sshd would honour it on every login
+			// while the fingerprint reports only the bare key.
+			name:           "AuthorizedKeysOptions",
+			publicKey:      `command="/bin/false",no-pty ` + testED25519PubKey,
+			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
+		},
+		{
+			// Only the first key would be validated and fingerprinted, yet both
+			// would be installed on the guest.
+			name:           "MultipleKeys",
+			publicKey:      testED25519PubKey + "\n" + testRSAPubKey,
+			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
+		},
+		{
+			// ParseAuthorizedKey skips leading junk, so the blob must be rejected
+			// before it reaches the parser.
+			name:           "LeadingCommentLine",
+			publicKey:      "# my key\n" + testED25519PubKey,
+			expectedErrMsg: awserrors.ErrorInvalidKeyFormat,
+		},
 	}

 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			// A name per case: a case that regressed into acceptance would
+			// otherwise store an object and fail every later case with
+			// Duplicate, pointing the failure at the wrong input.
+			keyName := "test-key-" + tt.name
 			_, err := svc.ImportKeyPair(context.Background(), &ec2.ImportKeyPairInput{
-				KeyName:           aws.String("test-key"),
+				KeyName:           aws.String(keyName),
 				PublicKeyMaterial: []byte(tt.publicKey),
 			}, testAccountID)
 			require.Error(t, err)
 			assert.Equal(t, tt.expectedErrMsg, err.Error())
+
+			// Rejection must precede the upload, or the guest is served material
+			// the API refused.
+			_, err = svc.GetPublicKeyMaterial(testAccountID, keyName)
+			require.Error(t, err)
+			assert.Equal(t, awserrors.ErrorInvalidKeyPairNotFound, err.Error())
 		})
 	}
 }
@@ -570,10 +643,6 @@ func TestGetPublicKeyMaterial_EmptyObject(t *testing.T) {
 	assert.NotEqual(t, awserrors.ErrorInvalidKeyPairNotFound, err.Error())
 }

-// ============================================================
-// formatFingerprint Tests
-// ============================================================
-
 // ============================================================
 // getKeyNameFromKeyPairId Tests
 // ============================================================
@@ -638,82 +707,67 @@ func TestGetKeyNameFromKeyPairId_Success(t *testing.T) {
 }

 // ============================================================
-// formatFingerprint Tests
+// keyFingerprint / keyPairType Tests
 // ============================================================

-func TestFormatFingerprint_MD5(t *testing.T) {
-	hash := []byte{0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99}
-	result := formatFingerprint(hash, "MD5")
-	assert.Equal(t, "aabbccddeeff00112233445566778899", result)
-	// Should be lowercase hex, no colons
-	assert.NotContains(t, result, ":")
-	assert.Equal(t, strings.ToLower(result), result)
-}
-
-func TestFormatFingerprint_SHA256(t *testing.T) {
-	hash := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08}
-	result := formatFingerprint(hash, "SHA256")
-	assert.True(t, strings.HasPrefix(result, "SHA256:"))
-	// Should be raw base64 (no padding)
-	assert.False(t, strings.HasSuffix(result, "="))
-}
-
-func TestFormatFingerprint_EmptyHash(t *testing.T) {
-	result := formatFingerprint([]byte{}, "MD5")
-	assert.Empty(t, result)
-
-	result = formatFingerprint([]byte{}, "SHA256")
-	assert.Equal(t, "SHA256:", result)
-}
-
-// ============================================================
-// calculateFingerprint Tests
-// ============================================================
-
-func TestCalculateFingerprint_ED25519(t *testing.T) {
-	svc, _ := newTestKeyService()
-
-	fp, err := svc.calculateFingerprint([]byte(testED25519PubKey), "ed25519")
+// parseTestPubKey parses an authorized-key line that the test asserts is valid.
+func parseTestPubKey(t *testing.T, material string) ssh.PublicKey {
+	t.Helper()
+	pubKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(material))
 	require.NoError(t, err)
-	assert.True(t, strings.HasPrefix(fp, "SHA256:"), "ed25519 should use SHA256 fingerprint")
-	assert.NotEqual(t, "SHA256:", fp, "fingerprint should have content after prefix")
+	return pubKey
 }

-func TestCalculateFingerprint_RSA(t *testing.T) {
-	svc, _ := newTestKeyService()
+func TestKeyFingerprint(t *testing.T) {
+	tests := []struct {
+		name     string
+		pubKey   string
+		expected string
+	}{
+		// ED25519 keys report the OpenSSH SHA256 digest.
+		{name: "ED25519", pubKey: testED25519PubKey, expected: testED25519Fingerprint},
+		// RSA keys report the legacy colon-separated MD5 digest, as AWS does.
+		{name: "RSA", pubKey: testRSAPubKey, expected: testRSAFingerprint},
+	}

-	fp, err := svc.calculateFingerprint([]byte(testRSAPubKey), "rsa")
-	require.NoError(t, err)
-	assert.False(t, strings.HasPrefix(fp, "SHA256:"), "RSA should use MD5 fingerprint")
-	assert.NotEmpty(t, fp)
-	// MD5 of RSA key data = 32 hex chars
-	assert.Len(t, fp, 32)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.expected, keyFingerprint(parseTestPubKey(t, tt.pubKey)))
+		})
+	}
 }

-func TestCalculateFingerprint_Deterministic(t *testing.T) {
-	svc, _ := newTestKeyService()
-
-	fp1, err := svc.calculateFingerprint([]byte(testED25519PubKey), "ed25519")
-	require.NoError(t, err)
-	fp2, err := svc.calculateFingerprint([]byte(testED25519PubKey), "ed25519")
-	require.NoError(t, err)
-	assert.Equal(t, fp1, fp2, "same key should produce same fingerprint")
+// A comment or surrounding whitespace is not part of the key, so it must not
+// move the fingerprint.
+func TestKeyFingerprint_IgnoresComment(t *testing.T) {
+	withComment := parseTestPubKey(t, testED25519PubKey+" user@host\n")
+	assert.Equal(t, testED25519Fingerprint, keyFingerprint(withComment))
 }

-func TestCalculateFingerprint_InvalidFormat(t *testing.T) {
-	svc, _ := newTestKeyService()
+func TestKeyPairType(t *testing.T) {
+	tests := []struct {
+		name      string
+		pubKey    string
+		expected  string
+		expectErr bool
+	}{
+		{name: "ED25519", pubKey: testED25519PubKey, expected: "ed25519"},
+		{name: "RSA", pubKey: testRSAPubKey, expected: "rsa"},
+		{name: "ECDSAUnsupported", pubKey: testECDSAPubKey, expectErr: true},
+		{name: "DSAUnsupported", pubKey: testDSAPubKey, expectErr: true},
+	}

-	_, err := svc.calculateFingerprint([]byte("no-space-here"), "ed25519")
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "invalid public key format")
-}
-
-func TestCalculateFingerprint_InvalidBase64(t *testing.T) {
-	svc, _ := newTestKeyService()
-
-	_, err := svc.calculateFingerprint([]byte("ssh-ed25519 !!!not-base64!!!"), "ed25519")
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "failed to decode public key")
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			keyType, err := keyPairType(parseTestPubKey(t, tt.pubKey))
+			if tt.expectErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, tt.expected, keyType)
+		})
+	}
 }

 // --- DescribeKeyPairs AWS filter tests ---