Added support for multi tenant authentication (#746)

## Summary of the Pull Request

_What is this about?_

## PR Checklist
* [x] Applies to work item: #562 
* [x] CLA signed. If not, go over [here](https://cla.opensource.microsoft.com/microsoft/onefuzz) and sign the CLI.
* [x] Tests added/passed
* [ ] Requires documentation to be updated
* [x] I've discussed this with core contributors already. If not checked, I'm ready to accept this work might be rejected in favor of a different grand plan. Issue number where discussion took place: #xxx

## Info on Pull Request

The end-to-end changes needed to have onefuzz deployed with multi-tenant authentication.

## Validation Steps Performed

_How does someone test & validate?_

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

@@ -25,6 +25,8 @@ pub struct StaticConfig {
 
     pub onefuzz_url: Url,
 
+    pub multi_tenant_domain: Option<String>,
+
     // TODO: remove the alias once the service has been updated to match
     #[serde(alias = "instrumentation_key")]
     pub instance_telemetry_key: Option<InstanceTelemetryKey>,
@@ -47,6 +49,8 @@ struct RawStaticConfig {
 
     pub onefuzz_url: Url,
 
+    pub multi_tenant_domain: Option<String>,
+
     // TODO: remove the alias once the service has been updated to match
     #[serde(alias = "instrumentation_key")]
     pub instance_telemetry_key: Option<InstanceTelemetryKey>,
@@ -73,14 +77,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,
             microsoft_telemetry_key: config.microsoft_telemetry_key,
             instance_telemetry_key: config.instance_telemetry_key,
             heartbeat_queue: config.heartbeat_queue,
@@ -102,6 +108,7 @@ 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 = std::env::var("ONEFUZZ_MULTI_TENANT_DOMAIN").ok();
         let onefuzz_url = Url::parse(&std::env::var("ONEFUZZ_URL")?)?;
         let pool_name = std::env::var("ONEFUZZ_POOL")?;
 
@@ -125,15 +132,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(),
+            tenant,
+            multi_tenant_domain.clone(),
+        )
+        .into();
 
         Ok(Self {
             instance_id,
             credentials,
             pool_name,
             onefuzz_url,
+            multi_tenant_domain,
             instance_telemetry_key,
             microsoft_telemetry_key,
             heartbeat_queue,

src/agent/onefuzz/src/auth.rs

@@ -88,25 +88,48 @@ pub struct ClientCredentials {
     client_secret: Secret<String>,
     resource: String,
     tenant: String,
+    multi_tenant_domain: Option<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,
+        tenant: String,
+        multi_tenant_domain: Option<String>,
+    ) -> Self {
         let client_secret = client_secret.into();
 
         Self {
             client_id,
             client_secret,
             resource,
             tenant,
+            multi_tenant_domain,
         }
     }
 
     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 +138,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,18 +170,34 @@ 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>) -> Result<Self> {
+        let resource = if let Some(domain) = multi_tenant_domain.clone() {
+            let resource_url = Url::parse(&resource)?;
+            let host = resource_url.host_str().ok_or_else(|| {
+                anyhow::format_err!("resource URL does not have a host string: {}", resource_url)
+            })?;
+            let instance: Vec<&str> = host.split('.').collect();
+            format!("https://{}/{}", domain, instance[0])
+        } else {
+            resource
+        };
+
+        Ok(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);
         url
@@ -198,3 +237,51 @@ impl From<ManagedIdentityAccessTokenBody> for AccessToken {
         AccessToken { secret }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use anyhow::Result;
+
+    #[test]
+    fn test_managed_creds_with_valid_single_tenant() -> Result<()> {
+        let resource = "https://host-26.azurewebsites.net";
+
+        let managed_creds = ManagedIdentityCredentials::new(resource.to_string(), None)?;
+
+        assert_eq!(managed_creds.resource, resource);
+        Ok(())
+    }
+
+    #[test]
+    fn test_managed_creds_with_valid_multi_tenant_domain() -> Result<()> {
+        let resource = "https://host-26.azurewebsites.net";
+        let multi_tenant_domain = "mycloud.contoso.com";
+
+        let managed_creds = ManagedIdentityCredentials::new(
+            resource.to_string(),
+            Some(multi_tenant_domain.to_string()),
+        )?;
+
+        let expected = "https://mycloud.contoso.com/host-26";
+
+        assert_eq!(managed_creds.resource, expected);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_managed_creds_with_invalid_single_tenant() -> Result<()> {
+        let resource = "invalid-url";
+        let multi_tenant_domain = "mycloud.contoso.com";
+
+        let managed_creds = ManagedIdentityCredentials::new(
+            resource.to_string(),
+            Some(multi_tenant_domain.to_string()),
+        );
+
+        assert!(managed_creds.is_err());
+
+        Ok(())
+    }
+}

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,11 @@ 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,7 +275,7 @@
                         "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')]"