Merge remote-tracking branch 'origin/main' into feat/org

shuttle-hq · Apr 5, 2024 · 70a87ef · 70a87ef
2 parents 126e8f7 + d536aa4
commit 70a87ef
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 113 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/common/Cargo.toml b/common/Cargo.toml
@@ -44,6 +44,7 @@ wiremock = { workspace = true, optional = true }
 [features]
 axum = ["dep:axum"]
 claims = [
+    "axum",
     "bytes",
     "chrono/clock",
     "headers",

diff --git a/common/src/claims.rs b/common/src/claims.rs
@@ -5,10 +5,11 @@ use std::{
     task::{Context, Poll},
 };
 
+use axum::extract::FromRequestParts;
 use bytes::Bytes;
 use chrono::{Duration, Utc};
 use headers::{Authorization, HeaderMapExt};
-use http::{Request, StatusCode};
+use http::{request::Parts, Request, StatusCode};
 use http_body::combinators::UnsyncBoxBody;
 use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
 use opentelemetry::global;
@@ -332,6 +333,26 @@ impl Claim {
     }
 }
 
+/// Extract the claim from the request and fail with unauthorized if the claim doesn't exist
+#[axum::async_trait]
+impl<S> FromRequestParts<S> for Claim {
+    type Rejection = StatusCode;
+
+    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
+        let claim = parts
+            .extensions
+            .get::<Claim>()
+            .ok_or(StatusCode::UNAUTHORIZED)?;
+
+        // Record current account name for tracing purposes
+        Span::current().record("account.user_id", &claim.sub);
+
+        trace!(?claim, "got user");
+
+        Ok(claim.clone())
+    }
+}
+
 // Future for layers that just return the inner response
 #[pin_project]
 pub struct ResponseFuture<F>(#[pin] pub F);

diff --git a/gateway/src/api/latest.rs b/gateway/src/api/latest.rs
@@ -26,7 +26,7 @@ use shuttle_backends::metrics::{Metrics, TraceLayer};
 use shuttle_backends::project_name::ProjectName;
 use shuttle_backends::request_span;
 use shuttle_backends::ClaimExt;
-use shuttle_common::claims::{Scope, EXP_MINUTES};
+use shuttle_common::claims::{Claim, Scope, EXP_MINUTES};
 use shuttle_common::models::error::{ErrorKind, InvalidOrganizationName};
 use shuttle_common::models::{admin::ProjectResponse, project, stats};
 use shuttle_common::models::{organization, service};
@@ -48,7 +48,7 @@ use x509_parser::time::ASN1Time;
 
 use crate::acme::{AccountWrapper, AcmeClient, CustomDomain};
 use crate::api::tracing::project_name_tracing_layer;
-use crate::auth::{ScopedUser, User};
+use crate::auth::ScopedUser;
 use crate::service::{ContainerSettings, GatewayService};
 use crate::task::{self, BoxedTask};
 use crate::tls::{GatewayCertResolver, RENEWAL_VALIDITY_THRESHOLD_IN_DAYS};
@@ -132,12 +132,12 @@ async fn check_project_name(
 }
 async fn get_projects_list(
     State(RouterState { service, .. }): State<RouterState>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<AxumJson<Vec<project::Response>>, Error> {
     let mut projects = vec![];
     for p in service
         .permit_client
-        .get_user_projects(&id)
+        .get_user_projects(&sub)
         .await
         .map_err(|_| Error::from(ErrorKind::Internal))?
     {
@@ -164,7 +164,7 @@ async fn create_project(
     State(RouterState {
         service, sender, ..
     }): State<RouterState>,
-    User { id, claim, .. }: User,
+    claim: Claim,
     CustomErrorPath(project_name): CustomErrorPath<ProjectName>,
     AxumJson(config): AxumJson<project::Config>,
 ) -> Result<AxumJson<project::Response>, Error> {
@@ -173,7 +173,7 @@ async fn create_project(
     // Check that the user is within their project limits.
     let can_create_project = claim.can_create_project(
         service
-            .get_project_count(&id)
+            .get_project_count(&claim.sub)
             .await?
             .saturating_sub(is_cch_project as u32),
     );
@@ -185,7 +185,7 @@ async fn create_project(
     let project = service
         .create_project(
             project_name.clone(),
-            &id,
+            &claim.sub,
             claim.is_admin(),
             can_create_project,
             if is_cch_project {
@@ -399,7 +399,7 @@ async fn override_create_service(
     scoped_user: ScopedUser,
     req: Request<Body>,
 ) -> Result<Response<Body>, Error> {
-    let user_id = scoped_user.user.id.clone();
+    let user_id = scoped_user.claim.sub.clone();
     let posthog_client = state.posthog_client.clone();
     tokio::spawn(async move {
         let event = async_posthog::Event::new("shuttle_api_start_deployment", &user_id);
@@ -461,9 +461,9 @@ async fn route_project(
     let project_name = scoped_user.scope;
     let is_cch_project = project_name.is_cch_project();
 
-    if !scoped_user.user.claim.is_admin() {
+    if !scoped_user.claim.is_admin() {
         service
-            .has_capacity(is_cch_project, &scoped_user.user.claim.tier)
+            .has_capacity(is_cch_project, &scoped_user.claim.tier)
             .await?;
     }
 
@@ -472,16 +472,16 @@ async fn route_project(
         .await?
         .0;
     service
-        .route(&project.state, &project_name, &scoped_user.user.id, req)
+        .route(&project.state, &project_name, &scoped_user.claim.sub, req)
         .await
 }
 
 #[instrument(skip_all)]
 async fn get_organizations(
     State(RouterState { service, .. }): State<RouterState>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<AxumJson<Vec<organization::Response>>, Error> {
-    let orgs = service.permit_client.get_organizations(&id).await?;
+    let orgs = service.permit_client.get_organizations(&sub).await?;
 
     Ok(AxumJson(orgs))
 }
@@ -490,7 +490,7 @@ async fn get_organizations(
 async fn create_organization(
     State(RouterState { service, .. }): State<RouterState>,
     CustomErrorPath(organization_name): CustomErrorPath<String>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<String, Error> {
     if organization_name.chars().count() > 30 {
         return Err(Error::from_kind(ErrorKind::InvalidOrganizationName(
@@ -503,7 +503,10 @@ async fn create_organization(
         display_name: organization_name.clone(),
     };
 
-    service.permit_client.create_organization(&id, &org).await?;
+    service
+        .permit_client
+        .create_organization(&sub, &org)
+        .await?;
 
     Span::current().record("shuttle.organization.id", &org.id);
 
@@ -514,11 +517,11 @@ async fn create_organization(
 async fn get_organization_projects(
     State(RouterState { service, .. }): State<RouterState>,
     CustomErrorPath(organization_id): CustomErrorPath<String>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<AxumJson<Vec<project::Response>>, Error> {
     let project_ids = service
         .permit_client
-        .get_organization_projects(&id, &organization_id)
+        .get_organization_projects(&sub, &organization_id)
         .await?;
 
     let mut projects = Vec::with_capacity(project_ids.len());
@@ -542,11 +545,11 @@ async fn get_organization_projects(
 async fn delete_organization(
     State(RouterState { service, .. }): State<RouterState>,
     CustomErrorPath(organization_id): CustomErrorPath<String>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<String, Error> {
     service
         .permit_client
-        .delete_organization(&id, &organization_id)
+        .delete_organization(&sub, &organization_id)
         .await?;
 
     Ok("Organization deleted".to_string())
@@ -556,11 +559,11 @@ async fn delete_organization(
 async fn transfer_project_to_organization(
     State(RouterState { service, .. }): State<RouterState>,
     CustomErrorPath((organization_id, project_id)): CustomErrorPath<(String, String)>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<String, Error> {
     service
         .permit_client
-        .transfer_project_to_org(&id, &project_id, &organization_id)
+        .transfer_project_to_org(&sub, &project_id, &organization_id)
         .await?;
 
     Ok("Project transfered".to_string())
@@ -570,11 +573,11 @@ async fn transfer_project_to_organization(
 async fn transfer_project_from_organization(
     State(RouterState { service, .. }): State<RouterState>,
     CustomErrorPath((organization_id, project_id)): CustomErrorPath<(String, String)>,
-    User { id, .. }: User,
+    Claim { sub, .. }: Claim,
 ) -> Result<String, Error> {
     service
         .permit_client
-        .transfer_project_from_org(&id, &project_id, &organization_id)
+        .transfer_project_from_org(&sub, &project_id, &organization_id)
         .await?;
 
     Ok("Project transfered".to_string())

diff --git a/gateway/src/api/project_caller.rs b/gateway/src/api/project_caller.rs
@@ -43,7 +43,7 @@ impl ProjectCaller {
         Ok(Self {
             project: project.state,
             project_name,
-            user_id: scoped_user.user.id,
+            user_id: scoped_user.claim.sub,
             service,
             headers: headers.clone(),
         })