From d82945bcc5a8168b9469b2e400dee71a9cbf9f22 Mon Sep 17 00:00:00 2001
From: Brian Aydemir <baydemir@morgridge.org>
Date: Fri, 7 Mar 2025 21:45:37 +0000
Subject: [PATCH] Fix `generate keygen` so that public keys included required
 fields

---
 cmd/generate_keygen.go      | 13 +++++--------
 cmd/generate_keygen_test.go |  8 ++++++++
 config/init_server_creds.go | 10 +++++-----
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/cmd/generate_keygen.go b/cmd/generate_keygen.go
index a26816fe2..bb8a4cc53 100644
--- a/cmd/generate_keygen.go
+++ b/cmd/generate_keygen.go
@@ -19,7 +19,6 @@
 package main
 
 import (
-	"crypto"
 	"crypto/elliptic"
 	"encoding/json"
 	"fmt"
@@ -34,11 +33,7 @@ import (
 	"github.com/pelicanplatform/pelican/config"
 )
 
-func createJWKS(privKey crypto.PrivateKey) (jwk.Set, error) {
-	key, err := jwk.FromRaw(privKey)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to generate JWK from private key")
-	}
+func createJWKS(key jwk.Key) (jwk.Set, error) {
 	jwks := jwk.NewSet()
 
 	pkey, err := jwk.PublicKeyOf(key)
@@ -47,7 +42,7 @@ func createJWKS(privKey crypto.PrivateKey) (jwk.Set, error) {
 	}
 
 	if err = jwks.AddKey(pkey); err != nil {
-		return nil, errors.Wrapf(err, "Failed to add public key %s to new JWKS", key.KeyID())
+		return nil, errors.Wrapf(err, "failed to add public key %s to new JWKS", key.KeyID())
 	}
 
 	return jwks, nil
@@ -91,7 +86,8 @@ func keygenMain(cmd *cobra.Command, args []string) error {
 	if err := config.GeneratePrivateKey(privateKeyPath, elliptic.P256(), false); err != nil {
 		return errors.Wrapf(err, "failed to generate new private key at %s", privateKeyPath)
 	}
-	privKey, err := config.LoadPrivateKey(privateKeyPath, false)
+
+	privKey, err := config.LoadSinglePEM(privateKeyPath)
 	if err != nil {
 		return errors.Wrapf(err, "failed to load private key from %s", privateKeyPath)
 	}
@@ -100,6 +96,7 @@ func keygenMain(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+
 	bytes, err := json.MarshalIndent(pubJWKS, "", "	")
 	if err != nil {
 		return errors.Wrap(err, "failed to generate json from jwks")
diff --git a/cmd/generate_keygen_test.go b/cmd/generate_keygen_test.go
index c839875e0..d57517ee0 100644
--- a/cmd/generate_keygen_test.go
+++ b/cmd/generate_keygen_test.go
@@ -60,6 +60,14 @@ func checkKeys(t *testing.T, privateKey, publicKey string) {
 	assert.True(t, ok)
 	err = key.Validate()
 	assert.NoError(t, err)
+
+	// The "alg" and "kid" keys must explicitly be added to the JWK.
+	// Thus, we test that this actually happened.
+	// See also: https://github.com/PelicanPlatform/pelican/issues/2084
+	_, ok = key.Get("alg")
+	assert.True(t, ok)
+	_, ok = key.Get("kid")
+	assert.True(t, ok)
 }
 
 func TestKeygenMain(t *testing.T) {
diff --git a/config/init_server_creds.go b/config/init_server_creds.go
index 80ed1d0e1..7aa5b9a38 100644
--- a/config/init_server_creds.go
+++ b/config/init_server_creds.go
@@ -554,7 +554,7 @@ func initKeysMap() {
 }
 
 // Helper function to load one .pem file from specified filename
-func loadSinglePEM(path string) (jwk.Key, error) {
+func LoadSinglePEM(path string) (jwk.Key, error) {
 	contents, err := os.ReadFile(path)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to read key file")
@@ -562,7 +562,7 @@ func loadSinglePEM(path string) (jwk.Key, error) {
 
 	key, err := jwk.ParseKey(contents, jwk.WithPEM(true))
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to parse issuer key file %v", path)
+		return nil, errors.Wrapf(err, "failed to parse key file %v", path)
 	}
 
 	// Add the algorithm to the key, needed for verifying tokens elsewhere
@@ -589,7 +589,7 @@ func loadPEMFiles(dir string) (jwk.Key, error) {
 	issuerKeyPath := param.IssuerKey.GetString()
 	if issuerKeyPath != "" {
 		if _, err := os.Stat(issuerKeyPath); err == nil {
-			issuerKey, err := loadSinglePEM(issuerKeyPath)
+			issuerKey, err := LoadSinglePEM(issuerKeyPath)
 			if err != nil {
 				log.Warnf("Failed to load key %s: %v", issuerKeyPath, err)
 			} else {
@@ -624,7 +624,7 @@ func loadPEMFiles(dir string) (jwk.Key, error) {
 			}
 			if dirEnt.Type().IsRegular() && filepath.Ext(dirEnt.Name()) == ".pem" {
 				// Parse the private key in this file and add to the in-memory keys map
-				key, err := loadSinglePEM(path)
+				key, err := LoadSinglePEM(path)
 				if err != nil {
 					log.Warnf("Failed to load key %s: %v", path, err)
 					return nil // Skip this file and continue
@@ -706,7 +706,7 @@ func GeneratePEM(dir string) (key jwk.Key, err error) {
 		return nil, errors.Wrapf(err, "failed to generate private key in file %s", fname)
 	}
 
-	if key, err = loadSinglePEM(fname); err != nil {
+	if key, err = LoadSinglePEM(fname); err != nil {
 		log.Errorf("Failed to load key %s: %v", fname, err)
 		err = errors.Wrapf(err, "failed to load key from %s", fname)
 		return