Support multiple jwt keys for authentication (#1374)

* Add auth::parsers::tests::parse_jwt_key* tests

* Support multiple JWT keys

* Get rid of the JWT key(s) enum and use Vec

* Add gen_jwt_test_assets.py script for generating example JWTs

libsql-server/Cargo.toml

@@ -47,6 +47,7 @@ moka = { version = "0.12.1", features = ["future", "sync"] }
 nix = { version = "0.26.2", features = ["fs"] }
 once_cell = "1.17.0"
 parking_lot = "0.12.1"
+pem = "3.0.4"
 pin-project-lite = "0.2.13"
 priority-queue = "1.3"
 prost = "0.12"

libsql-server/assets/test/auth/combined123.pem

@@ -0,0 +1,9 @@
+-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEAz78yKnNWXkSQJUdQW2TU3WdelH2KtifEzg27BdtIL7c=
+-----END PUBLIC KEY-----
+-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEAKg/aT1QW16CtOKFqEvj2kvxODjWOBA6iYPDRnzrbLJ0=
+-----END PUBLIC KEY-----
+-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEARi1RRwfK5qpUf1leN9Qtjt9lRIBuNDAz6AE8PdnoSCc=
+-----END PUBLIC KEY-----

libsql-server/assets/test/auth/example1.jwt

@@ -0,0 +1 @@
+eyJhbGciOiJFZERTQSJ9.eyJwIjp7InJvIjp7Im5zIjpbImV4YW1wbGUxYSIsImV4YW1wbGUxYiIsImV4YW1wbGUxYyJdfX0sImV4cCI6MTAwMTcxNDg1Mjg1M30.q1g7TI-byMxKKZrSnoNZJZs3JfC6on_rV_-yv4ExuEE0WAy5TUePv3fg43KicLi1UEERNBSUFH4wCTKIdwE_AA

libsql-server/assets/test/auth/example1.key

@@ -0,0 +1,3 @@
+-----BEGIN PRIVATE KEY-----
+MC4CAQAwBQYDK2VwBCIEIItDMR2x6H3nXdBV7/yV6NPpFXdxSjkjQmmsnDzRW+lN
+-----END PRIVATE KEY-----

libsql-server/assets/test/auth/example1.pem

@@ -0,0 +1,3 @@
+-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEAz78yKnNWXkSQJUdQW2TU3WdelH2KtifEzg27BdtIL7c=
+-----END PUBLIC KEY-----

libsql-server/assets/test/auth/example2.jwt

@@ -0,0 +1 @@
+eyJhbGciOiJFZERTQSJ9.eyJwIjp7InJvIjp7Im5zIjpbImV4YW1wbGUyZCJdfX0sImV4cCI6MTAwMTcxNDg1Mjg1M30.AcDHnd_TfHx1HIKsUwPDaYLLzJ8-oLuBmcplHJ3U80cNUhbDFF4G8TIufCoEvOL-gXNR5Gwzw6BaMxHuwOf-Bw

libsql-server/assets/test/auth/example2.key

@@ -0,0 +1,3 @@
+-----BEGIN PRIVATE KEY-----
+MC4CAQAwBQYDK2VwBCIEIN0wLpkCrLAR851D69wdCbGtbLq6w/QqZhB16xVn3rUz
+-----END PRIVATE KEY-----

libsql-server/assets/test/auth/example2.pem

@@ -0,0 +1,3 @@
+-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEAKg/aT1QW16CtOKFqEvj2kvxODjWOBA6iYPDRnzrbLJ0=
+-----END PUBLIC KEY-----

libsql-server/assets/test/auth/example3.jwt

@@ -0,0 +1 @@
+eyJhbGciOiJFZERTQSJ9.eyJwIjp7InJvIjp7Im5zIjpbImV4YW1wbGUzZSIsImV4YW1wbGUzZiJdfX0sImV4cCI6MTAwMTcxNDg1Mjg1NH0.vgrIk9W8XbMyeMWF-b7SnqecdmRIQYh3jcucjeSkv_jI6cTD94dmeMesk_Wh1ffpnvi-kUe14-IiPmnQ9mrOAQ

libsql-server/assets/test/auth/example3.key

@@ -0,0 +1,3 @@
+-----BEGIN PRIVATE KEY-----
+MC4CAQAwBQYDK2VwBCIEILBwHAWamD8tGsPX/WaVSbxrxgEUOSMusbETkO16zSbS
+-----END PRIVATE KEY-----

libsql-server/assets/test/auth/example3.pem

@@ -0,0 +1,3 @@
+-----BEGIN PUBLIC KEY-----
+MCowBQYDK2VwAyEARi1RRwfK5qpUf1leN9Qtjt9lRIBuNDAz6AE8PdnoSCc=
+-----END PUBLIC KEY-----

libsql-server/scripts/gen_jwt_test_assets.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""utility that generates Ed25519 key and a JWT for testing
+
+the public key is stored in jwt_key.pem (in PEM format) and jwt_key.base64 (raw
+base64 format) and the JWT is printed to stdout
+"""
+import base64
+import datetime
+import jwt
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+
+def update_example(name, namespaces):
+    privkey = Ed25519PrivateKey.generate()
+    pubkey = privkey.public_key()
+
+    pubkey_pem = pubkey.public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+
+    pubkey_base64 = base64.b64encode(
+        pubkey.public_bytes(
+            encoding=serialization.Encoding.Raw,
+            format=serialization.PublicFormat.Raw,
+        ),
+        altchars=b"-_",
+    )
+    while pubkey_base64[-1] == ord("="):
+        pubkey_base64 = pubkey_base64[:-1]
+
+    privkey_pem = privkey.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    exp = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=100_000)
+    claims = {
+        "p": { "ro": { "ns": namespaces } },
+        "exp": int(exp.timestamp()),
+    }
+    token = jwt.encode(claims, privkey_pem, "EdDSA")
+    open(f"libsql-server/assets/test/auth/{name}.key", "wb").write(privkey_pem)
+    open(f"libsql-server/assets/test/auth/{name}.pem", "wb").write(pubkey_pem)
+    open(f"libsql-server/assets/test/auth/{name}.jwt", "wb").write(token.encode())
+    open(f"libsql-server/assets/test/auth/combined123.pem", "ab").write(pubkey_pem)
+
+open(f"libsql-server/assets/test/auth/combined123.pem", "wb").write("".encode())
+update_example("example1", ["example1a", "example1b", "example1c"])
+update_example("example2", ["example2d"])
+update_example("example3", ["example3e", "example3f"])

libsql-server/src/auth/authorized.rs

@@ -175,7 +175,7 @@ impl Authorized {
     }
 }
 
-#[derive(Debug, serde::Deserialize, serde::Serialize, Default)]
+#[derive(Debug, PartialEq, Eq, serde::Deserialize, serde::Serialize, Default)]
 pub struct Scopes {
     #[serde(rename = "ns", default)]
     pub namespaces: Option<HashSet<NamespaceName>>,

libsql-server/src/auth/mod.rs

@@ -11,7 +11,7 @@ pub mod user_auth_strategies;
 pub use authenticated::Authenticated;
 pub use authorized::Authorized;
 pub use errors::AuthError;
-pub use parsers::{parse_http_auth_header, parse_http_basic_auth_arg, parse_jwt_key};
+pub use parsers::{parse_http_auth_header, parse_http_basic_auth_arg, parse_jwt_keys};
 pub use permission::Permission;
 pub use user_auth_strategies::{Disabled, HttpBasic, Jwt, UserAuthContext, UserAuthStrategy};
 

libsql-server/src/auth/parsers.rs

@@ -22,16 +22,21 @@ pub fn parse_http_basic_auth_arg(arg: &str) -> Result<Option<String>> {
     }
 }
 
-pub fn parse_jwt_key(data: &str) -> Result<jsonwebtoken::DecodingKey> {
-    if data.starts_with("-----BEGIN PUBLIC KEY-----") {
-        jsonwebtoken::DecodingKey::from_ed_pem(data.as_bytes())
-            .context("Could not decode Ed25519 public key from PEM")
-    } else if data.starts_with("-----BEGIN PRIVATE KEY-----") {
-        bail!("Received a private key, but a public key is expected")
-    } else if data.starts_with("-----BEGIN") {
-        bail!("Key is in unsupported PEM format")
+pub fn parse_jwt_keys(data: &str) -> Result<Vec<jsonwebtoken::DecodingKey>> {
+    if data.starts_with("-----BEGIN") {
+        let pems = pem::parse_many(data).context("Could not parse many certificates from PEM")?;
+
+        pems.iter()
+            .map(|pem| match pem.tag() {
+                "PUBLIC KEY" => jsonwebtoken::DecodingKey::from_ed_pem(pem.to_string().as_bytes())
+                    .context("Could not decode Ed25519 public key from PEM"),
+                "PRIVATE KEY" => bail!("Received a private key, but a public key is expected"),
+                _ => bail!("Key is in unsupported PEM format"),
+            })
+            .collect()
     } else {
         jsonwebtoken::DecodingKey::from_ed_components(data)
+            .map(|v| vec![v]) // Only supports a single key
             .map_err(|e| anyhow::anyhow!("Could not decode Ed25519 public key from base64: {e}"))
     }
 }
@@ -72,7 +77,9 @@ mod tests {
     use axum::http::HeaderValue;
     use hyper::header::AUTHORIZATION;
 
-    use crate::auth::{parse_http_auth_header, AuthError};
+    use crate::auth::authorized::Scopes;
+    use crate::auth::user_auth_strategies::jwt::Token;
+    use crate::auth::{parse_http_auth_header, parse_jwt_keys, AuthError};
 
     use super::{parse_grpc_auth_header, parse_http_basic_auth_arg};
 
@@ -160,4 +167,132 @@ mod tests {
         let out = parse_http_basic_auth_arg("always").unwrap();
         assert_eq!(out, Some("always".to_string()));
     }
+
+    // Examples created via libsql-server/scripts/gen_jwt_test_assets.py
+    const EXAMPLE_JWT_PUBLIC_KEY: &str = include_str!("../../assets/test/auth/example1.pem");
+    const EXAMPLE_JWT_PRIVATE_KEY: &str = include_str!("../../assets/test/auth/example1.key");
+    const EXAMPLE_JWT: &str = include_str!("../../assets/test/auth/example1.jwt");
+    const EXAMPLE3_JWT: &str = include_str!("../../assets/test/auth/example3.jwt");
+    const MULTI_JWT_PUBLIC_KEY: &str = include_str!("../../assets/test/auth/combined123.pem");
+
+    #[test]
+    fn parse_jwt_keys_single_pem() {
+        let keys = parse_jwt_keys(EXAMPLE_JWT_PUBLIC_KEY);
+        assert_eq!(keys.as_ref().map_or(0, |v| v.len()), 1);
+        let key = keys.unwrap().into_iter().next().unwrap();
+
+        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
+        validation.required_spec_claims.remove("exp");
+        let decoded = jsonwebtoken::decode::<Token>(&EXAMPLE_JWT, &key, &validation);
+
+        assert!(matches!(
+            decoded,
+            Ok(jsonwebtoken::TokenData {
+                header: _,
+                claims: _
+            })
+        ));
+
+        let jsonwebtoken::TokenData { header, claims } = decoded.unwrap();
+        assert_eq!(header.alg, jsonwebtoken::Algorithm::EdDSA);
+
+        assert!(claims.p.is_some());
+        let Some(authorized) = claims.p else {
+            panic!("Assertion should have already failed");
+        };
+
+        let scopes: Scopes =
+            serde_json::from_str(r##"{"ns":["example1a","example1b","example1c"]}"##)
+                .expect("JSON failed to parse");
+        assert_eq!(authorized.read_only, Some(scopes));
+    }
+
+    #[test]
+    fn parse_jwt_keys_multiple_pems() {
+        let keys = parse_jwt_keys(MULTI_JWT_PUBLIC_KEY);
+        assert_eq!(keys.as_ref().map_or(0, |v| v.len()), 3);
+        let keys = keys.unwrap();
+
+        let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::EdDSA);
+        validation.required_spec_claims.remove("exp");
+        let decoded = jsonwebtoken::decode::<Token>(&EXAMPLE_JWT, &keys[0], &validation);
+
+        assert!(matches!(
+            decoded,
+            Ok(jsonwebtoken::TokenData {
+                header: _,
+                claims: _
+            })
+        ));
+
+        let jsonwebtoken::TokenData { header, claims } = decoded.unwrap();
+        assert_eq!(header.alg, jsonwebtoken::Algorithm::EdDSA);
+
+        assert!(claims.p.is_some());
+        let Some(authorized) = claims.p else {
+            panic!("Assertion should have already failed");
+        };
+
+        let scopes: Scopes =
+            serde_json::from_str(r##"{"ns":["example1a","example1b","example1c"]}"##)
+                .expect("JSON failed to parse");
+        assert_eq!(authorized.read_only, Some(scopes));
+
+        let decoded = jsonwebtoken::decode::<Token>(&EXAMPLE3_JWT, &keys[2], &validation);
+
+        assert!(matches!(
+            decoded,
+            Ok(jsonwebtoken::TokenData {
+                header: _,
+                claims: _
+            })
+        ));
+
+        let jsonwebtoken::TokenData { header, claims } = decoded.unwrap();
+        assert_eq!(header.alg, jsonwebtoken::Algorithm::EdDSA);
+
+        assert!(claims.p.is_some());
+        let Some(authorized) = claims.p else {
+            panic!("Assertion should have already failed");
+        };
+
+        let scopes: Scopes = serde_json::from_str(r##"{"ns":["example3e","example3f"]}"##)
+            .expect("JSON failed to parse");
+        assert_eq!(authorized.read_only, Some(scopes));
+    }
+
+    #[test]
+    fn parse_jwt_keys_fail_when_multiple_contains_private_key() {
+        let keys = parse_jwt_keys(
+            format!("{}\n{}", MULTI_JWT_PUBLIC_KEY, EXAMPLE_JWT_PRIVATE_KEY).as_str(),
+        );
+        assert!(keys.is_err());
+        assert_eq!(
+            keys.err().unwrap().to_string(),
+            "Received a private key, but a public key is expected"
+        );
+    }
+
+    #[test]
+    fn parse_jwt_keys_fail_when_private_key() {
+        let keys = parse_jwt_keys(EXAMPLE_JWT_PRIVATE_KEY);
+        assert!(keys.is_err());
+        assert_eq!(
+            keys.err().unwrap().to_string(),
+            "Received a private key, but a public key is expected"
+        );
+    }
+
+    #[test]
+    fn parse_jwt_keys_fail_when_non_key_pem() {
+        let keys = parse_jwt_keys(
+            "-----BEGIN CERTIFICATE-----\nMIIKLwIBAzCCCesGCSqGSIb3DQE\n-----END CERTIFICATE-----\n",
+        );
+
+        assert!(keys.is_err());
+        assert_eq!(
+            keys.err().unwrap().to_string(),
+            "Could not parse many certificates from PEM"
+        );
+    }
 }