From d6e68ef52bcba80da44b0522956113c26962048a Mon Sep 17 00:00:00 2001
From: Kostis Triantafyllakis <ctriant@admin.grnet.gr>
Date: Fri, 13 May 2022 11:22:06 +0300
Subject: [PATCH] Handle missing token exchange requested_token_type parameter

---
 src/idpyoidc/server/oauth2/token.py           |  8 +++-
 src/idpyoidc/server/oauth2/token_helper.py    |  1 +
 tests/test_server_36_oauth2_token_exchange.py | 47 ++++++++++++++++++-
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/src/idpyoidc/server/oauth2/token.py b/src/idpyoidc/server/oauth2/token.py
index 1600db92..94d10789 100755
--- a/src/idpyoidc/server/oauth2/token.py
+++ b/src/idpyoidc/server/oauth2/token.py
@@ -129,7 +129,13 @@ def process_request(self, request: Optional[Union[Message, dict]] = None, **kwar
         _context = self.server_get("endpoint_context")
 
         if isinstance(request, TokenExchangeRequest):
-            _handler_key = TOKEN_TYPES_MAPPING[request["requested_token_type"]]
+            requested_token_type = request.get(
+                "requested_token_type",
+                self.helper["urn:ietf:params:oauth:grant-type:token-exchange"].config[
+                    "default_requested_token_type"
+                ],
+            )
+            _handler_key = TOKEN_TYPES_MAPPING[requested_token_type]
         else:
             _handler_key = "access_token"
 
diff --git a/src/idpyoidc/server/oauth2/token_helper.py b/src/idpyoidc/server/oauth2/token_helper.py
index 300aefdb..744654ea 100755
--- a/src/idpyoidc/server/oauth2/token_helper.py
+++ b/src/idpyoidc/server/oauth2/token_helper.py
@@ -392,6 +392,7 @@ def __init__(self, endpoint, config=None):
                     "urn:ietf:params:oauth:token-type:access_token",
                     "urn:ietf:params:oauth:token-type:refresh_token",
                 ],
+                "default_requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
                 "policy": {"": {"callable": validate_token_exchange_policy}},
             }
         else:
diff --git a/tests/test_server_36_oauth2_token_exchange.py b/tests/test_server_36_oauth2_token_exchange.py
index f1c34261..30baae1e 100644
--- a/tests/test_server_36_oauth2_token_exchange.py
+++ b/tests/test_server_36_oauth2_token_exchange.py
@@ -239,7 +239,52 @@ def _mint_code(self, grant, client_id):
             {"refresh_token": "urn:ietf:params:oauth:token-type:refresh_token"},
         ],
     )
-    def test_token_exchange(self, token):
+    def test_token_exchange1(self, token):
+        """
+        Test that token exchange requests work correctly with only the required parameters
+        present
+        """
+        if list(token.keys())[0] == "refresh_token":
+            AUTH_REQ["scope"] = ["openid", "offline_access"]
+        areq = AUTH_REQ.copy()
+
+        session_id = self._create_session(areq)
+        grant = self.endpoint_context.authz(session_id, areq)
+        code = self._mint_code(grant, areq["client_id"])
+
+        _token_request = TOKEN_REQ_DICT.copy()
+        _token_request["code"] = code.value
+        _req = self.endpoint.parse_request(_token_request)
+        _resp = self.endpoint.process_request(request=_req)
+        _token_value = _resp["response_args"][list(token.keys())[0]]
+
+        token_exchange_req = TokenExchangeRequest(
+            grant_type="urn:ietf:params:oauth:grant-type:token-exchange",
+            subject_token=_token_value,
+            subject_token_type=token[list(token.keys())[0]]
+        )
+
+        _req = self.endpoint.parse_request(
+            token_exchange_req.to_urlencoded(),
+            {"headers": {"authorization": "Basic {}".format("Y2xpZW50XzI6aGVtbGlndA==")}},
+        )
+        _resp = self.endpoint.process_request(request=_req)
+        assert set(_resp["response_args"].keys()) == {
+            "access_token",
+            "token_type",
+            "scope",
+            "expires_in",
+            "issued_token_type",
+        }
+
+    @pytest.mark.parametrize(
+        "token",
+        [
+            {"access_token": "urn:ietf:params:oauth:token-type:access_token"},
+            {"refresh_token": "urn:ietf:params:oauth:token-type:refresh_token"},
+        ],
+    )
+    def test_token_exchange2(self, token):
         """
         Test that token exchange requests work correctly
         """