audience added to token (#29)

* audience added to token

Author: YaroslavLitvinov

crates/api-snowflake-rest-sessions/src/error.rs

@@ -10,7 +10,7 @@ use snafu::prelude::*;
 pub type Result<T> = std::result::Result<T, Error>;
 
 #[derive(Snafu)]
-#[snafu(visibility(pub(crate)))]
+#[snafu(visibility(pub))]
 #[error_stack_trace::debug]
 pub enum Error {
     #[snafu(display("Can't add header to response: {error}"))]
@@ -42,6 +42,20 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
+    #[snafu(display("Can't authenticate request: Host is missing"))]
+    MissingHost {
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display("Extension error: {error}"))]
+    ExtensionRejection {
+        #[snafu(source)]
+        error: axum::extract::rejection::ExtensionRejection,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -53,7 +67,12 @@ pub struct ErrorResponse {
 impl IntoResponse for Error {
     fn into_response(self) -> axum::response::Response<axum::body::Body> {
         let message = self.to_string();
-        let code = StatusCode::INTERNAL_SERVER_ERROR;
+        let code = match self {
+            Self::BadAuthToken { .. }
+            | Self::MissingHost { .. }
+            | Self::ExtensionRejection { .. } => StatusCode::UNAUTHORIZED,
+            _ => StatusCode::INTERNAL_SERVER_ERROR,
+        };
 
         let error = ErrorResponse {
             message,

crates/api-snowflake-rest-sessions/src/helpers.rs

@@ -8,19 +8,21 @@ use uuid::Uuid;
 #[cfg_attr(test, derive(Debug))]
 pub struct Claims {
     pub sub: String, // token issued to a particular user
+    pub aud: String, // validate audience since as it can be deployed on multiple hosts
     pub iat: i64,    // Issued At
     pub exp: i64,    // Expiration Time
     pub session_id: String,
 }
 
 #[must_use]
-pub fn jwt_claims(username: &str, expiration: Duration) -> Claims {
+pub fn jwt_claims(username: &str, audience: &str, expiration: Duration) -> Claims {
     let now = Local::now();
     let iat = now.timestamp();
     let exp = now.timestamp() + expiration.whole_seconds();
 
     Claims {
         sub: username.to_string(),
+        aud: audience.to_string(),
         iat,
         exp,
         session_id: Uuid::new_v4().to_string(),
@@ -29,11 +31,13 @@ pub fn jwt_claims(username: &str, expiration: Duration) -> Claims {
 
 pub fn get_claims_validate_jwt_token(
     token: &str,
+    audience: &str,
     jwt_secret: &str,
 ) -> Result<Claims, jsonwebtoken::errors::Error> {
     let mut validation = Validation::default();
     validation.leeway = 5;
-    validation.set_required_spec_claims(&["exp"]);
+    validation.set_audience(&[audience]);
+    validation.set_required_spec_claims(&["exp", "aud"]);
 
     let decoding_key = DecodingKey::from_secret(jwt_secret.as_bytes());
 

crates/api-snowflake-rest-sessions/src/layer.rs

@@ -1,16 +1,44 @@
 use crate::error as session_error;
-use crate::error::Result;
+use crate::error::{Error, Result};
 use crate::session::{
     DFSessionId, SESSION_ID_COOKIE_NAME, SessionStore, extract_token_from_cookie,
 };
-use axum::extract::{Request, State};
+use axum::extract::{FromRequestParts, Request, State};
+use axum::http::{HeaderMap, HeaderName, request::Parts};
 use axum::middleware::Next;
 use axum::response::IntoResponse;
 use http::header::SET_COOKIE;
-use http::{HeaderMap, HeaderName};
 use snafu::ResultExt;
 use tower_sessions::cookie::{Cookie, SameSite};
 
+#[derive(Debug, Clone)]
+pub struct Host(pub String);
+
+impl<S> FromRequestParts<S> for Host
+where
+    S: Send + Sync,
+{
+    type Rejection = Error;
+
+    #[allow(clippy::unwrap_used)]
+    async fn from_request_parts(
+        req: &mut Parts,
+        state: &S,
+    ) -> std::result::Result<Self, Self::Rejection> {
+        let headers = HeaderMap::from_request_parts(req, state)
+            .await
+            .map_err(|err| match err {})
+            .unwrap(); // unwrap on Infallibe error is safe
+        let host = headers.get("host");
+        let host = host.and_then(|host| host.to_str().ok());
+        if let Some(host) = host {
+            Ok(Self(host.to_string()))
+        } else {
+            session_error::MissingHostSnafu.fail()
+        }
+    }
+}
+
 #[allow(clippy::unwrap_used, clippy::cognitive_complexity)]
 pub async fn propagate_session_cookie(
     State(state): State<SessionStore>,

crates/api-snowflake-rest-sessions/src/session.rs

@@ -8,7 +8,7 @@ use http::header::COOKIE;
 use http::request::Parts;
 use http::{HeaderMap, HeaderName};
 use regex::Regex;
-use snafu::ResultExt;
+use snafu::{OptionExt, ResultExt};
 use std::{collections::HashMap, sync::Arc};
 
 pub const SESSION_ID_COOKIE_NAME: &str = "session_id";
@@ -52,10 +52,24 @@ where
     async fn from_request_parts(req: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
         let execution_svc = state.get_execution_svc();
 
+        // Not using Host extractor as for some reason it extracts host without port
+        // use axum::RequestPartsExt;
+        // use axum::extract::Extension;
+        // use crate::layer::Host;
+        // let Extension(Host(host)) = req.extract::<Extension<Host>>()
+        //     .await
+        //     .context(session_error::ExtensionRejectionSnafu)?;
+        // tracing::info!("Host '{host}' extracted from DFSessionId");
+
         let (session_id, located_at) = if let Some(token) = extract_token_from_auth(&req.headers) {
+            // host is require to check token audience claim
+            let host = req.headers.get("host");
+            let host = host.and_then(|host| host.to_str().ok());
+            let host = host.context(session_error::MissingHostSnafu)?;
+
             let jwt_secret = state.jwt_secret();
-            let jwt_claims =
-                get_claims_validate_jwt_token(&token, jwt_secret).context(BadAuthTokenSnafu)?;
+            let jwt_claims = get_claims_validate_jwt_token(&token, host, jwt_secret)
+                .context(BadAuthTokenSnafu)?;
 
             (jwt_claims.session_id, "auth header")
         } else {

crates/api-snowflake-rest/src/server/error.rs

@@ -93,6 +93,14 @@ pub enum Error {
         #[snafu(implicit)]
         location: Location,
     },
+
+    #[snafu(display("Auth session error: {source}"))]
+    AuthSession {
+        #[snafu(source(from(api_snowflake_rest_sessions::error::Error, Box::new)))]
+        source: Box<api_snowflake_rest_sessions::error::Error>,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 impl IntoResponse for Error {
@@ -216,7 +224,8 @@ impl Error {
             | Self::InvalidAuthToken { .. }
             | Self::NoJwtSecret { .. }
             | Self::CreateJwt { .. }
-            | Self::BadAuthToken { .. } => (
+            | Self::BadAuthToken { .. }
+            | Self::AuthSession { .. } => (
                 http::StatusCode::UNAUTHORIZED,
                 SqlState::Success,
                 ErrorCode::Other,

crates/api-snowflake-rest/src/server/handlers.rs

@@ -5,6 +5,7 @@ use crate::models::{
 use crate::server::error::Result;
 use crate::server::logic::{handle_login_request, handle_query_request};
 use api_snowflake_rest_sessions::DFSessionId;
+use api_snowflake_rest_sessions::layer::Host;
 use axum::Json;
 use axum::extract::{ConnectInfo, Query, State};
 use executor::RunningQueryId;
@@ -23,10 +24,11 @@ pub struct SessionQueryParams {
 
 #[tracing::instrument(name = "api_snowflake_rest::login", level = "debug", skip(state), err, ret(level = tracing::Level::TRACE))]
 pub async fn login(
+    Host(host): Host,
     State(state): State<AppState>,
     Json(login_request): Json<LoginRequestBody>,
 ) -> Result<Json<LoginResponse>> {
-    let response = handle_login_request(&state, login_request.data).await?;
+    let response = handle_login_request(&state, host, login_request.data).await?;
     Ok(Json(response))
 }
 

crates/api-snowflake-rest/src/server/layer.rs

@@ -3,6 +3,7 @@ use crate::server::error::{BadAuthTokenSnafu, NoJwtSecretSnafu};
 use api_snowflake_rest_sessions::helpers::{
     ensure_jwt_secret_is_valid, get_claims_validate_jwt_token,
 };
+use api_snowflake_rest_sessions::layer::Host;
 use api_snowflake_rest_sessions::session::extract_token_from_auth;
 use axum::extract::{Request, State};
 use axum::middleware::Next;
@@ -19,6 +20,7 @@ use snafu::{OptionExt, ResultExt};
 )]
 pub async fn require_auth(
     State(state): State<AppState>,
+    Host(host): Host,
     req: Request,
     next: Next,
 ) -> error::Result<impl IntoResponse> {
@@ -35,7 +37,7 @@ pub async fn require_auth(
         ensure_jwt_secret_is_valid(&state.config.auth.jwt_secret).context(NoJwtSecretSnafu)?;
 
     let jwt_claims =
-        get_claims_validate_jwt_token(&token, &jwt_secret).context(BadAuthTokenSnafu)?;
+        get_claims_validate_jwt_token(&token, &host, &jwt_secret).context(BadAuthTokenSnafu)?;
 
     // Record the result as part of the current span.
     tracing::Span::current().record("session_id", jwt_claims.session_id.as_str());

crates/api-snowflake-rest/src/server/logic.rs

@@ -24,6 +24,7 @@ pub const JWT_TOKEN_EXPIRATION_SECONDS: u32 = 24 * 60 * 60;
 )]
 pub async fn handle_login_request(
     state: &AppState,
+    host: String,
     credentials: LoginRequestData,
 ) -> Result<LoginResponse> {
     let LoginRequestData {
@@ -36,14 +37,18 @@ pub async fn handle_login_request(
         return api_snowflake_rest_error::InvalidAuthDataSnafu.fail();
     }
 
+    // host is required to check token audience claim
     let jwt_secret = &*state.config.auth.jwt_secret;
     let _ = ensure_jwt_secret_is_valid(jwt_secret).context(NoJwtSecretSnafu)?;
 
     let jwt_claims = jwt_claims(
         &login_name,
+        &host,
         Duration::seconds(JWT_TOKEN_EXPIRATION_SECONDS.into()),
     );
 
+    tracing::info!("Host '{host}' for token creation");
+
     let session_id = jwt_claims.session_id.clone();
     let _ = state.execution_svc.create_session(&session_id).await?;
 

crates/embucket-lambda/src/main.rs

@@ -7,7 +7,9 @@ use api_snowflake_rest::server::layer::require_auth;
 use api_snowflake_rest::server::router::{create_auth_router, create_router};
 use api_snowflake_rest::server::server_models::Config as SnowflakeServerConfig;
 use api_snowflake_rest::server::state::AppState;
+use api_snowflake_rest_sessions::layer::Host;
 use api_snowflake_rest_sessions::session::{SESSION_EXPIRATION_SECONDS, SessionStore};
+use axum::Extension;
 use axum::body::Body as AxumBody;
 use axum::extract::connect_info::ConnectInfo;
 use axum::{Router, middleware};
@@ -119,10 +121,12 @@ impl LambdaApp {
         let snowflake_router = create_router()
             .with_state(state.clone())
             .layer(compression_layer.clone())
+            .layer(Extension(Host(String::default())))
             .layer(middleware::from_fn_with_state(state.clone(), require_auth));
         let snowflake_auth_router = create_auth_router()
             .with_state(state.clone())
-            .layer(compression_layer);
+            .layer(compression_layer)
+            .layer(Extension(Host(String::default())));
         let router = Router::new().merge(snowflake_router.merge(snowflake_auth_router));
 
         Ok(Self { router, state })

crates/embucketd/src/main.rs

@@ -12,7 +12,9 @@ use api_snowflake_rest::server::router::create_auth_router as create_snowflake_a
 use api_snowflake_rest::server::router::create_router as create_snowflake_router;
 use api_snowflake_rest::server::server_models::Config;
 use api_snowflake_rest::server::state::AppState as SnowflakeAppState;
+use api_snowflake_rest_sessions::layer::Host;
 use api_snowflake_rest_sessions::session::{SESSION_EXPIRATION_SECONDS, SessionStore};
+use axum::Extension;
 use axum::middleware;
 use axum::{
     Json, Router,
@@ -189,13 +191,15 @@ async fn async_main(
     let snowflake_router = create_snowflake_router()
         .with_state(snowflake_state.clone())
         .layer(compression_layer.clone())
+        .layer(Extension(Host(String::default())))
         .layer(middleware::from_fn_with_state(
             snowflake_state.clone(),
             snowflake_require_auth,
         ));
     let snowflake_auth_router = create_snowflake_auth_router()
         .with_state(snowflake_state.clone())
-        .layer(compression_layer);
+        .layer(compression_layer)
+        .layer(Extension(Host(String::default())));
     let snowflake_router = snowflake_router.merge(snowflake_auth_router);
 
     // --- OpenAPI specs ---