[String] allow passing a string of custom characters to ByteString::fromRandom

azjezz · azjezz · commit 9f8878f1c60c · 2020-04-16T19:31:13.000+01:00
diff --git a/src/Symfony/Component/String/ByteString.php b/src/Symfony/Component/String/ByteString.php
@@ -25,20 +25,58 @@
  */
 class ByteString extends AbstractString
 {
+    private const ALPHABET_ALPHANUMERIC = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+
     public function __construct(string $string = '')
     {
         $this->string = $string;
     }
 
-    public static function fromRandom(int $length = 16): self
+    public static function fromRandom(int $length = 16, string $alphabet = null): self
     {
-        $string = '';
+        if (0 === $length) {
+            return new static('');
+        }
+
+        if ($length < 0) {
+            throw new \RuntimeException(sprintf('Expected positive length value, got "%d".', $length));
+        }
+
+        $alphabet = $alphabet ?? self::ALPHABET_ALPHANUMERIC;
+        $alphabet_size = \strlen($alphabet);
+        $bits = (int) ceil(log($alphabet_size, 2.0));
+        if ($bits <= 0 || $bits > 56) {
+            throw new \RuntimeException('Expected $alphabet\'s length to be in [2^1, 2^56].');
+        }
 
-        do {
-            $string .= str_replace(['/', '+', '='], '', base64_encode(random_bytes($length)));
-        } while (\strlen($string) < $length);
+        $ret = '';
+        while ($length > 0) {
+            $urandom_length = (int) ceil(2 * $length * $bits / 8.0);
+            $data = random_bytes($urandom_length);
+            $unpacked_data = 0;
+            $unpacked_bits = 0;
+            for ($i = 0; $i < $urandom_length && $length > 0; ++$i) {
+                // Unpack 8 bits
+                $unpacked_data = ($unpacked_data << 8) | \ord($data[$i]);
+                $unpacked_bits += 8;
+
+                // While we have enough bits to select a character from the alphabet, keep
+                // consuming the random data
+                for (; $unpacked_bits >= $bits && $length > 0; $unpacked_bits -= $bits) {
+                    $index = ($unpacked_data & ((1 << $bits) - 1));
+                    $unpacked_data >>= $bits;
+                    // Unfortunately, the alphabet size is not necessarily a power of two.
+                    // Worst case, it is 2^k + 1, which means we need (k+1) bits and we
+                    // have around a 50% chance of missing as k gets larger
+                    if ($index < $alphabet_size) {
+                        $ret .= $alphabet[$index];
+                        --$length;
+                    }
+                }
+            }
+        }
 
-        return new static(substr($string, 0, $length));
+        return new static($ret);
     }
 
     public function bytesAt(int $offset): array
diff --git a/src/Symfony/Component/String/CHANGELOG.md b/src/Symfony/Component/String/CHANGELOG.md
@@ -12,6 +12,7 @@ CHANGELOG
    depending of the input string UTF-8 compliancy
  * added `$cut` parameter to `Symfony\Component\String\AbstractString::truncate()`
  * added `AbstractString::containsAny()`
+ * allow passing a string of custom characters to ByteString::fromRandom
 
 5.0.0
 -----
diff --git a/src/Symfony/Component/String/Tests/ByteStringTest.php b/src/Symfony/Component/String/Tests/ByteStringTest.php
@@ -12,6 +12,7 @@
 namespace Symfony\Component\String\Tests;
 
 use Symfony\Component\String\AbstractString;
+use function Symfony\Component\String\b;
 use Symfony\Component\String\ByteString;
 
 class ByteStringTest extends AbstractAsciiTestCase
@@ -21,6 +22,47 @@ protected static function createFromString(string $string): AbstractString
         return new ByteString($string);
     }
 
+    public function testFromRandom(): void
+    {
+        $random = ByteString::fromRandom(32);
+
+        self::assertSame(32, $random->length());
+        foreach ($random->chunk() as $char) {
+            self::assertNotNull(b('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789')->indexOf($char));
+        }
+    }
+
+    public function testFromRandomWithSpecificChars(): void
+    {
+        $random = ByteString::fromRandom(32, 'abc');
+
+        self::assertSame(32, $random->length());
+        foreach ($random->chunk() as $char) {
+            self::assertNotNull(b('abc')->indexOf($char));
+        }
+    }
+
+    public function testFromRandomEarlyReturnForZeroLength(): void
+    {
+        self::assertSame('', ByteString::fromRandom(0));
+    }
+
+    public function testFromRandomThrowsForNegativeLength(): void
+    {
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Expected positive length value, got -1');
+
+        ByteString::fromRandom(-1);
+    }
+
+    public function testFromRandomAlphabetMin(): void
+    {
+        $this->expectException(\RuntimeException::class);
+        $this->expectExceptionMessage('Expected $alphabet\'s length to be in [2^1, 2^56]');
+
+        ByteString::fromRandom(32, 'a');
+    }
+
     public static function provideBytesAt(): array
     {
         return array_merge(
