feat: add general consumption enforcement for JWT

514-labs · Sep 25, 2024 · ce960f6 · ce960f6
1 parent 6a987c1
commit ce960f6
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/apps/framework-cli/src/framework/typescript/consumption.rs b/apps/framework-cli/src/framework/typescript/consumption.rs
@@ -26,15 +26,22 @@ pub fn run(
         .as_ref()
         .map(|jwt| jwt.secret.clone())
         .unwrap_or("".to_string());
+
     let jwt_issuer = jwt_config
         .as_ref()
         .map(|jwt| jwt.issuer.clone())
         .unwrap_or("".to_string());
+
     let jwt_audience = jwt_config
         .as_ref()
         .map(|jwt| jwt.audience.clone())
         .unwrap_or("".to_string());
 
+    let enforce_on_all_consumptions_apis = jwt_config
+        .as_ref()
+        .map(|jwt| jwt.enforce_on_all_consumptions_apis.to_string())
+        .unwrap_or("false".to_string());
+
     let args = vec![
         consumption_path.to_str().unwrap(),
         &clickhouse_config.db_name,
@@ -46,6 +53,7 @@ pub fn run(
         &jwt_secret,
         &jwt_issuer,
         &jwt_audience,
+        &enforce_on_all_consumptions_apis,
     ];
 
     let mut consumption_process = bin::run(CONSUMPTION_RUNNER_BIN, project_path, &args)?;

diff --git a/apps/framework-cli/src/project.rs b/apps/framework-cli/src/project.rs
@@ -116,6 +116,8 @@ pub struct Project {
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct JwtConfig {
+    #[serde(default)]
+    pub enforce_on_all_consumptions_apis: bool,
     pub secret: String,
     pub issuer: String,
     pub audience: String,

diff --git a/packages/ts-moose-lib/src/consumption-apis/runner.ts b/packages/ts-moose-lib/src/consumption-apis/runner.ts
@@ -21,6 +21,7 @@ const [
   JWT_SECRET, // Optional we will need to bring a proper cli parsing tool to help to make sure this is more resilient. or make it one json object
   JWT_ISSUER, // Optional
   JWT_AUDIENCE, // Optional
+  ENFORCE_ON_ALL_CONSUMPTIONS_APIS, // Optional
 ] = process.argv;
 
 const clickhouseConfig = {
@@ -53,8 +54,21 @@ const apiHandler =
             jwtPayload = payload;
           } catch (error) {
             console.log("JWT verification failed");
+            if (ENFORCE_ON_ALL_CONSUMPTIONS_APIS === "true") {
+              res.writeHead(401, { "Content-Type": "application/json" });
+              res.end(JSON.stringify({ error: "Unauthorized" }));
+              return;
+            }
           }
+        } else if (ENFORCE_ON_ALL_CONSUMPTIONS_APIS === "true") {
+          res.writeHead(401, { "Content-Type": "application/json" });
+          res.end(JSON.stringify({ error: "Unauthorized" }));
+          return;
         }
+      } else if (ENFORCE_ON_ALL_CONSUMPTIONS_APIS === "true") {
+        res.writeHead(401, { "Content-Type": "application/json" });
+        res.end(JSON.stringify({ error: "Unauthorized" }));
+        return;
       }
 
       const pathName = createPath(fileName);