Merge pull request #452 from canonical/token_expiration

Prevent token from expiring while waiting for agent

cli/testflinger_cli/__init__.py

@@ -487,6 +487,17 @@ def pack_attachments(self, archive: str, attachment_data: dict):
                     attachment["agent"] = str(agent_path)
                     del attachment["local"]
 
+    def build_auth_headers(self):
+        """
+        Gets a JWT from the server and creates an authorization header from it
+        """
+        jwt = self.authenticate_with_server()
+        if jwt is not None:
+            auth_headers = {"Authorization": jwt}
+        else:
+            auth_headers = None
+        return auth_headers
+
     def submit(self):
         """Submit a new test job to the server"""
         if self.args.filename == "-":
@@ -500,15 +511,6 @@ def submit(self):
             except FileNotFoundError:
                 sys.exit(f"File not found: {self.args.filename}")
         job_dict = yaml.safe_load(data)
-        jwt = self.authenticate_with_server()
-        if jwt is not None:
-            auth_headers = {"Authorization": jwt}
-        else:
-            if "job_priority" in job_dict:
-                sys.exit(
-                    "Must provide client id and secret key for priority jobs"
-                )
-            auth_headers = None
 
         # Check if agents are available to handle this queue
         # and warn or exit depending on options
@@ -518,15 +520,15 @@ def submit(self):
         attachments_data = self.extract_attachment_data(job_dict)
         if attachments_data is None:
             # submit job, no attachments
-            job_id = self.submit_job_data(job_dict, headers=auth_headers)
+            job_id = self.submit_job_data(job_dict)
         else:
             with tempfile.NamedTemporaryFile(suffix="tar.gz") as archive:
                 archive_path = Path(archive.name)
                 # create attachments archive prior to job submission
                 logger.info("Packing attachments into %s", archive_path)
                 self.pack_attachments(archive_path, attachments_data)
                 # submit job, followed by the submission of the archive
-                job_id = self.submit_job_data(job_dict, headers=auth_headers)
+                job_id = self.submit_job_data(job_dict)
                 try:
                     logger.info("Submitting attachments for %s", job_id)
                     self.submit_job_attachments(job_id, path=archive_path)
@@ -570,44 +572,59 @@ def check_online_agents_available(self, queue: str):
             "Waiting for agents to become available..."
         )
 
-    def submit_job_data(self, data: dict, headers: dict = None):
+    def submit_job_data(self, data: dict):
         """Submit data that was generated or read from a file as a test job"""
-        try:
-            job_id = self.client.submit_job(data, headers=headers)
-        except client.HTTPError as exc:
-            if exc.status == 400:
-                sys.exit(
-                    "The job you submitted contained bad data or "
-                    "bad formatting, or did not specify a "
-                    "job_queue."
-                )
-            if exc.status == 404:
-                sys.exit(
-                    "Received 404 error from server. Are you "
-                    "sure this is a testflinger server?"
-                )
-            if exc.status == 401:
-                sys.exit(
-                    "Received 401 error from server. You are "
-                    "attempting to use a feature that requires "
-                    "client authorisation without using client "
-                    "credentials. See https://testflinger.readthedocs"
-                    ".io/en/latest/how-to/authentication/ for more details"
-                )
-            if exc.status == 403:
-                sys.exit(
-                    "Received 403 error from server with reason "
-                    f"{exc.msg}"
-                    "The specified client credentials do "
-                    "not have sufficient permissions for the resource(s) "
-                    "you are trying to access."
-                )
+        retry_count = 0
+        while True:
+            try:
+                auth_headers = self.build_auth_headers()
+                job_id = self.client.submit_job(data, headers=auth_headers)
+                break
+            except client.HTTPError as exc:
+                if exc.status == 400:
+                    sys.exit(
+                        "The job you submitted contained bad data or "
+                        "bad formatting, or did not specify a "
+                        "job_queue."
+                    )
+                if exc.status == 404:
+                    sys.exit(
+                        "Received 404 error from server. Are you "
+                        "sure this is a testflinger server?"
+                    )
 
-            # This shouldn't happen, so let's get more information
-            sys.exit(
-                "Unexpected error status from testflinger "
-                f"server: [{exc.status}] {exc.msg}"
-            )
+                if exc.status == 403:
+                    sys.exit(
+                        "Received 403 error from server with reason "
+                        f"{exc.msg}"
+                        "The specified client credentials do not have "
+                        "sufficient permissions for the resource(s) "
+                        "you are trying to access."
+                    )
+                if exc.status == 401:
+                    if "expired" in exc.msg:
+                        if retry_count < 2:
+                            retry_count += 1
+                        else:
+                            sys.exit(
+                                "Received 401 error from server due to "
+                                "expired authorization token."
+                            )
+                    else:
+                        sys.exit(
+                            "Received 401 error from server with reason "
+                            f"{exc.msg} You are attempting to use a feature "
+                            "that requires client authorisation "
+                            "without using client credentials. "
+                            "See https://testflinger.readthedocs.io/en/latest"
+                            "/how-to/authentication/ for more details"
+                        )
+                else:
+                    # This shouldn't happen, so let's get more information
+                    sys.exit(
+                        "Unexpected error status from testflinger "
+                        f"server: [{exc.status}] {exc.msg}"
+                    )
         return job_id
 
     def submit_job_attachments(self, job_id: str, path: Path):

cli/testflinger_cli/tests/test_cli.py

@@ -503,11 +503,15 @@ def test_submit_with_priority(tmp_path, requests_mock):
         json=[{"name": "fake_agent", "state": "waiting"}],
     )
     tfcli.submit()
-    assert requests_mock.last_request.headers.get("Authorization") == fake_jwt
+    history = requests_mock.request_history
+    assert len(history) == 3
+    assert history[1].path == "/v1/oauth2/token"
+    assert history[2].path == "/v1/job"
+    assert history[2].headers.get("Authorization") == fake_jwt
 
 
-def test_submit_priority_no_credentials(tmp_path):
-    """Tests priority jobs rejected with no specified credentials"""
+def test_submit_token_timeout_retry(tmp_path, requests_mock):
+    """Tests job submission retries 3 times when token has expired"""
     job_data = {
         "job_queue": "fake",
         "job_priority": 100,
@@ -517,12 +521,30 @@ def test_submit_priority_no_credentials(tmp_path):
 
     sys.argv = ["", "submit", str(job_file)]
     tfcli = testflinger_cli.TestflingerCli()
+    tfcli.client_id = "my_client_id"
+    tfcli.secret_key = "my_secret_key"
+
+    fake_jwt = "my_jwt"
+    requests_mock.post(f"{URL}/v1/oauth2/token", text=fake_jwt)
+    requests_mock.post(
+        f"{URL}/v1/job", text="Token has expired", status_code=401
+    )
+    requests_mock.get(
+        URL + "/v1/queues/fake/agents",
+        json=[{"name": "fake_agent", "state": "waiting"}],
+    )
     with pytest.raises(SystemExit) as exc_info:
         tfcli.submit()
-        assert (
-            "Must provide client id and secret key for priority jobs"
-            in exc_info.value
-        )
+        assert "Token has expired" in exc_info.value
+
+    history = requests_mock.request_history
+    assert len(history) == 7
+    assert history[1].path == "/v1/oauth2/token"
+    assert history[2].path == "/v1/job"
+    assert history[3].path == "/v1/oauth2/token"
+    assert history[4].path == "/v1/job"
+    assert history[5].path == "/v1/oauth2/token"
+    assert history[6].path == "/v1/job"
 
 
 def test_show(capsys, requests_mock):

server/src/api/v1.py

@@ -120,7 +120,7 @@ def decode_jwt_token(auth_token: str, secret_key: str) -> dict:
             options={"require": ["exp", "iat", "sub"]},
         )
     except jwt.exceptions.ExpiredSignatureError:
-        abort(403, "Token has expired")
+        abort(401, "Token has expired")
     except jwt.exceptions.InvalidTokenError:
         abort(403, "Invalid Token")
 
@@ -801,7 +801,7 @@ def generate_token(allowed_resources, secret_key):
     See retrieve_token for more information on the contents of
     the token payload
     """
-    expiration_time = datetime.utcnow() + timedelta(seconds=2)
+    expiration_time = datetime.utcnow() + timedelta(seconds=30)
     token_payload = {
         "exp": expiration_time,
         "iat": datetime.now(timezone.utc),  # Issued at time

server/tests/test_v1_authorization.py

@@ -156,7 +156,7 @@ def test_priority_expired_token(mongo_app_with_permissions):
     job_response = app.post(
         "/v1/job", json=job, headers={"Authorization": token}
     )
-    assert 403 == job_response.status_code
+    assert 401 == job_response.status_code
     assert "Token has expired" in job_response.text