From 3ab34f9dd4644f441366970ed96d47b7dac194fc Mon Sep 17 00:00:00 2001
From: David Butenhof <dbutenho@redhat.com>
Date: Sat, 8 Jul 2023 08:20:38 -0400
Subject: [PATCH] Gracefully handle connection errors in RELAY API (#3482)

PBENCH-1205

The requests package reports connection failures by exception, and the relay
module wasn't catching those, causing them to be handled "upstream" and result
in internal server errors.

Instead, catch the exceptions locally and report that the URI didn't respond.
---
 lib/pbench/server/api/resources/relay.py  | 20 +++++++--
 lib/pbench/test/unit/server/test_relay.py | 52 +++++++++++++++++++++++
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/lib/pbench/server/api/resources/relay.py b/lib/pbench/server/api/resources/relay.py
index a44c58e83e..e5306631ad 100644
--- a/lib/pbench/server/api/resources/relay.py
+++ b/lib/pbench/server/api/resources/relay.py
@@ -76,7 +76,12 @@ def _identify(self, args: ApiParams, request: Request) -> Intake:
             APIAbort on failure
         """
         uri = args.uri["uri"]
-        response = requests.get(uri, headers={"Accept": "application/json"})
+        try:
+            response = requests.get(uri, headers={"Accept": "application/json"})
+        except Exception as e:
+            raise APIAbort(
+                HTTPStatus.BAD_GATEWAY, f"Unable to connect to manifest URI: {str(e)!r}"
+            )
         if not response.ok:
             raise APIAbort(
                 HTTPStatus.BAD_GATEWAY,
@@ -121,9 +126,16 @@ def _stream(self, intake: Intake, request: Request) -> Access:
         Raises:
             APIAbort on failure
         """
-        response: requests.Response = requests.get(
-            url=intake.uri, stream=True, headers={"Accept": "application/octet-stream"}
-        )
+        try:
+            response: requests.Response = requests.get(
+                url=intake.uri,
+                stream=True,
+                headers={"Accept": "application/octet-stream"},
+            )
+        except Exception as e:
+            raise APIAbort(
+                HTTPStatus.BAD_GATEWAY, f"Unable to connect to results URI: {str(e)!r}"
+            )
         if not response.ok:
             raise APIAbort(
                 response.status_code,
diff --git a/lib/pbench/test/unit/server/test_relay.py b/lib/pbench/test/unit/server/test_relay.py
index ae9affbcd5..ecbdb288e7 100644
--- a/lib/pbench/test/unit/server/test_relay.py
+++ b/lib/pbench/test/unit/server/test_relay.py
@@ -215,6 +215,24 @@ def test_relay_tar_fail(self, client, server_config, pbench_drb_token, tarball):
             "message": "Unable to retrieve relay tarball: 'Not Found'"
         }
 
+    @responses.activate
+    def test_relay_manifest_connection(self, client, server_config, pbench_drb_token):
+        """Verify behavior when the primary relay URI doesn't respond"""
+        responses.add(
+            responses.GET,
+            "https://relay.example.com/uri1",
+            body=ConnectionError("nobody holme"),
+        )
+        response = client.post(
+            self.gen_uri(server_config, "https://relay.example.com/uri1"),
+            headers=self.gen_headers(pbench_drb_token),
+        )
+        assert response.status_code == HTTPStatus.BAD_GATEWAY
+        assert (
+            response.json["message"]
+            == "Unable to connect to manifest URI: 'nobody holme'"
+        )
+
     @responses.activate
     def test_relay_no_manifest(self, client, server_config, pbench_drb_token):
         """Verify behavior when the primary relay URI isn't found"""
@@ -321,3 +339,37 @@ def throw(self, args: ApiParams, request: Request) -> Intake:
         assert (
             response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
         ), f"Unexpected result, {response.text}"
+
+    @responses.activate
+    def test_relay_tarball_connection(
+        self, client, server_config, pbench_drb_token, monkeypatch
+    ):
+        """Verify behavior when we get a connection error reading the tarball"""
+
+        responses.add(
+            responses.GET,
+            "https://relay.example.com/uri1",
+            status=HTTPStatus.OK,
+            json={
+                "name": "tarball.tar.xz",
+                "md5": "anmd5",
+                "access": "private",
+                "metadata": [],
+                "uri": "https://relay.example.com/uri2",
+            },
+        )
+        responses.add(
+            responses.GET,
+            "https://relay.example.com/uri2",
+            body=ConnectionError("leaky wire"),
+        )
+        response = client.post(
+            self.gen_uri(server_config, "https://relay.example.com/uri1"),
+            headers=self.gen_headers(pbench_drb_token),
+        )
+        assert (
+            response.status_code == HTTPStatus.BAD_GATEWAY
+        ), f"Unexpected result, {response.text}"
+        assert (
+            response.json["message"] == "Unable to connect to results URI: 'leaky wire'"
+        )