From ab14372a6cfcdf45602383824b04972aa11a74ef Mon Sep 17 00:00:00 2001
From: jrconlin <jr+git@mozilla.com>
Date: Tue, 21 Jul 2020 22:38:53 +0000
Subject: [PATCH] bug: fix header_cleanup for aesgcm

* Remove embedded " and = from Encryption and Crypto-Key header values

Closes #1419
---
 autopush/tests/test_integration.py | 45 +++++++++++++++++++++---------
 autopush/utils.py                  |  6 ++--
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/autopush/tests/test_integration.py b/autopush/tests/test_integration.py
index cfbb486d..b4bbdbaf 100644
--- a/autopush/tests/test_integration.py
+++ b/autopush/tests/test_integration.py
@@ -91,7 +91,7 @@ def __init__(self, url, sslcontext=None):
         self.messages = {}
         self.notif_response = None  # type: Optional[HTTPResponse]
         self._crypto_key = """\
-keyid="http://example.org/bob/keys/123;salt="XZwpw6o37R-6qoZjw6KwAw"\
+keyid="http://example.org/bob/keys/123";salt="XZwpw6o37R-6qoZjw6KwAw=="\
 """
         self.sslcontext = sslcontext
         self.headers = {
@@ -679,7 +679,10 @@ def test_basic_delivery(self):
         data = str(uuid.uuid4())
         client = yield self.quick_register()
         result = yield client.send_notification(data=data)
-        assert result["headers"]["encryption"] == client._crypto_key
+        # the following presumes that only `salt` is padded.
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         yield self.shut_down(client)
@@ -694,7 +697,9 @@ def test_topic_basic_delivery(self):
         data = str(uuid.uuid4())
         client = yield self.quick_register()
         result = yield client.send_notification(data=data, topic="Inbox")
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         yield self.shut_down(client)
@@ -710,7 +715,9 @@ def test_topic_replacement_delivery(self):
         yield client.connect()
         yield client.hello()
         result = yield client.get_notification()
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data2)
         assert result["messageType"] == "notification"
         result = yield client.get_notification()
@@ -726,7 +733,9 @@ def test_topic_no_delivery_on_reconnect(self):
         yield client.connect()
         yield client.hello()
         result = yield client.get_notification(timeout=10)
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         yield client.ack(result["channelID"], result["version"])
@@ -746,7 +755,9 @@ def test_basic_delivery_with_vapid(self):
         client = yield self.quick_register()
         vapid_info = _get_vapid()
         result = yield client.send_notification(data=data, vapid=vapid_info)
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         assert self.logs.logged_ci(lambda ci: 'router_key' in ci)
@@ -1013,7 +1024,9 @@ def test_ttl_not_present_connected(self):
         client = yield self.quick_register()
         result = yield client.send_notification(data=data, ttl=None)
         assert result is not None
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         yield self.shut_down(client)
@@ -1024,7 +1037,9 @@ def test_ttl_not_present_connected_no_ack(self):
         client = yield self.quick_register()
         result = yield client.send_notification(data=data, ttl=None)
         assert result is not None
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         yield client.disconnect()
@@ -1040,7 +1055,9 @@ def test_ttl_0_connected(self):
         client = yield self.quick_register()
         result = yield client.send_notification(data=data, ttl=0)
         assert result is not None
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data)
         assert result["messageType"] == "notification"
         yield self.shut_down(client)
@@ -1091,7 +1108,9 @@ def test_ttl_batch_expired_and_good_one(self):
         yield client.hello()
         result = yield client.get_notification(timeout=4)
         assert result is not None
-        assert result["headers"]["encryption"] == client._crypto_key
+        clean_header = client._crypto_key.replace(
+            '"', '').rstrip('=')
+        assert result["headers"]["encryption"] == clean_header
         assert result["data"] == base64url_encode(data2)
         assert result["messageType"] == "notification"
         result = yield client.get_notification()
@@ -2267,9 +2286,9 @@ def test_bad_sends(self):
         ))
         assert response.code == 200
         jbody = json.loads(body)
-        crypto_key = ("keyid=p256dh;dh=BAFJxCIaaWyb4JSkZopERL9MjXBeh3WdBxew"
-                      "SYP0cZWNMJaT7YNaJUiSqBuGUxfRj-9vpTPz5ANmUYq3-u-HWOI")
-        salt = "keyid=p256dh;salt=S82AseB7pAVBJ2143qtM3A"
+        crypto_key = ("keyid=p256dh;dh=\"BAFJxCIaaWyb4JSkZopERL9MjXBeh3WdBxew"
+                      "SYP0cZWNMJaT7YNaJUiSqBuGUxfRj-9vpTPz5ANmUYq3-u-HWOI\"")
+        salt = "keyid=p256dh;salt=\"S82AseB7pAVBJ2143qtM3A==\""
         content_encoding = "aesgcm"
 
         # Test ADMAuth Error
diff --git a/autopush/utils.py b/autopush/utils.py
index fb3b53f0..071ddc60 100644
--- a/autopush/utils.py
+++ b/autopush/utils.py
@@ -375,9 +375,11 @@ def cleanup_headers(self):
         headers = self.headers
         # Strip crypto/encryption headers down
         for hdr in ["crypto-key", "encryption"]:
-            if STRIP_PADDING.search(headers.get(hdr, "")):
+            if hdr in headers:
                 head = headers[hdr].replace('"', '')
-                headers[hdr] = STRIP_PADDING.sub("", head)
+                head = STRIP_PADDING.sub("", head)
+                head.decode('ascii')
+                headers[hdr] = head
 
         # content-encoding header may already be stored as "encoding",
         # this is a failover to ensure that the proper value is pulled in.