From 17e885bf7412b7ca333d57081595ad4d6868152c Mon Sep 17 00:00:00 2001
From: jrconlin <jconlin+git@mozilla.com>
Date: Mon, 9 May 2016 15:33:11 -0700
Subject: [PATCH] bug: Normalize padding handling for restricted subscriptions

Since fernet passes it's content through base64 handling, it's just as
picky about padding as other base64 instances. For now, strip padding
from outbound endpoints and reapply it (if needed) to inbound endpoints.

Closes #466
---
 autopush/settings.py               | 14 +++++++++-----
 autopush/tests/test_endpoint.py    | 27 ++++++++++++++++++++++++++-
 autopush/tests/test_integration.py |  1 -
 autopush/utils.py                  | 12 +++++++++---
 4 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/autopush/settings.py b/autopush/settings.py
index ff38a643..74a63315 100644
--- a/autopush/settings.py
+++ b/autopush/settings.py
@@ -36,7 +36,12 @@
     SimpleRouter,
     WebPushRouter,
 )
-from autopush.utils import canonical_url, resolve_ip, base64url_decode
+from autopush.utils import (
+    canonical_url,
+    resolve_ip,
+    repad,
+    base64url_decode
+)
 from autopush.senderids import SENDERID_EXPRY, DEFAULT_BUCKET
 from autopush.crypto_key import (CryptoKey, CryptoKeyException)
 
@@ -303,8 +308,8 @@ def make_endpoint(self, uaid, chid, key=None):
             return root + 'v1/' + self.fernet.encrypt(base).strip('=')
 
         raw_key = base64url_decode(key.encode('utf8'))
-        return root + 'v2/' + self.fernet.encrypt(base +
-                                                  sha256(raw_key).digest())
+        ep = self.fernet.encrypt(base + sha256(raw_key).digest()).strip('=')
+        return root + 'v2/' + ep
 
     def parse_endpoint(self, token, version="v0", ckey_header=None):
         """Parse an endpoint into component elements of UAID, CHID and optional
@@ -320,8 +325,7 @@ def parse_endpoint(self, token, version="v0", ckey_header=None):
         :returns: a dict containing (uaid=UAID, chid=CHID, public_key=KEY)
 
         """
-
-        token = self.fernet.decrypt(token.encode('utf8'))
+        token = self.fernet.decrypt(repad(token).encode('utf8'))
         public_key = None
         if ckey_header:
             try:
diff --git a/autopush/tests/test_endpoint.py b/autopush/tests/test_endpoint.py
index 1d34b234..ddbfd317 100644
--- a/autopush/tests/test_endpoint.py
+++ b/autopush/tests/test_endpoint.py
@@ -3,6 +3,7 @@
 import time
 import uuid
 import random
+import base64
 
 from hashlib import sha256
 
@@ -178,6 +179,7 @@ def setUp(self, t):
             hostname="localhost",
             statsd_host=None,
         )
+        self.old_fernet = settings.fernet
         self.fernet_mock = settings.fernet = Mock(spec=Fernet)
         self.metrics_mock = settings.metrics = Mock(spec=Metrics)
         self.agent_mock = settings.agent = Mock(spec=Agent)
@@ -775,7 +777,7 @@ def _gen_jwt(self, header, payload):
         sk256p = ecdsa.SigningKey.generate(curve=ecdsa.NIST256p)
         vk = sk256p.get_verifying_key()
         sig = jws.sign(payload, sk256p, algorithm="ES256").strip('=')
-        crypto_key = utils.base64url_encode(vk.to_string())
+        crypto_key = utils.base64url_encode(vk.to_string()).strip('=')
         return (sig, crypto_key)
 
     def test_post_webpush_with_vapid_auth(self):
@@ -1230,6 +1232,29 @@ def test_padding(self):
         eq_(utils.base64url_decode("abc="), "\x69\xb7")
         eq_(utils.base64url_decode("abcd"), "\x69\xb7\x1d")
 
+    def test_v2_padded_fernet_decode(self):
+        # Generate a previously valid URL that would cause the #466 issue
+        self.endpoint.ap_settings.fernet = self.old_fernet
+        # stripped, but valid crypto_key
+        dummy_key = ("BHolMkL36ucQsRe_0KRS70JyHB55H4C5Igv2YQEVNzCILN"
+                     "nedxFHSPtzI4KhzNtN2YPqHe7-mWW6_uvaIc5yEDk")
+        # Generate an endpoint, since the fernet key is not locked during
+        # testing.
+        while True:
+            ep = self.endpoint.ap_settings.make_endpoint(
+                dummy_uaid,
+                dummy_chid,
+                dummy_key)
+            token = ep.split("/")[-1]
+            if len(token) % 4:
+                break
+        reply = self.settings.parse_endpoint(token, "v2",
+                                             "p256ecdsa=" + dummy_key)
+        eq_(reply, {'public_key': base64.urlsafe_b64decode(
+                    utils.repad(dummy_key)),
+                    'chid': dummy_chid.replace('-', ''),
+                    'uaid': dummy_uaid.replace('-', '')})
+
     def test_parse_endpoint(self):
         v0_valid = dummy_uaid + ":" + dummy_chid
         uaid_strip = dummy_uaid.replace('-', '')
diff --git a/autopush/tests/test_integration.py b/autopush/tests/test_integration.py
index 3ec38b63..0958691e 100644
--- a/autopush/tests/test_integration.py
+++ b/autopush/tests/test_integration.py
@@ -1301,7 +1301,6 @@ def test_with_key(self):
         yield client.connect()
         yield client.hello()
         yield client.register(chid=chid, key=pk_hex)
-        # check that the client actually registered the key.
 
         # Send an update with a properly formatted key.
         yield client.send_notification(vapid=vapid)
diff --git a/autopush/utils.py b/autopush/utils.py
index f4c6ddfd..80c5457f 100644
--- a/autopush/utils.py
+++ b/autopush/utils.py
@@ -64,6 +64,14 @@ def base64url_encode(string):
     return base64.urlsafe_b64encode(string).strip('=')
 
 
+def repad(string):
+    """Adds padding to strings for base64 decoding"""
+
+    if len(string) % 4:
+        string = string + '===='[len(string) % 4:]
+    return string
+
+
 def base64url_decode(string):
     """Decodes a Base64 URL-encoded string per RFC 7515.
 
@@ -71,9 +79,7 @@ def base64url_decode(string):
     encoded strings, but Python's ``urlsafe_b64decode`` only accepts padded
     strings.
     """
-    if len(string) % 4:
-        string = string + '===='[len(string) % 4:]
-    return base64.urlsafe_b64decode(string)
+    return base64.urlsafe_b64decode(repad(string))
 
 
 def decipher_public_key(key_data):