Added support for multi-tenant authentication

src/agent/onefuzz-supervisor/src/auth.rs

@@ -87,26 +87,48 @@ pub struct ClientCredentials {
     client_id: Uuid,
     client_secret: Secret<String>,
     resource: String,
+    multi_tenant_domain: Option<String>,
     tenant: String,
 }
 
 impl ClientCredentials {
-    pub fn new(client_id: Uuid, client_secret: String, resource: String, tenant: String) -> Self {
+    pub fn new(
+        client_id: Uuid,
+        client_secret: String,
+        resource: String,
+        multi_tenant_domain: Option<String>,
+        tenant: String,
+    ) -> Self {
         let client_secret = client_secret.into();
 
         Self {
             client_id,
             client_secret,
             resource,
+            multi_tenant_domain,
             tenant,
         }
     }
 
     pub async fn access_token(&self) -> Result<AccessToken> {
+        let (authority, resource) = if let Some(domain) = &self.multi_tenant_domain {
+            let url = Url::parse(&self.resource.clone())?;
+            let host = url.host_str().ok_or_else(|| {
+                anyhow::format_err!("resource URL does not have a host string: {}", url)
+            })?;
+            let instance: Vec<&str> = host.split('.').collect();
+            (
+                String::from("common"),
+                format!("https://{}/{}/", &domain, instance[0]),
+            )
+        } else {
+            (self.tenant.clone(), self.resource.clone())
+        };
+
         let mut url = Url::parse("https://login.microsoftonline.com")?;
         url.path_segments_mut()
             .expect("Authority URL is cannot-be-a-base")
-            .extend(&[&self.tenant, "oauth2", "v2.0", "token"]);
+            .extend(&[&authority.clone(), "oauth2", "v2.0", "token"]);
 
         let response = reqwest::Client::new()
             .post(url)
@@ -115,8 +137,8 @@ impl ClientCredentials {
                 ("client_id", self.client_id.to_hyphenated().to_string()),
                 ("client_secret", self.client_secret.expose_ref().to_string()),
                 ("grant_type", "client_credentials".into()),
-                ("tenant", self.tenant.clone()),
-                ("scope", format!("{}.default", self.resource)),
+                ("tenant", authority),
+                ("scope", format!("{}.default", resource)),
             ])
             .send_retry_default()
             .await?
@@ -147,20 +169,32 @@ impl From<ClientAccessTokenBody> for AccessToken {
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
 pub struct ManagedIdentityCredentials {
     resource: String,
+    multi_tenant_domain: Option<String>,
 }
 
 const MANAGED_IDENTITY_URL: &str =
     "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01";
 
 impl ManagedIdentityCredentials {
-    pub fn new(resource: String) -> Self {
-        Self { resource }
+    pub fn new(resource: String, multi_tenant_domain: Option<String>) -> Self {
+        Self {
+            resource,
+            multi_tenant_domain,
+        }
     }
 
     fn url(&self) -> Url {
         let mut url = Url::parse(MANAGED_IDENTITY_URL).unwrap();
-        url.query_pairs_mut()
-            .append_pair("resource", &self.resource);
+        let resource = if let Some(domain) = &self.multi_tenant_domain {
+            let uri = Url::parse(&self.resource).unwrap();
+            let host = uri.host_str().unwrap();
+            let instance: Vec<&str> = host.split('.').collect();
+            format!("https://{}/{}", domain, instance[0])
+        } else {
+            self.resource.clone()
+        };
+
+        url.query_pairs_mut().append_pair("resource", &resource);
         url
     }
 

src/agent/onefuzz-supervisor/src/config.rs

@@ -25,6 +25,8 @@ pub struct StaticConfig {
 
     pub onefuzz_url: Url,
 
+    pub multi_tenant_domain: Option<String>,
+
     pub instrumentation_key: Option<Uuid>,
 
     pub telemetry_key: Option<Uuid>,
@@ -43,6 +45,8 @@ struct RawStaticConfig {
 
     pub onefuzz_url: Url,
 
+    pub multi_tenant_domain: Option<String>,
+
     pub instrumentation_key: Option<Uuid>,
 
     pub telemetry_key: Option<Uuid>,
@@ -65,14 +69,16 @@ impl StaticConfig {
                     .to_string()
                     .trim_end_matches('/')
                     .to_owned();
-                let managed = ManagedIdentityCredentials::new(resource);
+                let managed =
+                    ManagedIdentityCredentials::new(resource, config.multi_tenant_domain.clone());
                 managed.into()
             }
         };
         let config = StaticConfig {
             credentials,
             pool_name: config.pool_name,
             onefuzz_url: config.onefuzz_url,
+            multi_tenant_domain: config.multi_tenant_domain,
             instrumentation_key: config.instrumentation_key,
             telemetry_key: config.telemetry_key,
             heartbeat_queue: config.heartbeat_queue,
@@ -94,6 +100,10 @@ impl StaticConfig {
         let client_id = Uuid::parse_str(&std::env::var("ONEFUZZ_CLIENT_ID")?)?;
         let client_secret = std::env::var("ONEFUZZ_CLIENT_SECRET")?;
         let tenant = std::env::var("ONEFUZZ_TENANT")?;
+        let multi_tenant_domain = match std::env::var("ONEFUZZ_MULTI_TENANT_DOMAIN") {
+            Ok(v) => Some(v),
+            Err(_e) => None,
+        };
         let onefuzz_url = Url::parse(&std::env::var("ONEFUZZ_URL")?)?;
         let pool_name = std::env::var("ONEFUZZ_POOL")?;
 
@@ -115,15 +125,21 @@ impl StaticConfig {
             None
         };
 
-        let credentials =
-            ClientCredentials::new(client_id, client_secret, onefuzz_url.to_string(), tenant)
-                .into();
+        let credentials = ClientCredentials::new(
+            client_id,
+            client_secret,
+            onefuzz_url.to_string(),
+            multi_tenant_domain.clone(),
+            tenant,
+        )
+        .into();
 
         Ok(Self {
             instance_id,
             credentials,
             pool_name,
             onefuzz_url,
+            multi_tenant_domain,
             instrumentation_key,
             telemetry_key,
             heartbeat_queue,

src/api-service/__app__/onefuzzlib/extension.py

@@ -113,6 +113,10 @@ def build_pool_config(pool: Pool) -> str:
         instance_id=get_instance_id(),
     )
 
+    multi_tenant_domain = os.environ.get("MULTI_TENANT_DOMAIN")
+    if multi_tenant_domain:
+        config.multi_tenant_domain = multi_tenant_domain
+
     filename = f"{pool.name}/config.json"
 
     save_blob(

src/api-service/__app__/pool/__init__.py

@@ -40,6 +40,10 @@ def set_config(pool: Pool) -> Pool:
         ),
         instance_id=get_instance_id(),
     )
+
+    multi_tenant_domain = os.environ.get("MULTI_TENANT_DOMAIN")
+    if multi_tenant_domain:
+        pool.config.multi_tenant_domain = multi_tenant_domain
     return pool
 
 

src/deployment/azuredeploy.json

@@ -17,6 +17,15 @@
         "signedExpiry": {
             "type": "string"
         },
+        "app_func_issuer": {
+            "type": "string"
+        },
+        "app_func_audience": {
+            "type": "string"
+        },
+        "multi_tenant_domain": {
+            "type": "string"
+        },
         "diagnosticsLogsLevel": {
             "type": "string",
             "defaultValue": "Verbose",
@@ -202,6 +211,10 @@
                             "name": "AzureWebJobsStorage",
                             "value": "[concat('DefaultEndpointsProtocol=https;AccountName=',variables('storageAccountNameFunc'),';AccountKey=',listKeys(resourceId('Microsoft.Storage/storageAccounts', variables('storageAccountNameFunc')), '2019-06-01').keys[0].value,';EndpointSuffix=','core.windows.net')]"
                         },
+                        {
+                            "name": "MULTI_TENANT_DOMAIN",
+                            "value": "[parameters('multi_tenant_domain')]"
+                        },
                         {
                             "name": "AzureWebJobsDisableHomepage",
                             "value": "true"
@@ -262,11 +275,9 @@
                         "tokenStoreEnabled": true,
                         "clientId": "[parameters('clientId')]",
                         "clientSecret": "[parameters('clientSecret')]",
-                        "issuer": "[concat('https://sts.windows.net/', subscription().tenantId, '/')]",
+                        "issuer": "[parameters('app_func_issuer')]",
                         "defaultProvider": "AzureActiveDirectory",
-                        "allowedAudiences": [
-                            "[concat('https://', parameters('name'), '.azurewebsites.net')]"
-                        ],
+                        "allowedAudiences": ["[parameters('app_func_audience')]"],
                         "isAadAutoProvisioned": false
                     }
                 },