From 03d29f0fde139ea6d7485cf70b3d0ffc094a5119 Mon Sep 17 00:00:00 2001
From: Nathan Gillett <ngillett@redhat.com>
Date: Fri, 27 Mar 2020 15:20:27 -0400
Subject: [PATCH] Support sha-256 want-digest request headers [DELIVERY-8988]

This commit provides support in origin_response for requests with
want-digest headers for the sha256 algorithm.
---
 .../origin_response/origin_response.py        |  6 +++
 tests/functions/test_origin_response.py       | 43 +++++++++++++------
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/exodus_lambda/functions/origin_response/origin_response.py b/exodus_lambda/functions/origin_response/origin_response.py
index c146c06b..7b61d3c3 100755
--- a/exodus_lambda/functions/origin_response/origin_response.py
+++ b/exodus_lambda/functions/origin_response/origin_response.py
@@ -24,6 +24,12 @@ def handler(self, event, context):
         request = event["Records"][0]["cf"]["request"]
         response = event["Records"][0]["cf"]["response"]
 
+        # TODO: Investigate cf-s3 response errors, ensure no digest if error
+        if "headers" in request and "want-digest" in request["headers"]:
+            response["headers"]["digest"] = [
+                {"key": "Digest", "value": f"sha-256={request['uri']}",}
+            ]
+
         max_age = self.conf["headers"]["max_age"]
         max_age_pattern_whitelist = [
             ".+/PULP_MANIFEST",
diff --git a/tests/functions/test_origin_response.py b/tests/functions/test_origin_response.py
index 8a92cef0..fb809671 100644
--- a/tests/functions/test_origin_response.py
+++ b/tests/functions/test_origin_response.py
@@ -13,30 +13,29 @@
 
 
 @pytest.mark.parametrize(
-    "test_input",
+    "original_uri, want_digest",
     [
-        "/some/repo/listing",
-        "/some/repo/repodata/repomd.xml",
-        "/some/repo/ostree/repo/refs/heads/ok/test",
+        ("/some/repo/listing", True),
+        ("/some/repo/repodata/repomd.xml", True),
+        ("/some/repo/ostree/repo/refs/heads/ok/test", False),
     ],
 )
-def test_origin_response_valid_headers(
-    test_input,
-    expected=[{"key": "Cache-Control", "value": "max-age={}".format(max_age)}],
-):
+def test_origin_response_valid_headers(original_uri, want_digest):
     event = {
         "Records": [
             {
                 "cf": {
                     "request": {
+                        "uri": "be7f3007df3e51fb48fff57da9c01c52e6b8e60eceacab"
+                        "7aaf0e05b57578493a",
                         "headers": {
                             "exodus-original-uri": [
                                 {
                                     "key": "exodus-original-uri",
-                                    "value": test_input,
+                                    "value": original_uri,
                                 }
                             ]
-                        }
+                        },
                     },
                     "response": {"headers": {}},
                 }
@@ -44,15 +43,33 @@ def test_origin_response_valid_headers(
         ]
     }
 
+    expected_headers = {
+        "cache-control": [
+            {"key": "Cache-Control", "value": f"max-age={max_age}"}
+        ]
+    }
+
+    if want_digest:
+        event["Records"][0]["cf"]["request"]["headers"]["want-digest"] = [
+            {"key": "Want-Digest", "value": "sha-256"}
+        ]
+        expected_headers["digest"] = [
+            {
+                "key": "Digest",
+                "value": "sha-256=be7f3007df3e51fb48fff57da9c01c52e6b8e60eceac"
+                "ab7aaf0e05b57578493a",
+            }
+        ]
+
     response = LambdaClient(conf_file=CONF_PATH).handler(event, context=None)
-    assert response["headers"]["cache-control"] == expected
+    assert response["headers"] == expected_headers
 
 
 @pytest.mark.parametrize(
     "test_input",
     ["/some/repo/some-rpm.rpm", "/some/repo/ostree/repo/refs/heads/nope"],
 )
-def test_origin_response_empty_headers(test_input, expected={}):
+def test_origin_response_empty_headers(test_input):
     event = {
         "Records": [
             {"cf": {"request": {"headers": {}}, "response": {"headers": {}}}}
@@ -60,7 +77,7 @@ def test_origin_response_empty_headers(test_input, expected={}):
     }
 
     response = LambdaClient(conf_file=CONF_PATH).handler(event, context=None)
-    assert response["headers"] == expected
+    assert response["headers"] == {}
 
 
 def test_origin_response_missing_headers():