matrix/authentication-service
1
0
Fork 0
mirror of https://github.com/matrix-org/matrix-authentication-service.git synced 2025-07-29 22:01:14 +03:00
Code Activity

Save the post auth action during upstream OAuth login

Browse Source
This commit is contained in:
Quentin Gliech
2022-12-05 18:27:56 +01:00
parent 4d93f4d4f0
commit 23fd833d45
18 changed files with 142 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/axum-utils/Cargo.toml
View File

@ -9,7 +9,6 @@ license = "Apache-2.0"
async-trait = "0.1.59"
axum = { version = "0.6.1", features = ["headers"] }
axum-extra = { version = "0.4.2", features = ["cookie-private"] }
bincode = "1.3.3"
chrono = "0.4.23"
data-encoding = "2.3.2"
futures-util = "0.3.25"

12
crates/axum-utils/src/cookies.rs
View File

@ -14,15 +14,13 @@
//! Private (encrypted) cookie jar, based on axum-extra's cookie jar
use data_encoding::BASE64URL_NOPAD;
use serde::{de::DeserializeOwned, Serialize};
use thiserror::Error;
#[derive(Debug, Error)]
#[error("could not decode cookie")]
pub enum CookieDecodeError {
Deserialize(#[from] bincode::Error),
Decode(#[from] data_encoding::DecodeError),
Deserialize(#[from] serde_json::Error),
}
pub trait CookieExt {
@ -41,10 +39,7 @@ impl<'a> CookieExt for axum_extra::extract::cookie::Cookie<'a> {
where
T: DeserializeOwned,
{
let bytes = BASE64URL_NOPAD.decode(self.value().as_bytes())?;
let decoded = bincode::deserialize(&bytes)?;
let decoded = serde_json::from_str(self.value())?;
Ok(decoded)
}
@ -52,8 +47,7 @@ impl<'a> CookieExt for axum_extra::extract::cookie::Cookie<'a> {
where
T: Serialize,
{
let bytes = bincode::serialize(t).unwrap();
let encoded = BASE64URL_NOPAD.encode(&bytes);
let encoded = serde_json::to_string(t).unwrap();
self.set_value(encoded);
self
}

7
crates/handlers/src/upstream_oauth2/authorize.rs
View File

@ -13,7 +13,7 @@
// limitations under the License.
use axum::{
extract::{Path, State},
extract::{Path, Query, State},
response::{IntoResponse, Redirect},
};
use axum_extra::extract::PrivateCookieJar;
@ -28,7 +28,7 @@ use thiserror::Error;
use ulid::Ulid;
use super::UpstreamSessionsCookie;
use crate::impl_from_error_for_route;
use crate::{impl_from_error_for_route, views::shared::OptionalPostAuthAction};
#[derive(Debug, Error)]
pub(crate) enum RouteError {
@ -68,6 +68,7 @@ pub(crate) async fn get(
State(url_builder): State<UrlBuilder>,
cookie_jar: PrivateCookieJar<Encrypter>,
Path(provider_id): Path<Ulid>,
Query(query): Query<OptionalPostAuthAction>,
) -> Result<impl IntoResponse, RouteError> {
let (clock, mut rng) = crate::rng_and_clock()?;
@ -115,7 +116,7 @@ pub(crate) async fn get(
.await?;
let cookie_jar = UpstreamSessionsCookie::load(&cookie_jar)
.add(session.id, provider.id, data.state)
.add(session.id, provider.id, data.state, query.post_auth_action)
.save(cookie_jar, clock.now());
txn.commit().await?;

2
crates/handlers/src/upstream_oauth2/callback.rs
View File

@ -137,7 +137,7 @@ pub(crate) async fn get(
let mut txn = pool.begin().await?;
let sessions_cookie = UpstreamSessionsCookie::load(&cookie_jar);
let session_id = sessions_cookie
let (session_id, _post_auth_action) = sessions_cookie
.find_session(provider_id, &params.state)
.map_err(|_| RouteError::MissingCookie)?;

44
crates/handlers/src/upstream_oauth2/cookie.rs
View File

@ -17,6 +17,7 @@
use axum_extra::extract::{cookie::Cookie, PrivateCookieJar};
use chrono::{DateTime, Duration, NaiveDateTime, Utc};
use mas_axum_utils::CookieExt;
use mas_router::PostAuthAction;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use time::OffsetDateTime;
@ -28,12 +29,13 @@ static COOKIE_NAME: &str = "upstream-oauth2-sessions";
/// Sessions expire after 10 minutes
static SESSION_MAX_TIME_SECS: i64 = 60 * 10;
#[derive(Serialize, Deserialize)]
#[derive(Serialize, Deserialize, Debug)]
pub struct Payload {
session: Ulid,
provider: Ulid,
state: String,
link: Option<Ulid>,
post_auth_action: Option<PostAuthAction>,
}
impl Payload {
@ -46,7 +48,7 @@ impl Payload {
}
}
#[derive(Serialize, Deserialize, Default)]
#[derive(Serialize, Deserialize, Default, Debug)]
pub struct UpstreamSessions(Vec<Payload>);
#[derive(Debug, Error, PartialEq, Eq)]
@ -87,12 +89,19 @@ impl UpstreamSessions {
}
/// Add a new session, for a provider and a random state
pub fn add(mut self, session: Ulid, provider: Ulid, state: String) -> Self {
pub fn add(
mut self,
session: Ulid,
provider: Ulid,
state: String,
post_auth_action: Option<PostAuthAction>,
) -> Self {
self.0.push(Payload {
session,
provider,
state,
link: None,
post_auth_action,
});
self
}
@ -102,11 +111,11 @@ impl UpstreamSessions {
&self,
provider: Ulid,
state: &str,
) -> Result<Ulid, UpstreamSessionNotFound> {
) -> Result<(Ulid, Option<&PostAuthAction>), UpstreamSessionNotFound> {
self.0
.iter()
.find(|p| p.provider == provider && p.state == state && p.link.is_none())
.map(|p| p.session)
.map(|p| (p.session, p.post_auth_action.as_ref()))
.ok_or(UpstreamSessionNotFound)
}
@ -127,11 +136,14 @@ impl UpstreamSessions {
}
/// Find a session from its link
pub fn lookup_link(&self, link_id: Ulid) -> Result<Ulid, UpstreamSessionNotFound> {
pub fn lookup_link(
&self,
link_id: Ulid,
) -> Result<(Ulid, Option<&PostAuthAction>), UpstreamSessionNotFound> {
self.0
.iter()
.find(|p| p.link == Some(link_id))
.map(|p| p.session)
.map(|p| (p.session, p.post_auth_action.as_ref()))
.ok_or(UpstreamSessionNotFound)
}
@ -171,22 +183,22 @@ mod tests {
let first_session = Ulid::from_datetime_with_source(now.into(), &mut rng);
let first_state = "first-state";
let sessions = sessions.add(first_session, provider_a, first_state.into());
let sessions = sessions.add(first_session, provider_a, first_state.into(), None);
let now = now + Duration::minutes(5);
let second_session = Ulid::from_datetime_with_source(now.into(), &mut rng);
let second_state = "second-state";
let sessions = sessions.add(second_session, provider_b, second_state.into());
let sessions = sessions.add(second_session, provider_b, second_state.into(), None);
let sessions = sessions.expire(now);
assert_eq!(
sessions.find_session(provider_a, first_state),
Ok(first_session)
sessions.find_session(provider_a, first_state).unwrap().0,
first_session,
);
assert_eq!(
sessions.find_session(provider_b, second_state),
Ok(second_session)
sessions.find_session(provider_b, second_state).unwrap().0,
second_session
);
assert!(sessions.find_session(provider_b, first_state).is_err());
assert!(sessions.find_session(provider_a, second_state).is_err());
@ -196,8 +208,8 @@ mod tests {
let sessions = sessions.expire(now);
assert!(sessions.find_session(provider_a, first_state).is_err());
assert_eq!(
sessions.find_session(provider_b, second_state),
Ok(second_session)
sessions.find_session(provider_b, second_state).unwrap().0,
second_session
);
// Associate a link with the second
@ -210,7 +222,7 @@ mod tests {
assert!(sessions.find_session(provider_b, second_state).is_err());
// But it can be looked up by its link
assert_eq!(sessions.lookup_link(second_link), Ok(second_session));
assert_eq!(sessions.lookup_link(second_link).unwrap().0, second_session);
// And it can be consumed
let sessions = sessions.consume_link(second_link).unwrap();
// But only once

13
crates/handlers/src/upstream_oauth2/link.rs
View File

@ -24,7 +24,6 @@ use mas_axum_utils::{
SessionInfoExt,
};
use mas_keystore::Encrypter;
use mas_router::Route;
use mas_storage::{
upstream_oauth2::{
associate_link_to_user, consume_session, lookup_link, lookup_session_on_link,
@ -44,7 +43,7 @@ use thiserror::Error;
use ulid::Ulid;
use super::UpstreamSessionsCookie;
use crate::impl_from_error_for_route;
use crate::{impl_from_error_for_route, views::shared::OptionalPostAuthAction};
#[derive(Debug, Error)]
pub(crate) enum RouteError {
@ -114,7 +113,7 @@ pub(crate) async fn get(
let (clock, mut rng) = crate::rng_and_clock()?;
let sessions_cookie = UpstreamSessionsCookie::load(&cookie_jar);
let session_id = sessions_cookie
let (session_id, _post_auth_action) = sessions_cookie
.lookup_link(link_id)
.map_err(|_| RouteError::MissingCookie)?;
@ -213,10 +212,14 @@ pub(crate) async fn post(
let form = cookie_jar.verify_form(clock.now(), form)?;
let sessions_cookie = UpstreamSessionsCookie::load(&cookie_jar);
let session_id = sessions_cookie
let (session_id, post_auth_action) = sessions_cookie
.lookup_link(link_id)
.map_err(|_| RouteError::MissingCookie)?;
let post_auth_action = OptionalPostAuthAction {
post_auth_action: post_auth_action.cloned(),
};
let link = lookup_link(&mut txn, link_id)
.await
.to_option()?
@ -267,5 +270,5 @@ pub(crate) async fn post(
txn.commit().await?;
Ok((cookie_jar, mas_router::Index.go()))
Ok((cookie_jar, post_auth_action.go_next()))
}

6
crates/handlers/src/views/login.rs
View File

@ -22,7 +22,6 @@ use mas_axum_utils::{
FancyError, SessionInfoExt,
};
use mas_keystore::Encrypter;
use mas_router::Route;
use mas_storage::user::{login, LoginError};
use mas_templates::{
FieldError, FormError, LoginContext, LoginFormField, TemplateContext, Templates, ToFormState,
@ -160,10 +159,7 @@ async fn render(
} else {
ctx
};
let register_link = mas_router::Register::from(action.post_auth_action).relative_url();
let ctx = ctx
.with_register_link(register_link.to_string())
.with_csrf(csrf_token.form_value());
let ctx = ctx.with_csrf(csrf_token.form_value());
let content = templates.render_login(&ctx).await?;
Ok(content)

5
crates/handlers/src/views/register.rs
View File

@ -243,10 +243,7 @@ async fn render(
} else {
ctx
};
let login_link = mas_router::Login::from(action.post_auth_action).relative_url();
let ctx = ctx
.with_login_link(login_link.to_string())
.with_csrf(csrf_token.form_value());
let ctx = ctx.with_csrf(csrf_token.form_value());
let content = templates.render_register(&ctx).await?;
Ok(content)

32
crates/handlers/src/views/shared.rs
View File

@ -16,7 +16,7 @@ use mas_router::{PostAuthAction, Route};
use mas_storage::{
compat::get_compat_sso_login_by_id, oauth2::authorization_grant::get_grant_by_id,
};
use mas_templates::PostAuthContext;
use mas_templates::{PostAuthContext, PostAuthContextInner};
use serde::{Deserialize, Serialize};
use sqlx::PgConnection;
@ -41,23 +41,24 @@ impl OptionalPostAuthAction {
&self,
conn: &mut PgConnection,
) -> anyhow::Result<Option<PostAuthContext>> {
match &self.post_auth_action {
Some(PostAuthAction::ContinueAuthorizationGrant { data }) => {
let grant = get_grant_by_id(conn, *data).await?;
let Some(action) = self.post_auth_action.clone() else { return Ok(None) };
let ctx = match action {
PostAuthAction::ContinueAuthorizationGrant { data } => {
let grant = get_grant_by_id(conn, data).await?;
let grant = Box::new(grant.into());
Ok(Some(PostAuthContext::ContinueAuthorizationGrant { grant }))
PostAuthContextInner::ContinueAuthorizationGrant { grant }
}
Some(PostAuthAction::ContinueCompatSsoLogin { data }) => {
let login = get_compat_sso_login_by_id(conn, *data).await?;
PostAuthAction::ContinueCompatSsoLogin { data } => {
let login = get_compat_sso_login_by_id(conn, data).await?;
let login = Box::new(login.into());
Ok(Some(PostAuthContext::ContinueCompatSsoLogin { login }))
PostAuthContextInner::ContinueCompatSsoLogin { login }
}
Some(PostAuthAction::ChangePassword) => Ok(Some(PostAuthContext::ChangePassword)),
PostAuthAction::ChangePassword => PostAuthContextInner::ChangePassword,
Some(PostAuthAction::LinkUpstream { id }) => {
let link = mas_storage::upstream_oauth2::lookup_link(&mut *conn, *id).await?;
PostAuthAction::LinkUpstream { id } => {
let link = mas_storage::upstream_oauth2::lookup_link(&mut *conn, id).await?;
let provider =
mas_storage::upstream_oauth2::lookup_provider(&mut *conn, link.provider_id)
@ -65,10 +66,13 @@ impl OptionalPostAuthAction {
let provider = Box::new(provider);
let link = Box::new(link);
Ok(Some(PostAuthContext::LinkUpstream { provider, link }))
PostAuthContextInner::LinkUpstream { provider, link }
}
};
None => Ok(None),
}
Ok(Some(PostAuthContext {
params: action.clone(),
ctx,
}))
}
}

18
crates/router/src/endpoints.rs
View File

@ -532,17 +532,27 @@ impl Route for CompatLoginSsoComplete {
/// `GET /upstream/authorize/:id`
pub struct UpstreamOAuth2Authorize {
id: Ulid,
post_auth_action: Option<PostAuthAction>,
}
impl UpstreamOAuth2Authorize {
#[must_use]
pub const fn new(id: Ulid) -> Self {
Self { id }
Self {
id,
post_auth_action: None,
}
}
#[must_use]
pub fn and_then(mut self, action: PostAuthAction) -> Self {
self.post_auth_action = Some(action);
self
}
}
impl Route for UpstreamOAuth2Authorize {
type Query = ();
type Query = PostAuthAction;
fn route() -> &'static str {
"/upstream/authorize/:provider_id"
}
@ -550,6 +560,10 @@ impl Route for UpstreamOAuth2Authorize {
fn path(&self) -> std::borrow::Cow<'static, str> {
format!("/upstream/authorize/{}", self.id).into()
}
fn query(&self) -> Option<&Self::Query> {
self.post_auth_action.as_ref()
}
}
/// `GET /upstream/callback/:id`

1
crates/templates/Cargo.toml
View File

@ -20,6 +20,7 @@ serde_urlencoded = "0.7.1"
camino = "1.1.1"
chrono = "0.4.23"
url = "2.3.1"
http = "0.2.8"
ulid = { version = "1.0.0", features = ["serde"] }
oauth2-types = { path = "../oauth2-types" }

38
crates/templates/src/context.rs
View File

@ -244,10 +244,10 @@ impl FormField for LoginFormField {
}
}
/// Context used in login and reauth screens, for the post-auth action to do
/// Inner context used in login and reauth screens. See [`PostAuthContext`].
#[derive(Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum PostAuthContext {
pub enum PostAuthContextInner {
/// Continue an authorization grant
ContinueAuthorizationGrant {
/// The authorization grant that will be continued after authentication
@ -274,13 +274,23 @@ pub enum PostAuthContext {
},
}
/// Context used in login and reauth screens, for the post-auth action to do
#[derive(Serialize)]
pub struct PostAuthContext {
/// The post auth action params from the URL
pub params: PostAuthAction,
/// The loaded post auth context
#[serde(flatten)]
pub ctx: PostAuthContextInner,
}
/// Context used by the `login.html` template
#[derive(Serialize, Default)]
pub struct LoginContext {
form: FormState<LoginFormField>,
next: Option<PostAuthContext>,
providers: Vec<UpstreamOAuthProvider>,
register_link: String,
}
impl TemplateContext for LoginContext {
@ -293,7 +303,6 @@ impl TemplateContext for LoginContext {
form: FormState::default(),
next: None,
providers: Vec::new(),
register_link: "/register".to_owned(),
}]
}
}
@ -313,18 +322,9 @@ impl LoginContext {
/// Add a post authentication action to the context
#[must_use]
pub fn with_post_action(self, next: PostAuthContext) -> Self {
pub fn with_post_action(self, context: PostAuthContext) -> Self {
Self {
next: Some(next),
..self
}
}
/// Add a registration link to the context
#[must_use]
pub fn with_register_link(self, register_link: String) -> Self {
Self {
register_link,
next: Some(context),
..self
}
}
@ -361,7 +361,6 @@ impl FormField for RegisterFormField {
pub struct RegisterContext {
form: FormState<RegisterFormField>,
next: Option<PostAuthContext>,
login_link: String,
}
impl TemplateContext for RegisterContext {
@ -373,7 +372,6 @@ impl TemplateContext for RegisterContext {
vec![RegisterContext {
form: FormState::default(),
next: None,
login_link: "/login".to_owned(),
}]
}
}
@ -393,12 +391,6 @@ impl RegisterContext {
..self
}
}
/// Add a login link to the context
#[must_use]
pub fn with_login_link(self, login_link: String) -> Self {
Self { login_link, ..self }
}
}
/// Context used by the `consent.html` template

43
crates/templates/src/functions.rs
View File

@ -22,7 +22,9 @@ use url::Url;
pub fn register(tera: &mut Tera, url_builder: UrlBuilder) {
tera.register_tester("empty", self::tester_empty);
tera.register_function("add_params_to_uri", function_add_params_to_uri);
tera.register_filter("to_params", filter_to_params);
tera.register_filter("safe_get", filter_safe_get);
tera.register_function("add_params_to_url", function_add_params_to_url);
tera.register_function("merge", function_merge);
tera.register_function("dict", function_dict);
tera.register_function("static_asset", make_static_asset(url_builder));
@ -37,12 +39,42 @@ fn tester_empty(value: Option<&Value>, params: &[Value]) -> Result<bool, tera::E
}
}
fn filter_to_params(params: &Value, kv: &HashMap<String, Value>) -> Result<Value, tera::Error> {
let prefix = kv.get("prefix").and_then(Value::as_str).unwrap_or("");
let params = serde_urlencoded::to_string(params)
.map_err(|e| tera::Error::chain(e, "Could not serialize parameters"))?;
if params.is_empty() {
Ok(Value::String(String::new()))
} else {
Ok(Value::String(format!("{prefix}{params}")))
}
}
/// Alternative to `get` which does not crash on `None` and defaults to `None`
pub fn filter_safe_get(value: &Value, args: &HashMap<String, Value>) -> Result<Value, tera::Error> {
let default = args.get("default").unwrap_or(&Value::Null);
let key = args
.get("key")
.and_then(Value::as_str)
.ok_or_else(|| tera::Error::msg("Invalid parameter `uri`"))?;
match value.as_object() {
Some(o) => match o.get(key) {
Some(val) => Ok(val.clone()),
// If the value is not present, allow for an optional default value
None => Ok(default.clone()),
},
None => Ok(default.clone()),
}
}
enum ParamsWhere {
Fragment,
Query,
}
fn function_add_params_to_uri(params: &HashMap<String, Value>) -> Result<Value, tera::Error> {
fn function_add_params_to_url(params: &HashMap<String, Value>) -> Result<Value, tera::Error> {
use ParamsWhere::{Fragment, Query};
// First, get the `uri`, `mode` and `params` parameters
@ -77,12 +109,7 @@ fn function_add_params_to_uri(params: &HashMap<String, Value>) -> Result<Value,
.unwrap_or_default();
// Merge the exising and the additional parameters together
let params: HashMap<&String, &Value> = params
.iter()
// Filter out the `uri` and `mode` params
.filter(|(k, _v)| k != &"uri" && k != &"mode")
.chain(existing.iter())
.collect();
let params: HashMap<&String, &Value> = params.iter().chain(existing.iter()).collect();
// Transform them back to urlencoded
let params = serde_urlencoded::to_string(params)

6
crates/templates/src/lib.rs
View File

@ -48,9 +48,9 @@ pub use self::{
AccountContext, AccountEmailsContext, CompatSsoContext, ConsentContext, EmailAddContext,
EmailVerificationContext, EmailVerificationPageContext, EmptyContext, ErrorContext,
FormPostContext, IndexContext, LoginContext, LoginFormField, PolicyViolationContext,
PostAuthContext, ReauthContext, ReauthFormField, RegisterContext, RegisterFormField,
TemplateContext, UpstreamExistingLinkContext, UpstreamRegister, UpstreamSuggestLink,
WithCsrf, WithOptionalSession, WithSession,
PostAuthContext, PostAuthContextInner, ReauthContext, ReauthFormField, RegisterContext,
RegisterFormField, TemplateContext, UpstreamExistingLinkContext, UpstreamRegister,
UpstreamSuggestLink, WithCsrf, WithOptionalSession, WithSession,
},
forms::{FieldError, FormError, FormField, FormState, ToFormState},
};