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,7 +77,8 @@ 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()
             }
         };
@@ -81,6 +86,7 @@ impl 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,8 +132,13 @@ impl StaticConfig {
                 None
             };
 
-        let credentials =
-            ClientCredentials::new(client_id, client_secret, onefuzz_url.to_string(), tenant)
+        let credentials = ClientCredentials::new(
+            client_id,
+            client_secret,
+            onefuzz_url.to_string(),
+            tenant,
+            multi_tenant_domain.clone(),
+        )
         .into();
 
         Ok(Self {
@@ -134,6 +146,7 @@ impl StaticConfig {
             credentials,
             pool_name,
             onefuzz_url,
+            multi_tenant_domain,
             instance_telemetry_key,
             microsoft_telemetry_key,
             heartbeat_queue,

src/agent/onefuzz/src/auth.rs

@@ -88,10 +88,17 @@ 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 {
@@ -99,14 +106,30 @@ impl ClientCredentials {
             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')]"

src/deployment/deploy.py

@@ -68,6 +68,7 @@ from registration import (
     assign_scaleset_role,
     authorize_application,
     register_application,
+    set_app_audience,
     update_pool_registration,
 )
 
@@ -76,6 +77,7 @@ ONEFUZZ_CLI_APP = "72f1562a-8c0c-41ea-beb9-fa2b71c80134"
 ONEFUZZ_CLI_AUTHORITY = (
     "https://login.microsoftonline.com/72f988bf-86f1-41af-91ab-2d7cd011db47"
 )
+COMMON_AUTHORITY = "https://login.microsoftonline.com/common"
 TELEMETRY_NOTICE = (
     "Telemetry collection on stats and OneFuzz failures are sent to Microsoft. "
     "To disable, delete the ONEFUZZ_TELEMETRY application setting in the "
@@ -118,6 +120,7 @@ class Client:
         migrations: List[str],
         export_appinsights: bool,
         log_service_principal: bool,
+        multi_tenant_domain: str,
         upgrade: bool,
     ):
         self.resource_group = resource_group
@@ -130,14 +133,19 @@ class Client:
         self.instance_specific = instance_specific
         self.third_party = third_party
         self.create_registration = create_registration
+        self.multi_tenant_domain = multi_tenant_domain
         self.upgrade = upgrade
         self.results: Dict = {
             "client_id": client_id,
             "client_secret": client_secret,
         }
+        if self.multi_tenant_domain:
+            authority = COMMON_AUTHORITY
+        else:
+            authority = ONEFUZZ_CLI_AUTHORITY
         self.cli_config: Dict[str, Union[str, UUID]] = {
             "client_id": ONEFUZZ_CLI_APP,
-            "authority": ONEFUZZ_CLI_AUTHORITY,
+            "authority": authority,
         }
         self.migrations = migrations
         self.export_appinsights = export_appinsights
@@ -230,7 +238,6 @@ class Client:
     def setup_rbac(self) -> None:
         """
         Setup the client application for the OneFuzz instance.
-
         By default, Service Principals do not have access to create
         client applications in AAD.
         """
@@ -274,6 +281,13 @@ class Client:
 
         if not existing:
             logger.info("creating Application registration")
+
+            if self.multi_tenant_domain:
+                url = "https://%s/%s" % (
+                    self.multi_tenant_domain,
+                    self.application_name,
+                )
+            else:
                 url = "https://%s.azurewebsites.net" % self.application_name
 
             params = ApplicationCreateParameters(
@@ -291,6 +305,7 @@ class Client:
                 ],
                 app_roles=app_roles,
             )
+
             app = client.applications.create(params)
 
             logger.info("creating service principal")
@@ -349,6 +364,27 @@ class Client:
                     app.object_id, ApplicationUpdateParameters(app_roles=app_roles)
                 )
 
+        if self.multi_tenant_domain and app.sign_in_audience == "AzureADMyOrg":
+            url = "https://%s/%s" % (
+                self.multi_tenant_domain,
+                self.application_name,
+            )
+            client.applications.patch(
+                app.object_id, ApplicationUpdateParameters(identifier_uris=[url])
+            )
+            set_app_audience(app.object_id, "AzureADMultipleOrgs")
+        elif (
+            not self.multi_tenant_domain
+            and app.sign_in_audience == "AzureADMultipleOrgs"
+        ):
+            set_app_audience(app.object_id, "AzureADMyOrg")
+            url = "https://%s.azurewebsites.net" % self.application_name
+            client.applications.patch(
+                app.object_id, ApplicationUpdateParameters(identifier_uris=[url])
+            )
+        else:
+            logger.debug("No change to App Registration signInAudence setting")
+
             creds = list(client.applications.list_password_credentials(app.object_id))
             client.applications.update_password_credentials(app.object_id, creds)
 
@@ -366,9 +402,13 @@ class Client:
             app_info = register_application(
                 "onefuzz-cli", self.application_name, OnefuzzAppRole.CliClient
             )
+            if self.multi_tenant_domain:
+                authority = COMMON_AUTHORITY
+            else:
+                authority = app_info.authority
             self.cli_config = {
                 "client_id": app_info.client_id,
-                "authority": app_info.authority,
+                "authority": authority,
             }
 
         else:
@@ -397,12 +437,31 @@ class Client:
         expiry = (datetime.now(TZ_UTC) + timedelta(days=365)).strftime(
             "%Y-%m-%dT%H:%M:%SZ"
         )
+
+        if self.multi_tenant_domain:
+            # clear the value in the Issuer Url field:
+            # https://docs.microsoft.com/en-us/sharepoint/dev/spfx/use-aadhttpclient-enterpriseapi-multitenant
+            app_func_audience = "https://%s/%s" % (
+                self.multi_tenant_domain,
+                self.application_name,
+            )
+            app_func_issuer = ""
+            multi_tenant_domain = {"value": self.multi_tenant_domain}
+        else:
+            app_func_audience = "https://%s.azurewebsites.net" % self.application_name
+            tenant_oid = str(self.cli_config["authority"]).split("/")[-1]
+            app_func_issuer = "https://sts.windows.net/%s/" % tenant_oid
+            multi_tenant_domain = {"value": ""}
+
         params = {
+            "app_func_audience": {"value": app_func_audience},
             "name": {"value": self.application_name},
             "owner": {"value": self.owner},
             "clientId": {"value": self.results["client_id"]},
             "clientSecret": {"value": self.results["client_secret"]},
+            "app_func_issuer": {"value": app_func_issuer},
             "signedExpiry": {"value": expiry},
+            "multi_tenant_domain": multi_tenant_domain,
             "workbookData": {"value": self.workbook_data},
         }
         deployment = Deployment(
@@ -779,13 +838,17 @@ class Client:
             if "client_secret" in self.cli_config
             else ""
         )
+        multi_tenant_domain = ""
+        if self.multi_tenant_domain:
+            multi_tenant_domain = "--tenant_domain %s" % self.multi_tenant_domain
         logger.info(
             "Update your CLI config via: onefuzz config --endpoint "
-            "https://%s.azurewebsites.net --authority %s --client_id %s %s",
+            "https://%s.azurewebsites.net --authority %s --client_id %s %s %s",
             self.application_name,
             self.cli_config["authority"],
             self.cli_config["client_id"],
             client_secret_arg,
+            multi_tenant_domain,
         )
 
 
@@ -898,6 +961,12 @@ def main() -> None:
         action="store_true",
         help="display service prinipal with info log level",
     )
+    parser.add_argument(
+        "--multi_tenant_domain",
+        type=str,
+        default=None,
+        help="enable multi-tenant authentication with this tenant domain",
+    )
     args = parser.parse_args()
 
     if shutil.which("func") is None:
@@ -921,6 +990,7 @@ def main() -> None:
         migrations=args.apply_migrations,
         export_appinsights=args.export_appinsights,
         log_service_principal=args.log_service_principal,
+        multi_tenant_domain=args.multi_tenant_domain,
         upgrade=args.upgrade,
     )
     if args.verbose:

src/deployment/registration.py

@@ -516,6 +516,35 @@ def assign_scaleset_role(onefuzz_instance_name: str, scaleset_name: str) -> None
         assign_scaleset_role_manually(onefuzz_instance_name, scaleset_name)
 
 
+def set_app_audience(objectId: str, audience: str) -> None:
+    # typical audience values: AzureADMyOrg, AzureADMultipleOrgs
+    http_body = {"signInAudience": audience}
+    try:
+        query_microsoft_graph(
+            method="PATCH",
+            resource="applications/%s" % objectId,
+            body=http_body,
+        )
+    except GraphQueryError:
+        query = (
+            "az rest --method patch --url "
+            "https://graph.microsoft.com/v1.0/applications/%s "
+            "--body '%s' --headers \"Content-Type\"=application/json"
+            % (objectId, http_body)
+        )
+        logger.warning(
+            "execute the following query in the azure portal bash shell and run deploy.py again : \n%s"
+            % query
+        )
+        err_str = (
+            "Unable to set signInAudience using Microsoft Graph Query API. \n"
+            + "The user must enable single/multi tenancy in the 'Authentication' blade of the "
+            + "Application Registration in the AAD web portal, or use the azure bash shell "
+            + "using the command given above."
+        )
+        raise Exception(err_str)
+
+
 def main() -> None:
     formatter = argparse.ArgumentDefaultsHelpFormatter
 

src/pytypes/onefuzztypes/models.py

@@ -364,6 +364,7 @@ class AgentConfig(BaseModel):
     heartbeat_queue: Optional[str]
     instrumentation_key: Optional[str]
     telemetry_key: Optional[str]
+    multi_tenant_domain: Optional[str]
     instance_id: UUID