🔒 Make Base64 encoding agnostic to dirty trailing memory bits (#843)

.gas-snapshot

@@ -1,16 +1,16 @@
 Base64Test:testBase64DecodeSentenceGas() (gas: 3663)
 Base64Test:testBase64DecodeShortStringGas() (gas: 919)
-Base64Test:testBase64EncodeDecode(bytes) (runs: 256, μ: 13919, ~: 10697)
-Base64Test:testBase64EncodeDecodeAltModes(bytes) (runs: 256, μ: 831731, ~: 763591)
-Base64Test:testBase64EncodeEmptyString() (gas: 953)
-Base64Test:testBase64EncodeFileSafeAndNoPadding(bytes,bool,bool) (runs: 256, μ: 16096, ~: 13634)
-Base64Test:testBase64EncodeSentence() (gas: 4757)
-Base64Test:testBase64EncodeShortStrings() (gas: 8496)
-Base64Test:testBase64EncodeToStringWithDoublePadding() (gas: 1702)
-Base64Test:testBase64EncodeToStringWithNoPadding() (gas: 1680)
-Base64Test:testBase64EncodeToStringWithSinglePadding() (gas: 1636)
-Base64Test:testBase64WordBoundary() (gas: 12511)
-Base64Test:test__codesize() (gas: 7802)
+Base64Test:testBase64EncodeDecode(bytes) (runs: 256, μ: 13955, ~: 10733)
+Base64Test:testBase64EncodeDecodeAltModes(bytes) (runs: 256, μ: 845900, ~: 760490)
+Base64Test:testBase64EncodeEmptyString() (gas: 1281)
+Base64Test:testBase64EncodeFileSafeAndNoPadding(bytes,bool,bool) (runs: 256, μ: 16144, ~: 13706)
+Base64Test:testBase64EncodeSentence() (gas: 5130)
+Base64Test:testBase64EncodeShortStrings() (gas: 10704)
+Base64Test:testBase64EncodeToStringWithDoublePadding() (gas: 2069)
+Base64Test:testBase64EncodeToStringWithNoPadding() (gas: 2047)
+Base64Test:testBase64EncodeToStringWithSinglePadding() (gas: 2003)
+Base64Test:testBase64WordBoundary() (gas: 14361)
+Base64Test:test__codesize() (gas: 7958)
 CREATE3Test:testDeployERC20() (gas: 761926)
 CREATE3Test:testDeployERC20(bytes32,string,string,uint8) (runs: 256, μ: 805442, ~: 808154)
 CREATE3Test:testDeployedUpperBitsSafeForPlainSolidity() (gas: 657)

src/utils/Base64.sol

@@ -39,6 +39,10 @@ library Base64 {
                 let ptr := add(result, 0x20)
                 let end := add(ptr, encodedLength)
 
+                let dataEnd := add(add(0x20, data), dataLength)
+                let dataEndValue := mload(dataEnd) // Cache the value at the `dataEnd` slot.
+                mstore(dataEnd, 0x00) // Zeroize the `dataEnd` slot to clear dirty bits.
+
                 // Run over the input, 3 bytes at a time.
                 for {} 1 {} {
                     data := add(data, 3) // Advance 3 bytes.
@@ -54,6 +58,7 @@ library Base64 {
                     ptr := add(ptr, 4) // Advance 4 bytes.
                     if iszero(lt(ptr, end)) { break }
                 }
+                mstore(dataEnd, dataEndValue) // Restore the cached value at `dataEnd`.
                 mstore(0x40, add(end, 0x20)) // Allocate the memory.
                 // Equivalent to `o = [0, 2, 1][dataLength % 3]`.
                 let o := div(2, mod(dataLength, 3))

test/Base64.t.sol

@@ -49,7 +49,24 @@ contract Base64Test is SoladyTest {
     }
 
     function _testBase64Encode(string memory input, string memory output) private {
-        assertEq(Base64.encode(bytes(input)), output);
+        assertEq(Base64.encode(_withDirtyEndBits(bytes(input))), output);
+    }
+
+    function _withDirtyEndBits(bytes memory input) private view returns (bytes memory output) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            output := mload(0x40)
+            let n := mload(input)
+
+            mstore(0x00, gas())
+            let r := keccak256(0x00, 0x60)
+            mstore(add(output, add(0x20, n)), r)
+            mstore(add(output, add(0x40, n)), r)
+
+            pop(staticcall(gas(), 4, input, add(0x20, n), output, add(0x20, n)))
+
+            mstore(0x40, add(output, add(0x40, n)))
+        }
     }
 
     function testBase64EncodeDecode(bytes memory input) public {
@@ -75,6 +92,9 @@ contract Base64Test is SoladyTest {
     function testBase64EncodeDecodeAltModes(bytes memory input) public brutalizeMemory {
         for (uint256 i; i < 2; ++i) {
             _misalignFreeMemoryPointer();
+            if (_random() % 2 == 0) {
+                input = _withDirtyEndBits(input);
+            }
             string memory encoded = Base64.encode(input);
             _checkMemory(encoded);