review: update crypto tests

Author: archandatta

pkg/extensions/webbotauth.go

@@ -226,19 +226,27 @@ func buildWebBotAuthExtension(ctx context.Context, browserExtDir, hostURL, keyDa
 	pterm.Info.Println("Validating key...")
 	var pemData []byte
 	var err error
+	// JWK data is used for keyid signing in background.ts
+	var jwkData string
 
 	if util.IsPEMKey(keyData) {
 		// Key is already in PEM format, validate it
 		if err := util.ValidatePEMKey(keyData); err != nil {
 			return "", fmt.Errorf("failed to validate PEM key: %w", err)
 		}
+
+		jwkData, err = util.ConvertPEMToJWK(keyData)
+		if err != nil {
+			return "", fmt.Errorf("failed to convert PEM to JWK: %w", err)
+		}
 		pemData = []byte(keyData)
 	} else {
 		// Key is in JWK format, convert to PEM
 		pemData, err = util.ConvertJWKToPEM(keyData)
 		if err != nil {
 			return "", fmt.Errorf("failed to convert JWK to PEM: %w", err)
 		}
+		jwkData = keyData
 	}
 
 	privateKeyPath := filepath.Join(browserExtDir, "private_key.pem")
@@ -250,7 +258,7 @@ func buildWebBotAuthExtension(ctx context.Context, browserExtDir, hostURL, keyDa
 	// Inject the JWK into background.ts (replacing the hardcoded test key)
 	pterm.Info.Println("Injecting custom JWK into background.ts...")
 	backgroundTsPath := filepath.Join(browserExtDir, "src", "background.ts")
-	if err := injectJWKIntoBackgroundTs(backgroundTsPath, keyData); err != nil {
+	if err := injectJWKIntoBackgroundTs(backgroundTsPath, jwkData); err != nil {
 		return "", fmt.Errorf("failed to inject JWK: %w", err)
 	}
 	pterm.Success.Println("Custom JWK injected successfully")

pkg/util/crypto.go

@@ -86,3 +86,50 @@ func ConvertJWKToPEM(jwkJSON string) ([]byte, error) {
 
 	return pem.EncodeToMemory(pemBlock), nil
 }
+
+// ConvertPEMToJWK converts an Ed25519 PEM private key to JWK format
+func ConvertPEMToJWK(pemData string) (string, error) {
+	// Decode PEM block
+	block, _ := pem.Decode([]byte(pemData))
+	if block == nil {
+		return "", fmt.Errorf("failed to decode PEM block")
+	}
+
+	if block.Type != "PRIVATE KEY" {
+		return "", fmt.Errorf("invalid PEM type: expected PRIVATE KEY, got %s", block.Type)
+	}
+
+	// Parse PKCS#8 private key
+	privateKeyInterface, err := x509.ParsePKCS8PrivateKey(block.Bytes)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse PKCS#8 private key: %w", err)
+	}
+
+	// Ensure it's an Ed25519 key
+	privateKey, ok := privateKeyInterface.(ed25519.PrivateKey)
+	if !ok {
+		return "", fmt.Errorf("invalid key type: expected Ed25519 private key, got %T", privateKeyInterface)
+	}
+
+	// Extract seed (first 32 bytes of Ed25519 private key)
+	seed := privateKey.Seed()
+
+	// Extract public key (last 32 bytes of Ed25519 private key)
+	publicKey := privateKey.Public().(ed25519.PublicKey)
+
+	// Encode to base64url (without padding)
+	jwk := jwkKey{
+		Kty: "OKP",
+		Crv: "Ed25519",
+		D:   base64.RawURLEncoding.EncodeToString(seed),
+		X:   base64.RawURLEncoding.EncodeToString(publicKey),
+	}
+
+	// Marshal to JSON
+	jwkJSON, err := json.Marshal(jwk)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal JWK: %w", err)
+	}
+
+	return string(jwkJSON), nil
+}

pkg/util/crypto_test.go

@@ -4,6 +4,7 @@ import (
 	"crypto/ed25519"
 	"crypto/x509"
 	"encoding/base64"
+	"encoding/json"
 	"encoding/pem"
 	"testing"
 
@@ -193,3 +194,141 @@ func TestConvertJWKToPEM(t *testing.T) {
 		})
 	}
 }
+
+func TestConvertPEMToJWK(t *testing.T) {
+	tests := []struct {
+		name       string
+		pemData    string
+		wantErr    bool
+		errMsg     string
+		wantKty    string
+		wantCrv    string
+		validateFn func(t *testing.T, jwkJSON string)
+	}{
+		{
+			name: "valid Ed25519 PEM key",
+			pemData: `-----BEGIN PRIVATE KEY-----
+MC4CAQAwBQYDK2VwBCIEIJ+DYvh6SEqVTm50DFtMDoQikTmiCqirVv9mWG9qfSnF
+-----END PRIVATE KEY-----`,
+			wantErr: false,
+			wantKty: "OKP",
+			wantCrv: "Ed25519",
+			validateFn: func(t *testing.T, jwkJSON string) {
+				// Parse the JWK
+				var jwk jwkKey
+				err := json.Unmarshal([]byte(jwkJSON), &jwk)
+				require.NoError(t, err)
+
+				// Verify structure
+				require.NotEmpty(t, jwk.D, "Private key component 'd' should not be empty")
+				require.NotEmpty(t, jwk.X, "Public key component 'x' should not be empty")
+
+				// Verify the private key can be decoded
+				privKeyBytes, err := base64.RawURLEncoding.DecodeString(jwk.D)
+				require.NoError(t, err)
+				assert.Len(t, privKeyBytes, ed25519.SeedSize, "Private key should be 32 bytes")
+
+				// Verify the public key can be decoded
+				pubKeyBytes, err := base64.RawURLEncoding.DecodeString(jwk.X)
+				require.NoError(t, err)
+				assert.Len(t, pubKeyBytes, ed25519.PublicKeySize, "Public key should be 32 bytes")
+			},
+		},
+		{
+			name:    "invalid PEM format",
+			pemData: "not a pem key",
+			wantErr: true,
+			errMsg:  "failed to decode PEM block",
+		},
+		{
+			name: "wrong PEM type (public key)",
+			pemData: `-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEAJrQLj5P/89iXES9+vFgrIy29clF9CC/oPPsw3c5D0bs=
+-----END PUBLIC KEY-----`,
+			wantErr: true,
+			errMsg:  "invalid PEM type",
+		},
+		{
+			name: "invalid PKCS8 data",
+			pemData: `-----BEGIN PRIVATE KEY-----
+aW52YWxpZCBkYXRh
+-----END PRIVATE KEY-----`,
+			wantErr: true,
+			errMsg:  "failed to parse PKCS#8 private key",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			jwkJSON, err := ConvertPEMToJWK(tt.pemData)
+			if tt.wantErr {
+				require.Error(t, err)
+				if tt.errMsg != "" {
+					assert.Contains(t, err.Error(), tt.errMsg)
+				}
+				return
+			}
+
+			require.NoError(t, err)
+			require.NotEmpty(t, jwkJSON)
+
+			// Parse and verify JWK structure
+			var jwk jwkKey
+			err = json.Unmarshal([]byte(jwkJSON), &jwk)
+			require.NoError(t, err, "JWK should be valid JSON")
+
+			// Verify basic fields
+			assert.Equal(t, tt.wantKty, jwk.Kty, "Key type should be OKP")
+			assert.Equal(t, tt.wantCrv, jwk.Crv, "Curve should be Ed25519")
+
+			// Run custom validation if provided
+			if tt.validateFn != nil {
+				tt.validateFn(t, jwkJSON)
+			}
+		})
+	}
+}
+
+func TestPEMToJWKRoundtrip(t *testing.T) {
+	// Test that converting PEM -> JWK -> PEM produces equivalent keys
+	originalPEM := `-----BEGIN PRIVATE KEY-----
+MC4CAQAwBQYDK2VwBCIEIJ+DYvh6SEqVTm50DFtMDoQikTmiCqirVv9mWG9qfSnF
+-----END PRIVATE KEY-----`
+
+	// Convert PEM to JWK
+	jwkJSON, err := ConvertPEMToJWK(originalPEM)
+	require.NoError(t, err)
+	require.NotEmpty(t, jwkJSON)
+
+	// Convert JWK back to PEM
+	pemData, err := ConvertJWKToPEM(jwkJSON)
+	require.NoError(t, err)
+	require.NotEmpty(t, pemData)
+
+	// Parse both keys and verify they're the same
+	block1, _ := pem.Decode([]byte(originalPEM))
+	require.NotNil(t, block1)
+	key1, err := x509.ParsePKCS8PrivateKey(block1.Bytes)
+	require.NoError(t, err)
+	ed25519Key1 := key1.(ed25519.PrivateKey)
+
+	block2, _ := pem.Decode(pemData)
+	require.NotNil(t, block2)
+	key2, err := x509.ParsePKCS8PrivateKey(block2.Bytes)
+	require.NoError(t, err)
+	ed25519Key2 := key2.(ed25519.PrivateKey)
+
+	// Compare private keys
+	assert.Equal(t, ed25519Key1, ed25519Key2, "Keys should be identical after roundtrip")
+
+	// Verify signatures match
+	message := []byte("test message for signature verification")
+	sig1 := ed25519.Sign(ed25519Key1, message)
+	sig2 := ed25519.Sign(ed25519Key2, message)
+
+	pubKey1 := ed25519Key1.Public().(ed25519.PublicKey)
+	pubKey2 := ed25519Key2.Public().(ed25519.PublicKey)
+
+	assert.True(t, ed25519.Verify(pubKey1, message, sig2), "Key1 should verify signature from Key2")
+	assert.True(t, ed25519.Verify(pubKey2, message, sig1), "Key2 should verify signature from Key1")
+}