handle tokens from x-ms-token-aad-id-token (#531)

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

@@ -3,6 +3,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+from typing import Optional
 from uuid import UUID
 
 import azure.functions as func
@@ -11,32 +12,46 @@ from onefuzztypes.enums import ErrorCode
 from onefuzztypes.models import Error, Result, UserInfo
 
 
-def parse_jwt_token(request: func.HttpRequest) -> Result[UserInfo]:
-    """ Obtains the Access Token from the Authorization Header """
-
+def get_bearer_token(request: func.HttpRequest) -> Optional[str]:
     auth: str = request.headers.get("Authorization", None)
     if not auth:
-        return Error(
-            code=ErrorCode.INVALID_REQUEST, errors=["Authorization header is expected"]
-        )
+        return None
 
     parts = auth.split()
 
     if len(parts) != 2:
-        return Error(
-            code=ErrorCode.INVALID_REQUEST, errors=["Invalid authorization header"]
-        )
+        return None
 
     if parts[0].lower() != "bearer":
+        return None
+
+    return parts[1]
+
+
+def get_auth_token(request: func.HttpRequest) -> Optional[str]:
+    token = get_bearer_token(request)
+    if token is not None:
+        return token
+
+    token_header = request.headers.get("x-ms-token-aad-id-token", None)
+    if token_header is None:
+        return None
+    return str(token_header)
+
+
+def parse_jwt_token(request: func.HttpRequest) -> Result[UserInfo]:
+    """ Obtains the Access Token from the Authorization Header """
+    token_str = get_auth_token(request)
+    if token_str is None:
         return Error(
             code=ErrorCode.INVALID_REQUEST,
-            errors=["Authorization header must start with Bearer"],
+            errors=["unable to find authorization token"],
         )
 
     # This token has already been verified by the azure authentication layer
-    token = jwt.decode(parts[1], verify=False)
+    token = jwt.decode(token_str, verify=False)
 
-    application_id = UUID(token["appid"])
+    application_id = UUID(token["appid"]) if "appid" in token else None
     object_id = UUID(token["oid"]) if "oid" in token else None
     upn = token.get("upn")
     return UserInfo(application_id=application_id, object_id=object_id, upn=upn)