refactor: ensure X-Exodus-Version works for synthesized responses too

If the logic for setting X-Exodus-Version lives only in origin-response,
it can't work for responses generated directly without going to origin,
e.g. /listing files, /_/cookie, 404 errors and maybe some other things
in the future.

Implement it also in origin-request, in a way which ensures it can't be
forgotten as handler is extended.

exodus_lambda/functions/base.py

@@ -53,6 +53,11 @@ def max_age(self):
     def lambda_version(self):
         return self.conf["lambda_version"]
 
+    def set_lambda_version(self, response):
+        response.setdefault("headers", {})["x-exodus-version"] = [
+            {"key": "X-Exodus-Version", "value": self.lambda_version}
+        ]
+
     def set_cache_control(self, uri, response):
         max_age_pattern_whitelist = [
             ".+/PULP_MANIFEST",

exodus_lambda/functions/origin_request.py

@@ -1,3 +1,4 @@
+import functools
 import json
 import os
 import time
@@ -31,6 +32,7 @@ def __init__(self, conf_file=CONF_FILE):
             ).total_seconds(),
             timer=time.monotonic,
         )
+        self.handler = self.__wrap_version_check(self.handler)
 
     @property
     def db_client(self):
@@ -227,6 +229,24 @@ def handle_listing_request(self, uri):
 
         return {}
 
+    def __wrap_version_check(self, handler):
+        # Decorator wrapping every request to add x-exodus-version on responses
+        # which generated directly without going to origin-response.
+
+        @functools.wraps(handler)
+        def new_handler(event, context):
+            request = event["Records"][0]["cf"]["request"]
+            response = handler(event, context)
+
+            if "status" in response and "x-exodus-query" in (
+                request.get("headers") or {}
+            ):
+                self.set_lambda_version(response)
+
+            return response
+
+        return new_handler
+
     def handler(self, event, context):
         # pylint: disable=unused-argument
 

exodus_lambda/functions/origin_response.py

@@ -34,9 +34,7 @@ def handler(self, event, context):
             ]
 
         if "headers" in request and "x-exodus-query" in request["headers"]:
-            response["headers"]["x-exodus-version"] = [
-                {"key": "X-Exodus-Version", "value": self.lambda_version}
-            ]
+            self.set_lambda_version(response)
 
         try:
             original_uri = request["headers"]["exodus-original-uri"][0][

tests/functions/test_origin_request.py

@@ -499,3 +499,39 @@ def test_origin_request_cookie_uri_without_secret(mocked_boto3_client, caplog):
     assert "Attempting to get secret %s" % arn in caplog.text
     assert "Couldn't load secret %s" % arn in caplog.text
     assert "botocore.exceptions.ClientError: An error occurred" in caplog.text
+
+
+@mock.patch("boto3.client")
+@mock.patch("exodus_lambda.functions.origin_request.cachetools")
+@mock.patch("exodus_lambda.functions.origin_request.datetime")
+def test_origin_request_with_version_check(
+    mocked_datetime, mocked_cache, mocked_boto3_client, caplog
+):
+    req_uri = "/content/dist/rhel/server/7/listing"
+
+    mocked_datetime.now().isoformat.return_value = MOCKED_DT
+    mocked_cache.TTLCache.return_value = {"exodus-config": {}}
+    mocked_boto3_client().query.return_value = {"Items": []}
+
+    event = {
+        "Records": [
+            {
+                "cf": {
+                    "request": {
+                        "uri": req_uri,
+                        "headers": {"x-exodus-query": 1},
+                    }
+                }
+            }
+        ]
+    }
+    request = OriginRequest(conf_file=TEST_CONF).handler(event, context=None)
+    assert request == {
+        "status": "404",
+        "statusDescription": "Not Found",
+        "headers": {
+            "x-exodus-version": [
+                {"key": "X-Exodus-Version", "value": "fake version"}
+            ]
+        },
+    }