From 35876bbcfe58e43485e4af1f0ba28b427111186c Mon Sep 17 00:00:00 2001
From: Alex Shovlin <shovlia@amazon.com>
Date: Wed, 16 Oct 2024 15:59:35 -0400
Subject: [PATCH] Add timeout for the OAuth redirect

---
 awscli/customizations/sso/utils.py          | 18 ++++++++++++++++--
 tests/unit/customizations/sso/test_utils.py | 19 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/awscli/customizations/sso/utils.py b/awscli/customizations/sso/utils.py
index 21920a333f979..18fa2be4bcd99 100644
--- a/awscli/customizations/sso/utils.py
+++ b/awscli/customizations/sso/utils.py
@@ -15,13 +15,17 @@
 import logging
 import os
 import socket
+import time
 import webbrowser
 from functools import partial
 from http.server import HTTPServer, BaseHTTPRequestHandler
 
 from botocore.compat import urlparse, parse_qs
 from botocore.credentials import JSONFileCache
-from botocore.exceptions import AuthCodeFetcherError
+from botocore.exceptions import (
+    AuthCodeFetcherError,
+    PendingAuthorizationExpiredError,
+)
 from botocore.utils import SSOTokenFetcher, SSOTokenFetcherAuth
 from botocore.utils import original_ld_library_path
 
@@ -191,6 +195,11 @@ class AuthCodeFetcher:
     """Manages the local web server that will be used
     to retrieve the authorization code from the OAuth callback
     """
+    # How many seconds handle_request should wait for an incoming request
+    _REQUEST_TIMEOUT = 10
+    # How long we wait overall for the callback
+    _OVERALL_TIMEOUT = 60 * 10
+
     def __init__(self):
         self._auth_code = None
         self._state = None
@@ -201,6 +210,7 @@ def __init__(self):
         try:
             handler = partial(OAuthCallbackHandler, self)
             self.http_server = HTTPServer(('', 0), handler)
+            self.http_server.timeout = self._REQUEST_TIMEOUT
         except socket.error as e:
             raise AuthCodeFetcherError(error_msg=e)
 
@@ -214,10 +224,14 @@ def get_auth_code_and_state(self):
         """Blocks until the expected redirect request with either the
         authorization code/state or and error is handled
         """
-        while not self._is_done:
+        start = time.time()
+        while not self._is_done and time.time() < start + self._OVERALL_TIMEOUT:
             self.http_server.handle_request()
         self.http_server.server_close()
 
+        if not self._is_done:
+            raise PendingAuthorizationExpiredError
+
         return self._auth_code, self._state
 
     def set_auth_code_and_state(self, auth_code, state):
diff --git a/tests/unit/customizations/sso/test_utils.py b/tests/unit/customizations/sso/test_utils.py
index 818de315144cb..3efd3cb3a5738 100644
--- a/tests/unit/customizations/sso/test_utils.py
+++ b/tests/unit/customizations/sso/test_utils.py
@@ -13,9 +13,10 @@
 import os
 import threading
 import webbrowser
-
 import pytest
 import urllib3
+
+from botocore.exceptions import PendingAuthorizationExpiredError
 from botocore.session import Session
 
 from awscli.compat import StringIO
@@ -249,3 +250,19 @@ def test_error(self):
         assert response.status == 200
         assert actual_code is None
         assert actual_state is None
+
+
+@mock.patch(
+    'awscli.customizations.sso.utils.AuthCodeFetcher._REQUEST_TIMEOUT',
+    0.1
+)
+@mock.patch(
+    'awscli.customizations.sso.utils.AuthCodeFetcher._OVERALL_TIMEOUT',
+    0.1
+)
+def test_get_auth_code_and_state_timeout():
+    """Tests the timeout case separately of TestAuthCodeFetcher,
+    since we need to override the constants
+    """
+    with pytest.raises(PendingAuthorizationExpiredError):
+        AuthCodeFetcher().get_auth_code_and_state()