Increased proxy token lifetime and implemented token reissuing

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "identity"
-version = "0.1.1"
+version = "0.1.5"
 edition = "2021"
 
 [dependencies]

src/auth.rs

@@ -1,6 +1,8 @@
 use crate::app::AppState;
 use crate::http::{get_header, AppError, Cookies, Either, SetCookie};
-use crate::proxy::{UpstreamAuthorizeParams, PROXY_AUTHORIZE_ENDPOINT, PROXY_TOKEN_TTL};
+use crate::proxy::{
+    ProxyClient, UpstreamAuthorizeParams, PROXY_AUTHORIZE_ENDPOINT, PROXY_TOKEN_TTL,
+};
 use crate::store::User;
 use crate::uri::UriBuilder;
 
@@ -49,6 +51,34 @@ pub struct AuthorizeParams {
     pub redirect_to: String,
 }
 
+pub async fn issue_upstream_token(
+    state: &AppState,
+    upstream: &ProxyClient,
+    user: &User,
+) -> anyhow::Result<Option<String>> {
+    let claims = state
+        .store()
+        .get_user_claims(user.id)
+        .await
+        .map_err(|err| anyhow::format_err!("failed to get user {} claims: {}", user.id, err))?;
+
+    if let Some(claims) = upstream.filter_claims(claims) {
+        let token = state
+            .issuer()
+            .create_token(
+                upstream.name(),
+                &user,
+                claims.as_ref(),
+                (*PROXY_TOKEN_TTL).into(),
+            )
+            .map_err(|err| anyhow::format_err!("failed to create token: {}", err))?;
+
+        Ok(Some(token))
+    } else {
+        Ok(None)
+    }
+}
+
 #[axum_macros::debug_handler]
 pub async fn authorize(
     State(state): State<AppState>,
@@ -74,29 +104,14 @@ pub async fn authorize(
         Some(user) => match state.upstreams().find_by_name(params.client_id.as_str()) {
             None => Err(StatusCode::BAD_REQUEST),
             Some(upstream) => {
-                let claims = state
-                    .store()
-                    .get_user_claims(user.id)
+                let token = issue_upstream_token(&state, upstream, &user)
                     .await
                     .map_err(|err| {
                         error!("failed to get user {} claims: {}", user.id, err);
                         StatusCode::INTERNAL_SERVER_ERROR
                     })?;
 
-                if let Some(claims) = upstream.filter_claims(claims) {
-                    let token = state
-                        .issuer()
-                        .create_token(
-                            upstream.name(),
-                            &user,
-                            claims.as_ref(),
-                            (*PROXY_TOKEN_TTL).into(),
-                        )
-                        .map_err(|err| {
-                            error!("failed to create token: {}", err);
-                            StatusCode::INTERNAL_SERVER_ERROR
-                        })?;
-
+                if let Some(token) = token {
                     let upstream_authorize_uri = UriBuilder::new()
                         .set_origin(upstream.origin())
                         .set_path(PROXY_AUTHORIZE_ENDPOINT)

src/issuer.rs

@@ -21,17 +21,23 @@ use url::Url;
 
 #[derive(Serialize, Deserialize)]
 pub struct Claims {
-    iss: String,
-    aud: String,
-    sub: String,
-    iat: i64,
-    exp: i64,
-    nbf: i64,
+    pub iss: String,
+    pub aud: String,
+    pub sub: String,
+    pub iat: i64,
+    pub exp: i64,
+    pub nbf: i64,
 
     // extras
-    name: String,
+    pub name: String,
     #[serde(flatten)]
-    extra: HashMap<String, String>,
+    pub extra: HashMap<String, String>,
+}
+
+impl Claims {
+    pub fn valid_for(&self) -> Duration {
+        Duration::seconds(self.exp - chrono::Utc::now().timestamp())
+    }
 }
 
 pub struct Issuer {

src/proxy.rs

@@ -1,13 +1,17 @@
 use crate::app::{AppState, UpstreamConfig};
-use crate::auth::AuthorizeParams;
+use crate::auth::{AuthorizeParams, issue_upstream_token};
 use crate::http::{Cookies, SetCookie};
 use crate::store::UserClaim;
 use crate::uri::UriBuilder;
-use crate::utils::Duration;
 
-use std::collections::{HashMap, HashSet};
-use std::str::FromStr;
+use std::{
+    borrow::Cow,
+    collections::{HashMap, HashSet},
+    convert::Into,
+    str::FromStr,
+};
 
+use crate::utils;
 use axum::{
     extract::{FromRequest, FromRequestParts, Host, OriginalUri, Query, State},
     http::{
@@ -17,12 +21,14 @@ use axum::{
     middleware::Next,
     response::{IntoResponse, Redirect, Response},
 };
+use chrono::Duration;
 use cookie::Cookie;
 use hyper::{client::HttpConnector, Body};
 use hyper_tls::HttpsConnector;
 use log::{error, info, warn};
 use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
+use crate::issuer::Claims;
 
 pub async fn middleware(
     State(state): State<AppState>,
@@ -48,8 +54,10 @@ pub async fn middleware(
     }
 }
 
+pub const PROXY_COOKIE_NAME: &str = "_identity.im";
 pub const PROXY_AUTHORIZE_ENDPOINT: &str = "/.identity/authorize";
-pub static PROXY_TOKEN_TTL: Lazy<Duration> = Lazy::new(|| Duration::minutes(5));
+pub static PROXY_TOKEN_TTL: Lazy<Duration> = Lazy::new(|| Duration::hours(1));
+pub static PROXY_TOKEN_REFRESH_THRESHOLD: Lazy<Duration> = Lazy::new(|| Duration::minutes(15));
 
 pub struct ProxyClient {
     config: UpstreamConfig,
@@ -60,8 +68,6 @@ pub struct ProxyClient {
     modified_headers: Option<Vec<(HeaderName, HeaderValue)>>,
 }
 
-const COOKIE_NAME: &str = "_identity.im";
-
 #[derive(Serialize, Deserialize)]
 pub struct UpstreamAuthorizeParams {
     pub token: String,
@@ -167,6 +173,18 @@ impl ProxyClient {
         }
     }
 
+    async fn refresh_token_if_needed(&self, state: &AppState, claims: &Claims) -> anyhow::Result<Option<String>> {
+        let refresh_needed = claims.valid_for() < *PROXY_TOKEN_REFRESH_THRESHOLD;
+        if refresh_needed {
+            let user_id = i32::from_str(claims.sub.as_str())?;
+            let user = state.store().find_user_by_id(user_id).await?.ok_or(anyhow::format_err!("user by id {} not found", claims.sub))?;
+
+            issue_upstream_token(state, self, &user).await
+        } else {
+            Ok(None)
+        }
+    }
+
     pub async fn handle(
         &self,
         mut request: Request<Body>,
@@ -176,15 +194,23 @@ impl ProxyClient {
             return Ok(self.authorize(request, state).await?.into_response());
         }
 
-        if let Some(cookie) = Cookies::extract_one(request.headers(), COOKIE_NAME) {
-            if let Err(err) = state.issuer().validate_token(cookie.value()) {
-                warn!("token validation failed: {}", err);
-            } else {
-                let value: HeaderValue = format!("Bearer {}", cookie.value()).try_into()?;
+        if let Some(cookie) = Cookies::extract_one(request.headers(), PROXY_COOKIE_NAME) {
+            match state.issuer().validate_token(cookie.value()) {
+                Err(err) => {
+                    warn!("token validation failed: {}", err);
+                }
+                Ok(claims) => {
+                    let value: HeaderValue = format!("Bearer {}", cookie.value()).try_into()?;
+                    request.headers_mut().append(AUTHORIZATION, value);
 
-                request.headers_mut().append(AUTHORIZATION, value);
 
-                return self.forward(request).await;
+                    return if let Ok(Some(token)) = self.refresh_token_if_needed(state, &claims).await {
+                        let response= self.forward(request).await?;
+                        Ok((Self::set_cookie(token), response).into_response())
+                    } else {
+                        self.forward(request).await
+                    }
+                }
             }
         }
 
@@ -195,6 +221,17 @@ impl ProxyClient {
         }
     }
 
+    fn set_cookie<'c, T>(token: T) -> SetCookie<'c>
+    where
+        T: Into<Cow<'c, str>>,
+    {
+        SetCookie(Cookie::build(PROXY_COOKIE_NAME, token)
+            .path("/")
+            .http_only(true)
+            .max_age(utils::Duration::from(*PROXY_TOKEN_TTL).into())
+            .finish())
+    }
+
     async fn authorize(
         &self,
         request: Request<Body>,
@@ -203,13 +240,7 @@ impl ProxyClient {
         let query = Query::<UpstreamAuthorizeParams>::from_request(request, state).await?;
 
         Ok((
-            SetCookie(
-                Cookie::build(COOKIE_NAME, query.token.clone())
-                    .path("/")
-                    .http_only(true)
-                    .max_age((*PROXY_TOKEN_TTL).into())
-                    .finish(),
-            ),
+            Self::set_cookie(query.token.clone()),
             Redirect::to(query.redirect_to.as_str()),
         ))
     }

src/sql/find_user_by_id.sql

@@ -0,0 +1 @@
+SELECT id, username, password_hash, role_name FROM users WHERE id = ?

src/store.rs

@@ -144,6 +144,28 @@ impl UserStore {
         Ok(row.is_some())
     }
 
+    pub async fn find_user_by_id(&self, id: i32) -> anyhow::Result<Option<User>> {
+        #[derive(sqlx::FromRow)]
+        struct Row {
+            id: i32,
+            username: String,
+            password_hash: String,
+            role_name: String,
+        }
+
+        let row = sqlx::query_as::<_, Row>(include_str!("sql/find_user_by_id.sql"))
+            .bind(id)
+            .fetch_optional(self.get_pool())
+            .await?;
+
+        Ok(row.map(|row| User {
+            id: row.id,
+            username: row.username,
+            password_hash: row.password_hash,
+            role: UserRole::from_str(row.role_name.as_str()).unwrap_or(UserRole::User),
+        }))
+    }
+
     pub async fn find_user_by_username(&self, username: &str) -> anyhow::Result<Option<User>> {
         let (_, normalized_username) = Self::normalize_username(username);
 

src/utils.rs

@@ -1,11 +1,5 @@
 #[derive(Copy, Clone)]
-pub struct Duration(pub chrono::Duration);
-
-impl Duration {
-    pub fn minutes(value: i64) -> Duration {
-        chrono::Duration::minutes(value).into()
-    }
-}
+pub struct Duration(chrono::Duration);
 
 impl From<chrono::Duration> for Duration {
     fn from(value: chrono::Duration) -> Self {
@@ -21,14 +15,14 @@ impl From<cookie::time::Duration> for Duration {
     }
 }
 
-impl Into<chrono::Duration> for Duration {
-    fn into(self) -> chrono::Duration {
-        self.0
+impl From<Duration> for chrono::Duration {
+    fn from(value: Duration) -> Self {
+        value.0
     }
 }
 
-impl Into<cookie::time::Duration> for Duration {
-    fn into(self) -> cookie::time::Duration {
-        cookie::time::Duration::milliseconds(self.0.num_milliseconds())
+impl From<Duration> for cookie::time::Duration {
+    fn from(value: Duration) -> Self {
+        cookie::time::Duration::milliseconds(value.0.num_milliseconds())
     }
 }