JWK sets

[einrobin]

Addresses #734

Still a draft since I wanna do some more testing myself, add unit tests and also write some documentation on it.

The config looks something like this:

    [auth]
    jwks_uri = "https://my-website/keys.json"
    jwks_refresh_interval = 3600 # comment out to disable continous refreshing
    # key = "/app/public.jwk"
    public = "anon"

[coderabbitai]

Walkthrough

Adds JWKS-based JWT authentication: new doc doc/guide/authentication.md; introduces KeySet in moq-token with (de)serialization, file I/O, encode/decode, and a feature-gated async load_keys (jwks-loader, optional reqwest); enables that feature in moq-relay. Updates relay auth to accept JWKS URIs, store keys as Arc<Mutex<KeySet>>, add jwks_refresh_interval to AuthConfig, perform initial JWKS load, spawn an optional periodic refresh task, adjust verification to lock the key set, and implement Drop to abort the refresh task. Main startup now awaits async auth initialization.

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.91% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title 'JWK sets' is vague and generic, using a technical acronym without context about what the changeset accomplishes.	Use a more descriptive title like 'Add JWKS authentication support with refresh capability' to clearly convey the main objective.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description is related to the changeset, addressing issue #734 and providing example configuration for JWKS URI and refresh interval support.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 5

🤖 Fix all issues with AI agents

In `@doc/guide/authentication.md`:
- Line 82: Tighten the sentence that currently reads "a very simple static site
serving a JSON file, or a fully OIDC compliant Identity Provider." by removing
the weak qualifier "very" and applying hyphenation: change it to "a simple
static site serving a JSON file, or a fully OIDC-compliant identity provider."
Update the sentence that begins "Instead of storing a public key locally in a
file..." to use this revised phrase so the doc uses clearer wording and correct
hyphenation.
- Around line 14-17: The symmetric-key example currently uses the asymmetric
algorithm RS256 in the command `moq-token --key root.jwk generate --algorithm
RS256`; change that algorithm to a symmetric HMAC variant (e.g., replace RS256
with HS256, or use HS384/HS512 as appropriate) so the symmetric-key example uses
HS* algorithms instead of RS*. Ensure the `--algorithm` value in the example
command is updated accordingly.

In `@rs/moq-relay/src/auth.rs`:
- Line 71: The constant JWKS_REFRESH_ERROR_INTERVAL uses a non-existent
constructor Duration::from_mins; change it to use a valid std::time::Duration
constructor (e.g., replace Duration::from_mins(5) with Duration::from_secs(5 *
60) or Duration::from_secs(300)) so the constant compiles; update the
declaration of JWKS_REFRESH_ERROR_INTERVAL accordingly.
- Around line 80-85: The Drop impl on Auth currently unconditionally calls abort
on the shared refresh_task, which kills the task when any clone is dropped;
change it to only abort when the dropping clone is the last reference by
checking Arc::strong_count(handle) == 1. Locate the Drop impl for Auth and
replace the unconditional abort with: get the Arc via
self.refresh_task.as_ref(), call Arc::strong_count(handle) and only call
handle.abort() if the count is 1 (use handle.abort() on the Arc's inner
JoinHandle reference). Ensure you reference Auth, refresh_task,
Arc::strong_count and abort in the fix.

In `@rs/moq-token/src/set.rs`:
- Around line 121-135: The load_keys function uses reqwest::get which has no
timeout; change it to use a reqwest::Client with an explicit timeout (e.g.,
Duration::from_secs) and reuse that client rather than calling reqwest::get each
time: build or accept a Client and replace the reqwest::get(jwks_uri).await...
chain with client.get(jwks_uri).send().await... (then error_for_status/text as
before) so the JWKS fetch in load_keys (and the KeySet::from_str step) cannot
block indefinitely; consider providing the Client via argument or a
shared/static so subsequent calls reuse the same client.

♻️ Duplicate comments (1)

rs/moq-relay/src/auth.rs (1)

139-169: Ensure the initial JWKS load succeeds before serving requests.

Auth::new returns immediately with an empty KeySet, so early requests will fail auth until the async refresh completes; if the first load fails, the service continues with an empty set. Consider gating startup on a successful initial fetch (e.g., async constructor or explicit init path) and surfacing errors.

🧹 Nitpick comments (1)

rs/moq-token/Cargo.toml (1)

11-22: Verify the reqwest pre-release choice and TLS backend.

reqwest 0.13.0-rc.1 is a release candidate and default features pull in native-tls; please confirm this is intentional for the JWKS loader and consider pinning to a stable release and explicitly selecting the TLS backend to avoid surprises.

[coderabbitai]

Actionable comments posted: 2

🤖 Fix all issues with AI agents

In `@rs/moq-relay/src/auth.rs`:
- Around line 151-162: Auth::new currently spawns the background refresh task
(via Self::spawn_refresh_task) when a jwks_refresh_interval is provided but
returns immediately with an empty KeySet, causing requests to fail until the
first background refresh; change the flow so you perform an initial synchronous
key fetch/population on the KeySet (call whatever synchronous fetch/parse helper
used elsewhere, e.g., the same logic used by spawn_refresh_task or a shared
fetch_keys function) and populate key_set before assigning refresh_task and
returning; ensure the initial fetch honors jwks_refresh_interval/error handling
and only then spawn Self::spawn_refresh_task with key_set.clone() so Auth::new
never returns an empty KeySet.
- Around line 163-170: Auth::new currently spawns refresh_key_set(&uri,
key_set.as_ref()) when jwks_refresh_interval is None, returning immediately with
an empty KeySet and panicking on background failure; change this so the initial
JWKS load is awaited and errors are propagated or retried instead of
spawn+expect. Specifically, in the branch that now does tokio::spawn, call
Self::refresh_key_set(&uri, key_set.as_ref()).await (or return a future that
completes when refresh_key_set succeeds) and handle errors by returning them
from Auth::new or by implementing the same retry/backoff logic used in the
periodic refresh path so startup blocks until keys are loaded and no background
panic can occur. Ensure references to KeySet, uri, jwks_refresh_interval, and
refresh_key_set are used to locate the code to modify.

♻️ Duplicate comments (1)

doc/guide/authentication.md (1)

80-87: Minor grammar fix: hyphenate "OIDC-compliant".

Per standard English grammar, compound adjectives before a noun should be hyphenated.

✏️ Suggested fix

-Instead of storing a public key locally in a file, it may also be retrieved from a server hosting a JWK set. This can be a simple static site serving a JSON file, or a fully OIDC compliant Identity Provider. That way you can easily implement automatic key rotation.
+Instead of storing a public key locally in a file, it may also be retrieved from a server hosting a JWK set. This can be a simple static site serving a JSON file, or a fully OIDC-compliant Identity Provider. That way you can easily implement automatic key rotation.

🧹 Nitpick comments (2)

rs/moq-token/src/set.rs (1)

83-88: Consider simplifying the find_key predicate.

The is_some() check followed by unwrap() can be simplified.

♻️ Suggested simplification

 pub fn find_key(&self, kid: &str) -> Option<Arc<Key>> {
     self.keys
         .iter()
-        .find(|k| k.kid.is_some() && k.kid.as_deref().unwrap() == kid)
+        .find(|k| k.kid.as_deref() == Some(kid))
         .cloned()
 }

rs/moq-relay/src/auth.rs (1)

91-107: Consider logging removed keys for better observability.

compare_key_sets logs newly added keys but doesn't log when keys are removed from the set. This could help with debugging key rotation issues.

♻️ Suggested enhancement

 fn compare_key_sets(previous: &KeySet, new: &KeySet) {
     for new_key in new.keys.iter() {
         if new_key.kid.is_some() && !previous.keys.iter().any(|k| k.kid == new_key.kid) {
             tracing::info!("Found new JWK \"{}\"", new_key.kid.as_deref().unwrap())
         }
     }
+    for old_key in previous.keys.iter() {
+        if old_key.kid.is_some() && !new.keys.iter().any(|k| k.kid == old_key.kid) {
+            tracing::info!("Removed JWK \"{}\"", old_key.kid.as_deref().unwrap())
+        }
+    }
 }

[einrobin]

Finally found time again today, sorry for the delay. Everything works fine, tested with hang-demo and this config:

[auth]
key = "https://<imagine my auth service>"
jwks_refresh_interval = 30