Merge pull request #5327 from systeminit/jkeiser/local-module-index

Support local auth api on module-index

bin/module-index/src/args.rs

@@ -59,6 +59,10 @@ pub(crate) struct Args {
     )]
     pub(crate) log_json: bool,
 
+    /// Override for the auth api url
+    #[arg(long, env = "SI_AUTH_API_URL")]
+    pub(crate) auth_api_url: Option<String>,
+
     /// PostgreSQL connection pool dbname [example: myapp]
     #[arg(long, env)]
     pub(crate) pg_dbname: Option<String>,
@@ -136,6 +140,9 @@ impl TryFrom<Args> for Config {
 
     fn try_from(args: Args) -> Result<Self, Self::Error> {
         ConfigFile::layered_load(NAME, |config_map| {
+            if let Some(auth_api_url) = args.auth_api_url {
+                config_map.set("auth_api_url", auth_api_url);
+            }
             if let Some(dbname) = args.pg_dbname {
                 config_map.set("pg.dbname", dbname);
             }

dev/Tiltfile

@@ -329,13 +329,13 @@ si_buck2_resource(
     resource_deps = [
         "postgres",
     ],
-    readiness_probe = probe(
-        period_secs = 5,
-        http_get = http_get_action(
-            port = 9001,
-            path = "/",
-        ),
-    ),
+    # readiness_probe = probe(
+    #     period_secs = 5,
+    #     http_get = http_get_action(
+    #         port = 9001,
+    #         path = "/",
+    #     ),
+    # ),
     trigger_mode = TRIGGER_MODE_MANUAL,
 )
 

lib/module-index-server/src/app_state.rs

@@ -19,6 +19,7 @@ pub enum ShutdownSource {}
 pub struct AppState {
     /// A PostgreSQL connection pool.
     pg_pool: DatabaseConnection,
+    auth_api_url: String,
     jwt_public_signing_key_chain: JwtPublicSigningKeyChain,
     posthog_client: PosthogClient,
     aws_creds: AwsCredentials,
@@ -35,6 +36,7 @@ impl AppState {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         pg_pool: DatabaseConnection,
+        auth_api_url: String,
         jwt_public_signing_key_chain: JwtPublicSigningKeyChain,
         posthog_client: PosthogClient,
         aws_creds: AwsCredentials,
@@ -43,6 +45,7 @@ impl AppState {
     ) -> Self {
         Self {
             pg_pool,
+            auth_api_url,
             jwt_public_signing_key_chain,
             posthog_client,
             aws_creds,
@@ -57,6 +60,11 @@ impl AppState {
         &self.pg_pool
     }
 
+    /// Gets the URL to the auth api
+    pub fn auth_api_url(&self) -> &str {
+        &self.auth_api_url
+    }
+
     /// Gets a reference to the public key used to sign the JWT
     pub fn jwt_public_signing_key(&self) -> &JwtPublicSigningKeyChain {
         &self.jwt_public_signing_key_chain

lib/module-index-server/src/config.rs

@@ -40,6 +40,10 @@ fn get_default_socket_addr() -> SocketAddr {
     SocketAddr::from(([0, 0, 0, 0], 5157))
 }
 
+fn default_auth_api_url() -> String {
+    auth_api_client::PROD_AUTH_API_ENDPOINT.to_string()
+}
+
 #[derive(Debug, Builder)]
 pub struct Config {
     #[builder(default = "get_default_socket_addr()")]
@@ -51,6 +55,9 @@ pub struct Config {
     #[builder(default = "random_instance_id()")]
     instance_id: String,
 
+    #[builder(default = "default_auth_api_url()")]
+    auth_api_url: String,
+
     jwt_signing_public_key_path: CanonicalFile,
     jwt_signing_public_key_algo: JwtAlgo,
 
@@ -86,6 +93,11 @@ impl Config {
         self.instance_id.as_ref()
     }
 
+    /// Gets the auth API URL.
+    pub fn auth_api_url(&self) -> &str {
+        &self.auth_api_url
+    }
+
     /// Gets a reference to the config's jwt signing public key path.
     #[must_use]
     pub fn jwt_signing_public_key_path(&self) -> &CanonicalFile {
@@ -130,6 +142,8 @@ pub struct ConfigFile {
     socket_addr: SocketAddr,
     #[serde(default = "random_instance_id")]
     instance_id: String,
+    #[serde(default)]
+    auth_api_url: String,
     #[serde(default = "default_jwt_signing_public_key_path")]
     pub jwt_signing_public_key_path: String,
     #[serde(default = "default_jwt_signing_public_key_algo")]
@@ -156,6 +170,7 @@ impl Default for ConfigFile {
             },
             socket_addr: get_default_socket_addr(),
             instance_id: random_instance_id(),
+            auth_api_url: default_auth_api_url(),
             jwt_signing_public_key_path: default_jwt_signing_public_key_path(),
             jwt_signing_public_key_algo: default_jwt_signing_public_key_algo(),
             jwt_secondary_signing_public_key_path: None,
@@ -180,6 +195,7 @@ impl TryFrom<ConfigFile> for Config {
         config.pg_pool(value.pg);
         config.socket_addr(value.socket_addr);
         config.instance_id(value.instance_id);
+        config.auth_api_url(value.auth_api_url);
         config.jwt_signing_public_key_path(value.jwt_signing_public_key_path.try_into()?);
         config.jwt_signing_public_key_algo(value.jwt_signing_public_key_algo);
         config.posthog(value.posthog);

lib/module-index-server/src/routes/list_modules_route.rs

@@ -58,7 +58,8 @@ pub async fn list_module_route(
     let query = si_module::Entity::find();
 
     let su = request.su.unwrap_or(false)
-        && is_systeminit_auth_token(&auth_token, state.token_emails()).await?;
+        && is_systeminit_auth_token(state.auth_api_url(), &auth_token, state.token_emails())
+            .await?;
 
     let kind = request.kind.unwrap_or(si_module::ModuleKind::Module);
 

lib/module-index-server/src/routes/promote_builtin_route.rs

@@ -59,7 +59,7 @@ pub async fn promote_builtin_route(
     State(state): State<AppState>,
     mut multipart: Multipart,
 ) -> Result<Json<Option<ModuleDetailsResponse>>, PromoteModuleError> {
-    if !is_systeminit_auth_token(&auth_token, state.token_emails()).await? {
+    if !is_systeminit_auth_token(state.auth_api_url(), &auth_token, state.token_emails()).await? {
         return Ok(Json(None));
     }
 

lib/module-index-server/src/routes/reject_module_route.rs

@@ -56,7 +56,7 @@ pub async fn reject_module(
     State(state): State<AppState>,
     mut multipart: Multipart,
 ) -> Result<Json<Option<ModuleDetailsResponse>>, RejectModuleError> {
-    if !is_systeminit_auth_token(&auth_token, state.token_emails()).await? {
+    if !is_systeminit_auth_token(state.auth_api_url(), &auth_token, state.token_emails()).await? {
         return Ok(Json(None));
     }
 

lib/module-index-server/src/server.rs

@@ -115,6 +115,7 @@ impl Server<(), ()> {
 
         let (service, shutdown_rx, shutdown_broadcast_rx) = build_service(
             pg_pool,
+            config.auth_api_url().to_owned(),
             jwt_public_signing_key,
             posthog_client,
             aws_creds,
@@ -257,6 +258,7 @@ where
 
 pub fn build_service(
     pg_pool: DatabaseConnection,
+    auth_api_url: String,
     jwt_public_signing_key_chain: JwtPublicSigningKeyChain,
     posthog_client: PosthogClient,
     aws_creds: AwsCredentials,
@@ -267,6 +269,7 @@ pub fn build_service(
 
     let state = AppState::new(
         pg_pool,
+        auth_api_url,
         jwt_public_signing_key_chain,
         posthog_client,
         aws_creds,

lib/module-index-server/src/whoami.rs

@@ -17,6 +17,7 @@ pub enum WhoamiError {
 type WhoamiResult<T> = Result<T, WhoamiError>;
 
 pub async fn get_email_for_auth_token(
+    auth_api_url: &str,
     token: &str,
     token_map: Arc<Mutex<HashMap<String, String>>>,
 ) -> WhoamiResult<String> {
@@ -25,8 +26,7 @@ pub async fn get_email_for_auth_token(
     match token_map.get(token) {
         Some(email) => Ok(email.into()),
         None => {
-            let auth_api_client =
-                AuthApiClient::new(auth_api_client::PROD_AUTH_API_ENDPOINT.try_into()?, token);
+            let auth_api_client = AuthApiClient::new(auth_api_url.try_into()?, token);
 
             let whoami = auth_api_client.whoami().await?;
 
@@ -42,10 +42,11 @@ pub fn is_systeminit_email(email: &str) -> bool {
 }
 
 pub async fn is_systeminit_auth_token(
+    auth_api_url: &str,
     token: &str,
     token_map: Arc<Mutex<HashMap<String, String>>>,
 ) -> WhoamiResult<bool> {
     Ok(is_systeminit_email(
-        &get_email_for_auth_token(token, token_map).await?,
+        &get_email_for_auth_token(auth_api_url, token, token_map).await?,
     ))
 }