Filter out invalid base64 characters

Author: JukkaL

mypyc/lib-rt/librt_base64.c

@@ -6,6 +6,9 @@
 
 #ifdef MYPYC_EXPERIMENTAL
 
+static PyObject *
+b64decode_handle_invalid(PyObject *out_bytes, char *outbuf, size_t max_out, const char *src, size_t srclen);
+
 #define BASE64_MAXBIN ((PY_SSIZE_T_MAX - 3) / 2)
 
 #define STACK_BUFFER_SIZE 1024
@@ -63,6 +66,12 @@ b64encode(PyObject *self, PyObject *const *args, size_t nargs) {
     return b64encode_internal(args[0]);
 }
 
+static inline int
+is_valid_base64_char(char c) {
+    return ((c >= 'A' && c <= 'Z') | (c >= 'a' && c <= 'z') |
+            (c >= '0' && c <= '9') | (c == '+') | (c == '/') | (c == '='));
+}
+
 static PyObject *
 b64decode_internal(PyObject *arg) {
     const char *src;
@@ -91,6 +100,12 @@ b64decode_internal(PyObject *arg) {
         return PyBytes_FromStringAndSize(NULL, 0);
     }
 
+    // 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])) {
+        srclen_ssz--;
+    }
+
     // Compute an output capacity that's at least 3/4 of input, without overflow:
     // ceil(3/4 * N) == N - floor(N/4)
     size_t srclen = (size_t)srclen_ssz;
@@ -112,14 +127,14 @@ b64decode_internal(PyObject *arg) {
     char *outbuf = PyBytes_AS_STRING(out_bytes);
     size_t outlen = max_out;
 
-    // Decode (flags = 0 for plain input)
     int ret = base64_decode(src, srclen, outbuf, &outlen, 0);
 
     if (ret != 1) {
-        Py_DECREF(out_bytes);
         if (ret == 0) {
-            PyErr_SetString(PyExc_ValueError, "Only base64 data is allowed");
-        } else if (ret == -1) {
+            return b64decode_handle_invalid(out_bytes, outbuf, max_out, src, srclen);
+        }
+        Py_DECREF(out_bytes);
+        if (ret == -1) {
             PyErr_SetString(PyExc_NotImplementedError, "base64 codec not available in this build");
         } else {
             PyErr_SetString(PyExc_RuntimeError, "base64_decode failed");
@@ -149,6 +164,45 @@ b64decode_internal(PyObject *arg) {
 #endif
 }
 
+static PyObject *
+b64decode_handle_invalid(PyObject *out_bytes, char *outbuf, size_t max_out, const char *src, size_t srclen)
+{
+    size_t i;
+    char *newbuf = PyMem_Malloc(srclen);
+    size_t newbuf_len = 0;
+    for (i = 0; i < srclen; i++) {
+        char c = src[i];
+        if (is_valid_base64_char(c)) {
+            newbuf[newbuf_len++] = c;
+        }
+    }
+
+    size_t outlen = max_out;
+    int ret = base64_decode(newbuf, newbuf_len, outbuf, &outlen, 0);
+    PyMem_Free(newbuf);
+
+    if (ret != 1) {
+        Py_DECREF(out_bytes);
+        if (ret == 0) {
+            PyErr_SetString(PyExc_ValueError, "Only base64 data is allowed");
+        }
+        if (ret == -1) {
+            PyErr_SetString(PyExc_NotImplementedError, "base64 codec not available in this build");
+        } else {
+            PyErr_SetString(PyExc_RuntimeError, "base64_decode failed");
+        }
+        return NULL;
+    }
+
+    // Shrink in place to the actual decoded length
+    if (_PyBytes_Resize(&out_bytes, (Py_ssize_t)outlen) < 0) {
+        // _PyBytes_Resize sets an exception and may free the old object
+        return NULL;
+    }
+    return out_bytes;
+}
+
+
 static PyObject*
 b64decode(PyObject *self, PyObject *const *args, size_t nargs) {
     if (nargs != 1) {

mypyc/test-data/run-base64.test

@@ -82,6 +82,17 @@ def test_decode_different_strings() -> None:
     for b in b"", b"ab", b"bac", b"1234", b"xyz88", b"abc" * 200:
         check_decode(b)
 
+def test_decode_with_non_base64_chars() -> None:
+    # For stdlib compatibility, non-base64 characters should be ignored.
+
+    # Invalid characters as a suffix use a fast path.
+    check_decode(b"eA== ", encoded=True)
+    check_decode(b"eA==\n", encoded=True)
+    check_decode(b"eA== \t\n", encoded=True)
+    check_decode(b"\n", encoded=True)
+
+    check_decode(b" e  A = = ", encoded=True)
+
 def test_decode_wrapper() -> None:
     dec: Any = b64decode
     assert dec(b"eA==") == b"x"