initial commit

src/api/client_server/sso.rs

@@ -0,0 +1,648 @@
+use std::{borrow::Borrow, collections::HashMap, iter::Iterator, time::SystemTime};
+
+use crate::{
+    config::{
+        sso::{Registration, Template},
+        IdpConfig,
+    },
+    service::sso::{
+        templates, LoginToken, RegistrationInfo, RegistrationToken, ValidationData,
+        REGISTRATION_EXPIRATION_SECS, SESSION_EXPIRATION_SECS, SSO_AUTH_EXPIRATION_SECS,
+        SSO_SESSION_COOKIE,
+    },
+    services, utils, Error, Result, Ruma,
+};
+use axum::{
+    extract::RawQuery,
+    response::{AppendHeaders, IntoResponse, Redirect},
+    RequestExt,
+};
+use axum_extra::{
+    headers::{self, HeaderMapExt},
+    TypedHeader,
+};
+use http::header;
+use mas_oidc_client::{
+    requests::{
+        authorization_code::{self, AuthorizationRequestData, AuthorizationValidationData},
+        jose::{self, JwtVerificationData},
+        userinfo,
+    },
+    types::{
+        client_credentials::ClientCredentials,
+        errors::ClientError,
+        iana::jose::JsonWebSignatureAlg,
+        requests::{AccessTokenResponse, AuthorizationResponse},
+    },
+};
+use rand::{rngs::StdRng, SeedableRng};
+use ruma::{
+    api::client::{
+        error::ErrorKind,
+        session::{self, sso_login, sso_login_with_provider},
+    },
+    events::GlobalAccountDataEventType,
+    push, OwnedMxcUri, UserId,
+};
+use serde_json::Number;
+use tracing::error;
+use url::Url;
+
+pub const CALLBACK_PATH: &str = "_matrix/client/unstable/sso/callback";
+
+/// # `GET /_matrix/client/v3/login/sso/redirect`
+///
+/// Redirect the user to the SSO interface.
+/// TODO: this should be removed once Ruma supports trailing slashes.
+pub async fn get_sso_redirect(
+    body: Ruma<sso_login::v3::Request>,
+) -> Result<sso_login::v3::Response> {
+    let sso_login_with_provider::v3::Response { location, cookie } =
+        get_sso_redirect_with_provider(
+            Ruma {
+                body: sso_login_with_provider::v3::Request::new(
+                    Default::default(),
+                    body.redirect_url.clone(),
+                ),
+                ..body
+            }
+            .into(),
+        )
+        .await?;
+
+    Ok(sso_login::v3::Response { location, cookie })
+}
+
+/// # `GET /_matrix/client/v3/login/sso/redirect/{idpId}`
+///
+/// Redirects the user to the SSO interface.
+pub async fn get_sso_redirect_with_provider(
+    body: Ruma<sso_login_with_provider::v3::Request>,
+) -> Result<sso_login_with_provider::v3::Response> {
+    let idp_ids: Vec<&str> = services()
+        .globals
+        .config
+        .idps
+        .iter()
+        .map(Borrow::borrow)
+        .collect();
+
+    let provider = match &*idp_ids {
+        [] => {
+            return Err(Error::BadRequest(
+                ErrorKind::forbidden(),
+                "Single Sign-On is disabled.",
+            ));
+        }
+        [idp_id] => services().sso.get(idp_id).expect("we know it exists"),
+        [_, ..] => services().sso.get(&body.idp_id).ok_or_else(|| {
+            Error::BadRequest(ErrorKind::InvalidParam, "Unknown identity provider.")
+        })?,
+    };
+
+    let redirect_url = body
+        .redirect_url
+        .parse::<Url>()
+        .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Invalid redirect_url."))?;
+
+    let mut callback = services()
+        .globals
+        .well_known_client()
+        .parse::<Url>()
+        .map_err(|_| Error::bad_config("Invalid well_known_client url."))?;
+    callback.set_path(CALLBACK_PATH);
+
+    let (auth_url, validation_data) = authorization_code::build_authorization_url(
+        provider.metadata.authorization_endpoint().clone(),
+        AuthorizationRequestData::new(
+            provider.config.client_id.clone(),
+            provider.config.scopes.clone(),
+            redirect_url,
+        ),
+        &mut StdRng::from_entropy(),
+    )
+    .map_err(|_| Error::BadRequest(ErrorKind::Unknown, "Failed to build authorization_url."))?;
+
+    let signed = services().globals.sign_claims(&ValidationData::new(
+        provider.borrow().to_string(),
+        validation_data,
+    ));
+
+    Ok(sso_login_with_provider::v3::Response {
+        location: auth_url.to_string(),
+        cookie: Some(
+            utils::build_cookie(
+                SSO_SESSION_COOKIE,
+                &signed,
+                "/_conduit/client/sso/callback",
+                Some(SSO_AUTH_EXPIRATION_SECS),
+            )
+            .to_string(),
+        ),
+    })
+}
+
+/// # `GET /_conduit/client/sso/callback`
+///
+/// Validate the authorization response received from the identity provider.
+/// On success, generate a login token, add it to `redirectUrl` as a query and perform the redirect.
+/// If this is the first login, register the user, possibly interactively through a fallback page.
+pub async fn get_sso_callback(req: axum::extract::Request) -> Result<axum::response::Response> {
+    let query = req.uri().query().ok_or_else(|| {
+        Error::BadRequest(ErrorKind::MissingParam, "Empty authorization callback.")
+    })?;
+
+    let AuthorizationResponse {
+        code,
+        access_token,
+        token_type,
+        id_token,
+        expires_in,
+    } = serde_html_form::from_str::<AuthorizationResponse>(query).map_err(|_| {
+        serde_html_form::from_str::<ClientError>(query).unwrap_or_else(|_| {
+            error!("Failed to deserialize authorization callback: {}", callback);
+
+            Error::BadRequest(
+                ErrorKind::Unknown,
+                "Failed to deserialize authorization callback.",
+            )
+        })
+    })?;
+
+    let cookie = req
+        .extract::<Option<TypedHeader<headers::Cookie>>>()
+        .await
+        .map_err(|_| Error::BadRequest(ErrorKind::InvalidParam, "Invalid session cookie."))?
+        .ok_or_else(|_| Error::BadRequest(ErrorKind::MissingParam, "Missing session cookie."))?;
+
+    let ValidationData {
+        provider,
+        inner: validation_data,
+    } = services()
+        .globals
+        .validate_claims(
+            cookie.get(SSO_SESSION_COOKIE).ok_or_else(|| {
+                Error::BadRequest(ErrorKind::MissingParam, "Missing value for session cookie.")
+            })?,
+            None,
+        )
+        .map_err(|e| {
+            Error::BadRequest(ErrorKind::InvalidParam, "Invalid value for session cookie.")
+        })?;
+
+    let provider = services().sso.get(&provider).ok_or_else(|e| {
+        Error::BadRequest(
+            ErrorKind::InvalidParam,
+            "Unknown provider for session cookie.",
+        )
+    })?;
+
+    let IdpConfig {
+        client_id,
+        client_secret,
+        auth_method,
+        ..
+    } = provider.config.clone();
+
+    let credentials = match &auth_method {
+        "basic" => ClientCredentials::ClientSecretBasic {
+            client_id,
+            client_secret,
+        },
+        "post" => ClientCredentials::ClientSecretPost {
+            client_id,
+            client_secret,
+        },
+        _ => todo!(),
+    };
+    let (
+        AccessTokenResponse {
+            access_token,
+            refresh_token,
+            token_type,
+            expires_in,
+            scope,
+            ..
+        },
+        Some(id_token),
+    ) = authorization_code::access_token_with_authorization_code(
+        services().sso.service(),
+        method,
+        provider.metadata.token_endpoint(),
+        code,
+        validation_data,
+        jwt_verification_data,
+        SystemTime::now().into(),
+        &mut StdRng::from_entropy(),
+    )
+    .await
+    .map_err(|e| Error::bad_config("Failed to fetch access token."))?
+    else {
+        unreachable!("ID token should never be empty")
+    };
+
+    // let userinfo = provider.fetch_userinfo(&access_token, &id_token).await?;
+
+    let mut userinfo = HashMap::default();
+    if let Some(endpoint) = &provider.metadata.userinfo_endpoint {
+        let ref jwks = jose::fetch_jwks(services().sso.service(), provider.metadata.jwks_uri())
+            .await
+            .map_err(|e| Error::bad_config("Failed to fetch signing keys for token endpoint."))?;
+        let jwt_verification_data = Some(JwtVerificationData {
+            jwks,
+            issuer: &provider.config.issuer,
+            client_id: credentials.client_id(),
+            signing_algorithm: &JsonWebSignatureAlg::Rs256,
+        });
+
+        userinfo = userinfo::fetch_userinfo(
+            services().sso.service(),
+            endpoint,
+            &access_token,
+            jwt_verification_data,
+            &id_token,
+        )
+        .await
+        .map_err(|e| Error::bad_config("Failed to fetch claims for userinfo endpoint."))?;
+    };
+
+    let (_, mut claims) = id_token.into_parts();
+
+    let subject = claims.get("sub").ok_or_else(|| {
+        error!("Unique \"sub\" claim is missing from ID token: {claims:?}");
+
+        Error::bad_config("Unique \"sub\" claim is missing from ID token.")
+    })?;
+
+    let subject = &subject
+        .as_str()
+        .map(str::to_owned)
+        .or_else(|| subject.as_number().map(Number::to_string))
+        .expect("unique claim should be a string or number");
+
+    let redirect_uri = &validation_data.redirect_uri;
+
+    if let Some(user_id) = services()
+        .sso
+        .user_from_claim(&validation_data.provider_id, subject)?
+    {
+        let login_token = LoginToken::new(validation_data.provider_id.to_owned(), user_id);
+
+        let redirect_uri = redirect_with_login_token(redirect_uri.to_owned(), &login_token);
+
+        return Ok((
+            AppendHeaders(vec![(
+                header::SET_COOKIE,
+                utils::reset_cookie("sso-session").to_string(),
+            )]),
+            Redirect::temporary(redirect_uri.as_str()),
+        )
+            .into_response());
+    }
+
+    match provider.config.registration {
+        Registration::Disabled => {
+            return Err(Error::BadRequest(
+                ErrorKind::forbidden(),
+                "Single Sign-On registration is disabled.",
+            ))
+        }
+        Registration::Automated => todo!(),
+        Registration::Interactive => {}
+    };
+
+    let Template {
+        username,
+        displayname,
+        avatar_url,
+        email,
+    } = &provider.config.template;
+    let registration_info =
+        RegistrationInfo::new(&claims, username, displayname, avatar_url, email);
+
+    let signed = services()
+        .globals
+        .sign_macaroon(&RegistrationToken::new(
+            validation_data.provider_id.clone(),
+            subject.to_owned(),
+            redirect_uri.to_owned(),
+            registration_info,
+        ))
+        .expect("signing macaroons always works");
+
+    let cookie = utils::build_cookie(
+        "sso-registration",
+        &signed,
+        "/_conduit/client/sso/register",
+        REGISTRATION_EXPIRATION_SECS,
+    );
+
+    Ok((
+        AppendHeaders(vec![
+            (header::SET_COOKIE, cookie.to_string()),
+            (
+                header::SET_COOKIE,
+                utils::reset_cookie("sso-session").to_string(),
+            ),
+        ]),
+        Redirect::temporary("/_conduit/client/sso/register"),
+    )
+        .into_response())
+}
+
+/// # `GET /_conduit/client/sso/pick_idp`
+pub async fn pick_idp(RawQuery(query): RawQuery) -> impl IntoResponse {
+    let providers: Vec<_> = services()
+        .globals
+        .config
+        .sso
+        .iter()
+        .map(|p| p.inner.to_owned())
+        .collect();
+
+    let body = maud::html! {
+        header {
+            h1 { "Log in to " (services().globals.server_name()) }
+            p { "Choose an identity provider to continue" }
+        }
+        main {
+            ul .providers {
+                @for provider in providers {
+                    li {
+                        a href={ "/_matrix/client/v3/login/sso/redirect/" (provider.id) "?" (query.as_deref().unwrap_or_default()) } {
+                            @if let Some(url) = provider.icon.as_deref().and_then(utils::mxc_to_http) {
+                                img src=(url);
+                            }
+                        }
+                        span {
+                            (provider.name)
+                        }
+                    }
+                }
+            }
+        }
+    };
+
+    (
+        [(header::CONTENT_TYPE, "text/html; charset=utf-8")],
+        maud::html! {
+            (templates::base("Pick Identity Provider", body))
+
+            (templates::footer())
+        },
+    )
+}
+
+/// # `GET /_conduit/client/sso/register`
+///
+/// Serve a registration form with defaults based on the retrieved claims.
+/// This endpoint is only available when interactive registration is enabled.
+pub async fn get_sso_registration(
+    cookie: TypedHeader<headers::Cookie>,
+) -> Result<axum::response::Response> {
+    let token = cookie.get("sso-registration").ok_or_else(|| {
+        Error::BadRequest(
+            ErrorKind::MissingParam,
+            "Missing registration token cookie.",
+        )
+    })?;
+
+    let registration_token: RegistrationToken = services()
+        .globals
+        .validate_macaroon(token, None)
+        .map_err(|_| {
+        Error::BadRequest(
+            ErrorKind::InvalidParam,
+            "Invalid registration token cookie.",
+        )
+    })?;
+
+    let provider = services()
+        .sso
+        .get(&registration_token.provider_id)
+        .map(|p| p.config.inner.to_owned())?;
+    let server_name = services().globals.server_name();
+
+    let RegistrationInfo {
+        username,
+        displayname,
+        avatar_url,
+        email,
+    } = registration_token.info;
+
+    let additional_info = (&displayname, &avatar_url, &email) != (&None, &None, &None);
+
+    fn detail(title: &str, body: maud::Markup) -> maud::Markup {
+        maud::html! {
+            label .detail for=(title) {
+                div .check-row {
+                    span .name { (title) } " "
+                    span .use { "use" }
+                    input #(title) type="checkbox" name={(title)"-checkbox"} value=(true) checked;
+                }
+                (body)
+            }
+        }
+    }
+
+    let body = maud::html! {
+        header {
+            h1 { "Complete your registration at " (server_name) }
+            p { "Confirm your details to finish creating your account." }
+        }
+        main {
+            form .form #form method="post" {
+                div .username-div #username-div {
+                    label for="username-input" { "Username (required)" }
+                    div .prefix { "@" }
+                    input .username-input type="text" name="username"
+                        value=(username) autofocus autocorrect="off" autocapitalize="none";
+                    div .postfix { ":" (server_name) }
+                }
+                output .username-output for="username-input" { }
+
+                @if additional_info {
+                    section .additional-info {
+                        h2 {
+                            @if let Some(icon) = provider.icon.as_deref().and_then(utils::mxc_to_http) {
+                                img src=(icon.to_string());
+                            }
+                            "Optional data from " (provider.name)
+                        }
+                        @if let Some(avatar_url) = avatar_url.as_ref() {
+                            (detail("avatar", maud::html!{
+                                img .avatar src=(avatar_url);
+                            }))
+                        }
+                        @if let Some(displayname) = displayname.as_ref() {
+                            (detail("displayname", maud::html!{
+                                p .value { (displayname) };
+                            }))
+                        }
+                        @if let Some(email) = email.as_ref() {
+                            (detail("email", maud::html!{
+                                p .value { (email) };
+                            }))
+                        }
+                    }
+                }
+
+                input type="submit" value="Submit" .primary-button {}
+            }
+        }
+    };
+
+    Ok((
+        [(header::CONTENT_TYPE, "text/html; charset=utf-8")],
+        maud::html! {
+            (templates::base("Register Account", body))
+
+            (templates::footer())
+        },
+    )
+        .into_response())
+}
+
+/// # `POST /_conduit/client/sso/register`
+///
+/// Submit the registration form.
+pub async fn submit_sso_registration(
+    cookie: TypedHeader<headers::Cookie>,
+    axum::extract::Form(registration_info): axum::extract::Form<RegistrationInfo>,
+) -> Result<axum::response::Response> {
+    let token = cookie.get("sso-registration").ok_or_else(|| {
+        Error::BadRequest(
+            ErrorKind::MissingParam,
+            "Missing registration token cookie.",
+        )
+    })?;
+
+    let registration_token: RegistrationToken = services()
+        .globals
+        .validate_macaroon(token, None)
+        .map_err(|_| {
+        Error::BadRequest(
+            ErrorKind::MissingParam,
+            "Invalid registration token cookie.",
+        )
+    })?;
+
+    let RegistrationInfo {
+        username,
+        mut displayname,
+        avatar_url,
+        email: _,
+    } = registration_info;
+
+    let user_id =
+        UserId::parse_with_server_name(username.to_lowercase(), services().globals.server_name())
+            .map_err(|_| Error::BadRequest(ErrorKind::InvalidUsername, "Invalid username."))?;
+
+    if services().users.exists(&user_id)? {
+        return Err(Error::BadRequest(
+            ErrorKind::UserInUse,
+            "Desired UserId is already taken.",
+        ));
+    }
+
+    if services().appservice.is_exclusive_user_id(&user_id).await {
+        return Err(Error::BadRequest(
+            ErrorKind::Exclusive,
+            "Desired UserId reserved by appservice.",
+        ));
+    }
+
+    services().users.create(&user_id, None)?;
+    services().users.set_password_placeholder(&user_id)?;
+
+    if let Some(avatar_url) = avatar_url {
+        let request = services().globals.default_client().get(avatar_url.as_ref());
+
+        let res = request.send().await.map_err(|_| {
+            Error::BadRequest(ErrorKind::UserInUse, "Desired UserId is already taken.")
+        })?;
+
+        let filename = avatar_url.path_segments().and_then(Iterator::last);
+
+        let (content_type, body): (Option<headers::ContentType>, Vec<u8>) = (
+            res.headers().typed_get(),
+            res.bytes().await.map(Into::into).map_err(|_| {
+                Error::BadRequest(ErrorKind::UserInUse, "Desired UserId is already taken.")
+            })?,
+        );
+
+        let mxc = format!(
+            "mxc://{}/{}",
+            services().globals.server_name(),
+            utils::random_string(crate::api::client_server::MXC_LENGTH)
+        );
+
+        services()
+            .media
+            .create(
+                mxc.clone(),
+                filename
+                    .map(|filename| "inline; filename=".to_owned() + filename)
+                    .as_deref(),
+                content_type.map(|header| header.to_string()).as_deref(),
+                &body,
+            )
+            .await?;
+
+        services()
+            .users
+            .set_avatar_url(&user_id, Some(OwnedMxcUri::from(mxc)))?;
+    };
+
+    if let (Some(displayname), true) = (
+        displayname.as_mut(),
+        services().globals.config.enable_lightning_bolt,
+    ) {
+        displayname.push_str(" ⚡️");
+    }
+
+    services().users.set_displayname(&user_id, displayname)?;
+
+    services().sso.save_claim(
+        &registration_token.provider_id,
+        &user_id,
+        &registration_token.unique_claim,
+    )?;
+
+    services().account_data.update(
+        None,
+        &user_id,
+        GlobalAccountDataEventType::PushRules.to_string().into(),
+        &serde_json::to_value(ruma::events::push_rules::PushRulesEvent {
+            content: ruma::events::push_rules::PushRulesEventContent {
+                global: push::Ruleset::server_default(&user_id),
+            },
+        })
+        .expect("PushRulesEvent should always serialize"),
+    )?;
+
+    let login_token = LoginToken::new(registration_token.provider_id, user_id);
+    let redirect_uri = redirect_with_login_token(registration_token.redirect_uri, &login_token);
+
+    Ok((
+        AppendHeaders([(
+            header::SET_COOKIE,
+            utils::reset_cookie("sso-registration").to_string(),
+        )]),
+        Redirect::temporary(redirect_uri.as_str()),
+    )
+        .into_response())
+}
+
+fn redirect_with_login_token(mut redirect_uri: Url, login_token: &LoginToken) -> Url {
+    let signed = services()
+        .globals
+        .sign_macaroon(login_token)
+        .expect("signing macaroons should always works");
+
+    redirect_uri
+        .query_pairs_mut()
+        .append_pair("loginToken", &signed);
+
+    redirect_uri
+}