FMI Proof-of-concept

Author: rayluo

docs/fmi.md

@@ -0,0 +1,22 @@
+# FMI Flows
+
+```mermaid
+sequenceDiagram
+title Legend: -------- FMI Cred  ________ FMI Token
+autonumber
+    participant RMA as Resource Manager Authority (RMA)
+    participant sub-RMA as sub-RMA
+    participant eSTS
+    participant Resource
+
+    RMA-->>eSTS: client_id=rma_id<br>&client_assertion=SNI<br>&scope=api://AzureFMITokenExchange<br>&fmi_path=/eid1/c/<cloud>/t/<tenantid>/a/<rma>/<fmi_path><br>&...
+    eSTS-->>RMA: Return FMI cred
+
+    Note over RMA, sub-RMA: Somehow transfer the FMI cred to sub-RMA
+
+    sub-RMA->>eSTS: client_id=urn:microsoft:identity:fmi<br>&client_assertion=FMI cred<br>&scope=api://a1b2c3...<br>&fmi_path=/eid1/c/<cloud>/t/<tenantid>/a/<rma>/<fmi_path><sub_path><br>&...
+    eSTS->>sub-RMA: Return FMI token
+
+    sub-RMA->>Resource: Request with FMI token
+    Resource->>sub-RMA: Access granted
+```

msal/application.py

@@ -256,6 +256,7 @@ class ClientApplication(object):
     _TOKEN_CACHE_DATA: dict[str, str] = {  # field_in_data: field_in_cache
         "key_id": "key_id",  # Some token types (SSH-certs, POP) are bound to a key
         "req_ds_cnf": "req_ds_cnf",  # Used in CDT scenario
+        "fmi_path": "fmi_path",  # FMI artifacts are bound to their path
     }
 
     @functools.lru_cache(maxsize=2)
@@ -2386,6 +2387,7 @@ def acquire_token_for_client(
         delegation_constraints: Optional[list] = None,
         delegation_confirmation_key=None,  # A Cyprtography's RSAPrivateKey-like object
             # TODO: Support ECC key? https://github.com/pyca/cryptography/issues/4093
+        fmi_path: Optional[str] = None,
         **kwargs
     ):
         """Acquires token for the current confidential client, not for an end user.
@@ -2419,6 +2421,7 @@ def acquire_token_for_client(
                 kwargs.pop("data", {}),
                 req_ds_cnf=_build_req_cnf(jwk)  # It is part of token cache key
                     if delegation_constraints else None,
+                fmi_path=fmi_path,
                 ),
             **kwargs))
         if delegation_constraints and not result.get("error"):

msal/token_cache.py

@@ -84,6 +84,7 @@ def __init__(self):
                         # Note: New field(s) can be added here
                         #key_id=None,
                         req_ds_cnf=None,
+                        fmi_path=None,
                         **ignored_payload_from_a_real_token:
                     "-".join([  # Note: Could use a hash here to shorten key length
                         home_account_id or "",
@@ -100,6 +101,7 @@ def __init__(self):
                             #       instead of response scope,
                             #       so that a search() can probably have O(1) hit.
                             if req_ds_cnf else "",  # CDT
+                        fmi_path or "",  # See the TODO above
                         ]).lower(),
             self.CredentialType.ID_TOKEN:
                 lambda home_account_id=None, environment=None, client_id=None,

tests/test_application.py

@@ -931,3 +931,62 @@ def test_acquire_token_for_client_should_return_a_cdt(self):
             self.assertAppObtainsCdt(app, ["scope1", "scope2"])
             self.assertEqual(mocked_post.call_count, 2)
 
+
+@patch("msal.authority.tenant_discovery", new=Mock(return_value={
+    "authorization_endpoint": "https://contoso.com/placeholder",
+    "token_endpoint": "https://contoso.com/placeholder",
+    }))
+class FmiTestCase(unittest.TestCase):
+
+    def assertFmi(self, result: dict) -> None:
+        self.assertIsNotNone(
+            result.get("access_token"), "Encountered {}: {}".format(
+                result.get("error"), result.get("error_description")))
+        self.assertIn(result["token_type"], ("fmi_cred", "fmi_token"))
+
+    def assertAppObtainsAndCaches(self, app, scopes, fmi_path) -> dict:
+        result = app.acquire_token_for_client(scopes, fmi_path=fmi_path)
+        self.assertFmi(result)
+        self.assertEqual(result["token_source"], "identity_provider")
+
+        result = app.acquire_token_for_client(scopes, fmi_path=fmi_path)
+        self.assertFmi(result)
+        self.assertEqual(result["token_source"], "cache")
+
+        return result
+
+    def assertAppCachesByFmipath(self, app, scopes, fmi_path) -> dict:
+        with patch.object(app.http_client, "post", return_value=MinimalResponse(
+            status_code=200, text=json.dumps({
+                "token_type": "fmi_cred"
+                    if scopes == ["api://AzureFMITokenExchange"] else "fmi_token",
+                "access_token": "payload",
+                "expires_in": 3600,
+        }))) as mocked_post:
+            result = self.assertAppObtainsAndCaches(app, scopes, fmi_path)
+            self.assertAppObtainsAndCaches(app, scopes, fmi_path + "/subpath")
+            self.assertEqual(mocked_post.call_count, 2)
+            self.assertEqual(
+                2,
+                len(list(app.token_cache.search(
+                    msal.TokenCache.CredentialType.ACCESS_TOKEN))),
+                "Should cache both tokens")
+            logger.debug("cache=%s", json.dumps(app.token_cache._cache, indent=2))
+            return result
+
+    def test_fmi_with_rma_and_sub_rma(self):
+        fmi_cred = self.assertAppCachesByFmipath(
+            msal.ConfidentialClientApplication(
+                "rma_client_id",
+                client_credential={"client_assertion": "some assertion"},
+            ),
+            ["api://AzureFMITokenExchange"],
+            "/eid1/c/<cloud>/t/<tenantid>/a/<rma>/<fmi_path>")
+        self.assertAppCachesByFmipath(
+            msal.ConfidentialClientApplication(
+                "urn:microsoft:identity:fmi",
+                client_credential={"client_assertion": fmi_cred["access_token"]},
+            ),
+            ["https://graph.microsoft.com/.default"],
+            "/eid1/c/<cloud>/t/<tenantid>/a/<rma>/<fmi_path>/foo")
+