Fix urlsafe_b64decode bug and test more

Author: JukkaL

mypyc/lib-rt/librt_base64.c

@@ -16,9 +16,9 @@ b64decode_handle_invalid_input(
 #define STACK_BUFFER_SIZE 1024
 
 static void
-convert_encoded_to_urlsafe(char *buf, size_t actual_len) {
+convert_encoded_to_urlsafe(char *buf, size_t len) {
     // The loop is written to enable SIMD optimizations
-    for (size_t i = 0; i < actual_len; i++) {
+    for (size_t i = 0; i < len; i++) {
         char ch = buf[i];
         if (ch == '+') {
             ch = '-';
@@ -30,9 +30,9 @@ convert_encoded_to_urlsafe(char *buf, size_t actual_len) {
 }
 
 static void
-convert_urlsafe_to_encoded(const char *src, size_t actual_len, char *buf) {
+convert_urlsafe_to_encoded(const char *src, size_t len, char *buf) {
     // The loop is written to enable SIMD optimizations
-    for (size_t i = 0; i < actual_len; i++) {
+    for (size_t i = 0; i < len; i++) {
         char ch = src[i];
         if (ch == '-') {
             ch = '+';
@@ -144,6 +144,15 @@ b64decode_internal(PyObject *arg, bool urlsafe) {
         return PyBytes_FromStringAndSize(NULL, 0);
     }
 
+    if (urlsafe) {
+        char *new_src = PyMem_Malloc(srclen_ssz + 1);
+        if (new_src == NULL) {
+            return PyErr_NoMemory();
+        }
+        convert_urlsafe_to_encoded(src, srclen_ssz, new_src);
+        src = new_src;
+    }
+
     // 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], true)) {
@@ -162,15 +171,6 @@ b64decode_internal(PyObject *arg, bool urlsafe) {
         return NULL;
     }
 
-    if (urlsafe) {
-        char *new_src = PyMem_Malloc(srclen_ssz);
-        if (new_src == NULL) {
-            return PyErr_NoMemory();
-        }
-        convert_urlsafe_to_encoded(src, srclen_ssz, new_src);
-        src = new_src;
-    }
-
     // Allocate output bytes (uninitialized) of the max capacity
     PyObject *out_bytes = PyBytes_FromStringAndSize(NULL, (Py_ssize_t)max_out);
     if (out_bytes == NULL) {

mypyc/test-data/run-base64.test

@@ -2,11 +2,14 @@
 from typing import Any
 import base64
 import binascii
+import random
 
 from librt.base64 import b64encode, b64decode, urlsafe_b64encode, urlsafe_b64decode
 
 from testutil import assertRaises
 
+rand_values = [random.randbytes(random.randint(1, 2000)) for _ in range(2000)]
+
 def test_encode_basic() -> None:
     assert b64encode(b"x") == b"eA=="
 
@@ -35,6 +38,9 @@ def test_encode_different_strings() -> None:
     for b in b"", b"ab", b"bac", b"1234", b"xyz88", b"abc" * 200:
         check_encode(b)
 
+    for b in rand_values:
+        check_encode(b)
+
 def test_encode_wrapper() -> None:
     enc: Any = b64encode
     assert enc(b"x") == b"eA=="
@@ -84,6 +90,9 @@ def test_decode_different_strings() -> None:
     for b in b"", b"ab", b"bac", b"1234", b"xyz88", b"abc" * 200:
         check_decode(b)
 
+    for b in rand_values:
+        check_decode(b)
+
 def is_base64_char(x: int) -> bool:
     c = chr(x)
     return ('a' <= c <= 'z') or ('A' <= c <= 'Z') or ('0' <= c <= '9') or c in '+/='
@@ -150,18 +159,40 @@ def test_decode_wrapper() -> None:
     with assertRaises(TypeError):
         dec(b"x", b"y")
 
-def check_urlsafe_b64encode(b: bytes) -> None:
+def check_urlsafe_encode(b: bytes) -> None:
     assert urlsafe_b64encode(b) == getattr(base64, "urlsafe_b64encode")(b)
 
 def test_urlsafe_b64encode() -> None:
-    check_urlsafe_b64encode(bytes([x for x in range(256)]))
-
-def check_urlsafe_b64decode(b: bytes) -> None:
+    check_urlsafe_encode(b"")
+    check_urlsafe_encode(b"a")
+    check_urlsafe_encode(b"\xf8")
+    check_urlsafe_encode(b"\xfc")
+    check_urlsafe_encode(b"\xfcx")
+    check_urlsafe_encode(b"\xfcxy")
+    check_urlsafe_encode(b"\xfcxyz")
+    check_urlsafe_encode(bytes([x for x in range(256)]))
+
+    for b in rand_values:
+        check_urlsafe_encode(b)
+
+def check_urlsafe_decode(b: bytes) -> None:
     enc = urlsafe_b64encode(b)
     assert urlsafe_b64decode(enc) == getattr(base64, "urlsafe_b64decode")(enc)
+    enc2 = b64encode(b)
+    assert urlsafe_b64decode(enc2) == getattr(base64, "urlsafe_b64decode")(enc2)
 
 def test_urlsafe_b64decode() -> None:
-    check_urlsafe_b64decode(bytes([x for x in range(256)]))
+    check_urlsafe_decode(b"")
+    check_urlsafe_decode(b"a")
+    check_urlsafe_decode(b"\xf8")
+    check_urlsafe_decode(b"\xfc")
+    check_urlsafe_decode(b"\xfcx")
+    check_urlsafe_decode(b"\xfcxy")
+    check_urlsafe_decode(b"\xfcxyz")
+    check_urlsafe_decode(bytes([x for x in range(256)]))
+
+    for b in rand_values:
+        check_urlsafe_decode(b)
 
 [case testBase64FeaturesNotAvailableInNonExperimentalBuild_librt_base64]
 # This also ensures librt.base64 can be built without experimental features