Check for invalid padding

Author: JukkaL

mypyc/lib-rt/librt_base64.c

@@ -1,5 +1,6 @@
 #define PY_SSIZE_T_CLEAN
 #include <Python.h>
+#include <stdbool.h>
 #include "librt_base64.h"
 #include "libbase64.h"
 #include "pythoncapi_compat.h"
@@ -67,9 +68,9 @@ b64encode(PyObject *self, PyObject *const *args, size_t nargs) {
 }
 
 static inline int
-is_valid_base64_char(char c) {
+is_valid_base64_char(char c, bool allow_padding) {
     return ((c >= 'A' && c <= 'Z') | (c >= 'a' && c <= 'z') |
-            (c >= '0' && c <= '9') | (c == '+') | (c == '/') | (c == '='));
+            (c >= '0' && c <= '9') | (c == '+') | (c == '/') | (allow_padding && c == '='));
 }
 
 static PyObject *
@@ -102,7 +103,7 @@ b64decode_internal(PyObject *arg) {
 
     // Quickly ignore invalid characters at the end. Other invalid characters
     // are also accepted, but they need a slow path.
-    while (srclen_ssz > 0 && !is_valid_base64_char(src[srclen_ssz - 1])) {
+    while (srclen_ssz > 0 && !is_valid_base64_char(src[srclen_ssz - 1], true)) {
         srclen_ssz--;
     }
 
@@ -172,11 +173,38 @@ b64decode_handle_invalid(PyObject *out_bytes, char *outbuf, size_t max_out, cons
     size_t newbuf_len = 0;
     for (i = 0; i < srclen; i++) {
         char c = src[i];
-        if (is_valid_base64_char(c)) {
+        if (is_valid_base64_char(c, false)) {
             newbuf[newbuf_len++] = c;
+        } else if (c == '=') {
+            // Copy necessary amount of padding
+            int remainder = newbuf_len % 4;
+            if (remainder == 0) {
+                // No padding needed -- ignore padding
+                break;
+            }
+            int numpad = 4 - remainder;
+            // Check that there is at least the required amount padding (CPython ignores
+            // extra padding)
+            while (numpad > 0) {
+                if (i == srclen || src[i] != '=') {
+                    break;
+                }
+                newbuf[newbuf_len++] = '=';
+                i++;
+                numpad--;
+                while (i < srclen && !is_valid_base64_char(src[i], true)) {
+                    i++;
+                }
+            }
+            break;
         }
     }
 
+    if (newbuf_len % 4 != 0) {
+        PyErr_SetString(PyExc_ValueError, "Incorrect padding");
+        return NULL;
+    }
+
     size_t outlen = max_out;
     int ret = base64_decode(newbuf, newbuf_len, outbuf, &outlen, 0);
     PyMem_Free(newbuf);

mypyc/test-data/run-base64.test

@@ -1,6 +1,7 @@
 [case testAllBase64Features_librt_experimental]
 from typing import Any
 import base64
+import binascii
 
 from librt.base64 import b64encode, b64decode
 
@@ -93,6 +94,39 @@ def test_decode_with_non_base64_chars() -> None:
 
     check_decode(b" e  A = = ", encoded=True)
 
+    # Special case: Two different encodings of the same data
+    check_decode(b"eAa=", encoded=True)
+    check_decode(b"eAY=", encoded=True)
+
+def check_decode_error(b: bytes, ignore_stdlib: bool = False) -> None:
+    if not ignore_stdlib:
+        with assertRaises(binascii.Error):
+            getattr(base64, "b64decode")(b)
+
+    # The raised error is different, since librt shouldn't depend on binascii
+    with assertRaises(ValueError):
+        b64decode(b)
+
+def test_decode_with_invalid_padding() -> None:
+    check_decode_error(b"eA")
+    check_decode_error(b"eA=")
+    check_decode_error(b"eHk")
+    check_decode_error(b"eA = ")
+
+    # Here stdlib behavior seems nonsensical, so we don't try to duplicate it
+    check_decode_error(b"eA=a=", ignore_stdlib=True)
+
+def test_decode_with_extra_data_after_padding() -> None:
+    check_decode(b"=", encoded=True)
+    check_decode(b"==", encoded=True)
+    check_decode(b"===", encoded=True)
+    check_decode(b"====", encoded=True)
+    check_decode(b"eA===", encoded=True)
+    check_decode(b"eHk==", encoded=True)
+    check_decode(b"eA==x", encoded=True)
+    check_decode(b"eHk=x", encoded=True)
+    check_decode(b"eA==abc=======efg", encoded=True)
+
 def test_decode_wrapper() -> None:
     dec: Any = b64decode
     assert dec(b"eA==") == b"x"