Merge pull request #10 from Stormheg/feature/add-padding-correction

Correct and consider truncated base64 padding valid

Author: Stormheg

src/drf_base64_binaryfield/fields.py

@@ -87,6 +87,7 @@ def from_base64(self, data: str) -> bytes:
 
         Override this method if you want to control how base64 is decoded.
         """
+        data = self._correct_padding(data)
         if self.url_safe:
             # Base64 uses + and /, but these characters are not safe to use in URLs.
             # We allow these characters to be replaced with - and _.
@@ -96,6 +97,16 @@ def from_base64(self, data: str) -> bytes:
             return base64.b64decode(data, altchars=websafe_substitution_chars, validate=True)
         return base64.b64decode(data, validate=True)
 
+    def _correct_padding(self, data: str) -> str:
+        """Correct base64 padding, if necessary.
+
+        Padding could have been truncated, so we add it back to make sure the data is valid base64.
+        """
+        padding = len(data) % 4
+        if padding:
+            data += "=" * (4 - padding)
+        return data
+
     def to_base64(self, data: bytes) -> bytes:
         """Convert binary data to base64.
 

tests/test_fields.py

@@ -26,7 +26,7 @@ def test_base64_binary_field_to_internal_value_happy(input, url_safe, expected_o
 
 @pytest.mark.parametrize(
     "input",
-    ["#(*@)", "aGVsbG8", "asd/sd/", "123098"],
+    ["#(*@)", "hello", "aGVsbG8==", "aGVsbG8==", "aGVsbG8===", "aGVsbG8====", "aGVsbG8====="],
 )
 @pytest.mark.parametrize("url_safe", [True, False])
 def test_base64_binary_field_to_internal_value_invalid(input, url_safe):
@@ -64,6 +64,20 @@ def test_base64_binary_field_to_representation_invalid():
     assert exc_info.value.detail[0].code == "invalid_type"
 
 
+@pytest.mark.parametrize(
+    "input,corrected_input",
+    [
+        ("WcUTqPwauLDnGdb9uCZ0VQ", "WcUTqPwauLDnGdb9uCZ0VQ=="),
+        ("JWmxDeJ8VCtmukFItR683RM", "JWmxDeJ8VCtmukFItR683RM="),
+        ("402aaUKt/TlX41z5tcQLdv3WMvw", "402aaUKt/TlX41z5tcQLdv3WMvw="),
+    ],
+)
+def test_base64_binary_field_corrects_truncated_padding(input, corrected_input):
+    """Check that the field corrects truncated padding to form valid base64."""
+    field = Base64BinaryField()
+    assert field._correct_padding(input) == corrected_input
+
+
 @pytest.mark.parametrize(
     "url_safe",
     [True, False],