archify-vscode-ext-sync/crates/server/src/oidc.rs

//! Single sign-on through an OpenID Connect provider (Keycloak, Authentik,
//! Entra ID, …). The editor's sign-in stays as it is: the browser page offers
//! "Sign in with …", the provider signs the person in and comes back to
//! `/oidc/callback`, and the server hands the editor its one-time code as for
//! a password sign-in. One provider client serves every address the server is
//! reached at; each address only needs its callback listed at the provider.

use std::collections::HashMap;
use std::time::{Duration, Instant};

use anyhow::{Context, Result, anyhow, bail, ensure};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::{DecodingKey, Validation, decode, decode_header};
use serde::Deserialize;
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;

use crate::auth::random_token;
use crate::pages::AuthorizeQuery;

/// How long the provider may take to send the person back.
const PENDING_TTL: Duration = Duration::from_secs(10 * 60);
const PENDING_MAX: usize = 10_000;
const CACHE_TTL: Duration = Duration::from_secs(60 * 60);
/// Unknown signing keys trigger a fresh key set at most this often.
const JWKS_REFRESH_MIN: Duration = Duration::from_secs(60);

#[derive(Debug, Clone)]
pub struct OidcConfig {
    /// `https://keycloak.example.com/realms/home`, as the provider names itself.
    pub issuer: String,
    pub client_id: String,
    /// Confidential clients (Keycloak's default) have one; public clients rely on PKCE alone.
    pub client_secret: Option<String>,
    /// The button: "Sign in with {label}".
    pub label: String,
    pub scopes: String,
    /// The claim that lists the person's groups (Keycloak's group mapper: `groups`).
    pub groups_claim: String,
    /// When set, only members of one of these groups may sign in, and they
    /// get an account on first sign-in without an invite.
    pub allowed_groups: Vec<String>,
    /// Members of these groups are admins of the server.
    pub admin_groups: Vec<String>,
}

/// Who the provider says signed in.
#[derive(Debug, Clone)]
pub struct Person {
    pub subject: String,
    /// `preferred_username`, or the part of the e-mail before the `@`.
    pub username: Option<String>,
    pub display_name: Option<String>,
    pub email: Option<String>,
    pub groups: Vec<String>,
}

impl Person {
    pub fn in_any(&self, groups: &[String]) -> bool {
        self.groups.iter().any(|g| groups.contains(g))
    }
}

#[derive(Debug, Clone, Deserialize)]
struct Discovery {
    issuer: String,
    authorization_endpoint: String,
    token_endpoint: String,
    jwks_uri: String,
}

struct Pending {
    query: AuthorizeQuery,
    nonce: String,
    verifier: String,
    created: Instant,
}

#[derive(Deserialize)]
struct TokenReply {
    id_token: String,
}

pub struct Oidc {
    pub config: OidcConfig,
    http: reqwest::Client,
    discovery: Mutex<Option<(Instant, Discovery)>>,
    jwks: Mutex<Option<(Instant, JwkSet)>>,
    pending: Mutex<HashMap<String, Pending>>,
}

impl Oidc {
    pub fn new(config: OidcConfig) -> Result<Self> {
        let http = reqwest::Client::builder()
            .timeout(Duration::from_secs(10))
            .redirect(reqwest::redirect::Policy::none())
            .build()?;
        Ok(Self {
            config,
            http,
            discovery: Mutex::new(None),
            jwks: Mutex::new(None),
            pending: Mutex::new(HashMap::new()),
        })
    }

    /// Where the provider sends people back; list it in the provider's client.
    pub fn callback_url(public_url: &str) -> String {
        format!("{public_url}/oidc/callback")
    }

    async fn discovery(&self) -> Result<Discovery> {
        let mut cached = self.discovery.lock().await;
        if let Some((at, d)) = cached.as_ref()
            && at.elapsed() < CACHE_TTL
        {
            return Ok(d.clone());
        }
        let url = format!(
            "{}/.well-known/openid-configuration",
            self.config.issuer.trim_end_matches('/')
        );
        let d: Discovery = self
            .http
            .get(&url)
            .send()
            .await
            .and_then(|r| r.error_for_status())
            .with_context(|| format!("the identity provider does not answer at {url}"))?
            .json()
            .await
            .context("the identity provider's configuration does not read")?;
        ensure!(
            d.issuer.trim_end_matches('/') == self.config.issuer.trim_end_matches('/'),
            "the identity provider calls itself {}, not {}",
            d.issuer,
            self.config.issuer
        );
        *cached = Some((Instant::now(), d.clone()));
        Ok(d)
    }

    async fn key_for(&self, kid: Option<&str>) -> Result<DecodingKey> {
        let d = self.discovery().await?;
        let mut cached = self.jwks.lock().await;
        let find = |set: &JwkSet| match kid {
            Some(kid) => set.find(kid).cloned(),
            None => set.keys.first().cloned(),
        };
        let stale = cached.as_ref().is_none_or(|(at, set)| {
            at.elapsed() >= CACHE_TTL || (find(set).is_none() && at.elapsed() >= JWKS_REFRESH_MIN)
        });
        if stale {
            let set: JwkSet = self
                .http
                .get(&d.jwks_uri)
                .send()
                .await
                .and_then(|r| r.error_for_status())
                .context("the identity provider's signing keys do not load")?
                .json()
                .await
                .context("the identity provider's signing keys do not read")?;
            *cached = Some((Instant::now(), set));
        }
        let (_, set) = cached.as_ref().expect("filled above");
        let jwk = find(set).ok_or_else(|| anyhow!("the sign-in was signed with an unknown key"))?;
        Ok(DecodingKey::from_jwk(&jwk)?)
    }

    /// The provider's sign-in address for this editor sign-in.
    pub async fn start(&self, public_url: &str, query: AuthorizeQuery) -> Result<String> {
        let d = self.discovery().await?;
        let state = random_token();
        let nonce = random_token();
        let verifier = random_token();
        let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
        {
            let mut pending = self.pending.lock().await;
            pending.retain(|_, p| p.created.elapsed() < PENDING_TTL);
            ensure!(
                pending.len() < PENDING_MAX,
                "too many sign-ins are waiting; try again in a few minutes"
            );
            pending.insert(
                state.clone(),
                Pending {
                    query,
                    nonce: nonce.clone(),
                    verifier,
                    created: Instant::now(),
                },
            );
        }
        let mut url = url::Url::parse(&d.authorization_endpoint)
            .context("the identity provider's sign-in address is not valid")?;
        url.query_pairs_mut()
            .append_pair("response_type", "code")
            .append_pair("client_id", &self.config.client_id)
            .append_pair("redirect_uri", &Self::callback_url(public_url))
            .append_pair("scope", &self.config.scopes)
            .append_pair("state", &state)
            .append_pair("nonce", &nonce)
            .append_pair("code_challenge", &challenge)
            .append_pair("code_challenge_method", "S256");
        Ok(url.into())
    }

    /// The provider sent the person back: exchanges the code and checks the ID token.
    pub async fn finish(
        &self,
        public_url: &str,
        code: &str,
        state: &str,
    ) -> Result<(AuthorizeQuery, Person)> {
        let pending = self
            .pending
            .lock()
            .await
            .remove(state)
            .filter(|p| p.created.elapsed() < PENDING_TTL)
            .ok_or_else(|| anyhow!("this sign-in has expired; start again from VS Code"))?;
        let d = self.discovery().await?;
        let redirect_uri = Self::callback_url(public_url);
        let mut form = vec![
            ("grant_type", "authorization_code"),
            ("code", code),
            ("redirect_uri", redirect_uri.as_str()),
            ("client_id", self.config.client_id.as_str()),
            ("code_verifier", pending.verifier.as_str()),
        ];
        if let Some(secret) = &self.config.client_secret {
            form.push(("client_secret", secret.as_str()));
        }
        let response = self
            .http
            .post(&d.token_endpoint)
            .form(&form)
            .send()
            .await
            .context("the identity provider does not answer")?;
        if !response.status().is_success() {
            let status = response.status();
            let body: Value = response.json().await.unwrap_or(Value::Null);
            // The provider's own words (invalid_client, …); never the request itself.
            bail!(
                "the identity provider refused the sign-in ({status}{})",
                body.get("error")
                    .and_then(Value::as_str)
                    .map(|e| format!(", {e}"))
                    .unwrap_or_default()
            );
        }
        let reply: TokenReply = response
            .json()
            .await
            .context("the identity provider sent no ID token")?;
        let person = self
            .check_id_token(&d, &reply.id_token, &pending.nonce)
            .await?;
        Ok((pending.query, person))
    }

    async fn check_id_token(&self, d: &Discovery, token: &str, nonce: &str) -> Result<Person> {
        let header = decode_header(token).context("the ID token does not read")?;
        let key = self.key_for(header.kid.as_deref()).await?;
        let mut validation = Validation::new(header.alg);
        validation.set_issuer(&[d.issuer.as_str()]);
        validation.set_audience(&[self.config.client_id.as_str()]);
        validation.set_required_spec_claims(&["exp", "iss", "aud", "sub"]);
        validation.leeway = 60;
        let claims: HashMap<String, Value> = decode(token, &key, &validation)
            .context("the ID token is not valid")?
            .claims;
        ensure!(
            claims.get("nonce").and_then(Value::as_str) == Some(nonce),
            "the ID token belongs to another sign-in"
        );
        Ok(person_from(&claims, &self.config.groups_claim))
    }
}

fn person_from(claims: &HashMap<String, Value>, groups_claim: &str) -> Person {
    let text = |key: &str| {
        claims
            .get(key)
            .and_then(Value::as_str)
            .map(str::trim)
            .filter(|s| !s.is_empty())
            .map(str::to_owned)
    };
    let email = text("email");
    let groups = match claims.get(groups_claim) {
        Some(Value::Array(items)) => items
            .iter()
            .filter_map(Value::as_str)
            // Keycloak sends "/team" when the mapper keeps full paths.
            .map(|g| g.trim_start_matches('/').to_owned())
            .collect(),
        Some(Value::String(one)) => vec![one.trim_start_matches('/').to_owned()],
        _ => vec![],
    };
    Person {
        subject: text("sub").unwrap_or_default(),
        username: text("preferred_username").or_else(|| {
            email
                .as_deref()
                .and_then(|e| e.split('@').next())
                .map(str::to_owned)
        }),
        display_name: text("name"),
        email,
        groups,
    }
}

/// A sign-in name from what the provider calls the person: letters, digits and `. _ - @`.
pub fn account_name(person: &Person) -> String {
    let raw = person.username.as_deref().unwrap_or("user");
    let mut name: String = raw
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-' | '@') {
                c.to_ascii_lowercase()
            } else {
                '-'
            }
        })
        .take(56)
        .collect();
    if name.len() < 2 {
        name = format!("user-{name}");
    }
    name
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn people_from_keycloak_claims() {
        let claims: HashMap<String, Value> = serde_json::from_value(serde_json::json!({
            "sub": "f3c1",
            "preferred_username": "Ana Šmit",
            "name": "Ana Šmit",
            "email": "ana@example.com",
            "groups": ["/archify-users", "admins"],
        }))
        .unwrap();
        let person = person_from(&claims, "groups");
        assert_eq!(person.subject, "f3c1");
        assert_eq!(person.groups, ["archify-users", "admins"]);
        assert!(person.in_any(&["archify-users".into()]));
        assert!(!person.in_any(&["others".into()]));
        assert_eq!(account_name(&person), "ana--mit");

        let bare: HashMap<String, Value> =
            serde_json::from_value(serde_json::json!({ "sub": "x", "email": "bo@example.com" }))
                .unwrap();
        let person = person_from(&bare, "groups");
        assert_eq!(account_name(&person), "bo");
        assert!(person.groups.is_empty());
    }
}