新的 arc4random_uniform 实现

2024-12-30 23:30 • 本文约 3381 字，阅读大致需要 7 分钟 | Security

本月初， Robert Clausecker 替换了 FreeBSD 的 arc4random_uniform(3)。

arc4random_uniform(3) 是 arc4random(3) 之上封装的一个生成一个较小范围伪随机数的函数。 arc4random(3) 采用密码学安全的伪随机数生成一个在 32-bit 范围，即 $[0,\ 2^{32}-1]$ 内均匀分布的伪随机整数，此处的随机分布是依靠对称加密算法（目前采用的是 Chacha20）中用于实现加密的伪随机置换（Pseudorandom Permutation）来保证的。

伪随机数满足均匀分布是许多应用场景中需要的。举例来说，如果模拟掷骰子，我们一定是希望骰子六个面出现的概率基本相同。注意到 $2 ^ {32}$ 并不能被 6 整除，直接将 arc4random(3) 取模，或：

arc4random() % 6

得到的并不是我们所希望的 [0, 5] 范围内的均匀分布的随机数，因为 arc4random() 有可能返回 [4294967292, 4294967295] 这几个数，从而使得结果中 0, 1, 2, 3 这四个数出现的概率略微比 4 和 5 高一些。

为了使 0, 1, 2, 3 这几个数字不要出现的概率比其他数字多，我们可以在样本中直接剔除掉开始或结束的这一段产生 0, 1, 2, 3 这几个在样本空间中多出现过一次的数字的原始输出。这可以是 [0, 3]，也可以是 [4294967292, 4294967295]，具体做法是如果 arc4random() 返回了这一区间的数字，则重新生成一个新的伪随机数。

原算法

FreeBSD 此前采用的是 OpenBSD 的算法：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
uint32_t
arc4random_uniform(uint32_t upper_bound)
{
	uint32_t r, min;

	if (upper_bound < 2)
		return 0;

	/* 2**32 % x == (2**32 - x) % x */
	min = -upper_bound % upper_bound;

	/*
	 * This could theoretically loop forever but each retry has
	 * p > 0.5 (worst case, usually far better) of selecting a
	 * number inside the range we need, so it should rarely need
	 * to re-roll.
	 */
	for (;;) {
		r = arc4random();
		if (r >= min)
			break;
	}

	return r % upper_bound;
}

即，排除靠近 0 的这一份总数为 min 的导致采样偏差的样本。

上述代码中 min = -upper_bound % upper_bound 利用了 $k$ 为整数、 $m$ 为正整数时任意数 $x$ 与 $x \pm km$ 同余，即 $x \equiv x \pm km \pmod{m}$ 的性质，将 $2^{32} \mod upper\_bound$ 改为更容易表示的 $(2^{32}-upper\_bound) \mod upper\_bound$ ，即 -upper_bound % upper_bound。

但这样做有一个问题：无论如何我们都要进行至少两次除法运算：第一次是算出希望丢弃的最低阈值 min，第二次是返回结果。更严重的问题是，这两次的除法中的除数均为 upper_bound，这个数未必是二的整数次方幂 ( $2^n$ )。对于计算机硬件来说，计算除数不是二的整数次方幂的除法比较复杂，而乘法或除数为二的整数次方幂的计算则相对要简单许多，因此 CPU 在处理这些计算时更为高效，我们希望能够避免发生这类任意除数除法的计算。

新算法

Daniel Lemire 教授在 2019 年发表了一个新的算法。

考虑我们有一个可以生成 $[0, 2^L-1]$ 范围内均匀分布的伪随机数，希望基于此生成 $[0, s-1]$ 范围内的均匀分布的伪随机数，其中 $0 \leq s \leq 2^L$ 。此处的 $s$ 就是前面算法中的 $upper\_bound$ 。

新算法具体来说是这样的：

首先，使用 $[0, 2^L-1]$ 范围内均匀分布的伪随机数来生成一个随机数 $x$ 。

为了便于理解，不妨将这个数 $x$ 除以 $2^L$ ，这样我们就得到一个在 $[0, 1)$ 范围内、间隔为 $1 \over 2^{L}$ 的均匀分布的定点小数。具体而言，由于 $0 \leq x \leq 2^L-1$ ，因此 ${0 \over 2^L} \leq {x \over 2^L} \leq {{2^L-1} \over {2^L}}$ 。使用一个包含 $2L$ 个比特的的整数，我们就可以将这个定点数表达为其高 $L$ 位作为其整数部分（全0），而其低 $L$ 位作为其小数部分（ $x$ ）。

定点小数的乘法规则与整数乘法无异。因此，在将长度为 $L$ 的整数 $x$ 扩展到 $2L$ 个比特之后，将 $x$ 乘以 $s$ 我们便可以得到一个整数部分取值范围为 $[0, s-1]$ 的数。从函数角度，这样一来，我们便把所有满足 $0 \leq x \times s < 2^L$ 的 $x$ 对应的输出都映射到 0、所有满足 $2^L \leq x \times s < 2 \times 2^L$ 的 $x$ 对应的输出都映射到 1，或更一般地，将所有满足 $i \times 2^L \leq x \times s < (i+1) \times 2^L$ 的 $x$ 值所对应的输出都映射到 $i$ 。

与原算法类似，这样一来依然会出现一些采样偏差，我们需要消除这些偏差。

下文中，将小于或等于 $x$ 的最大整数记作 $\lfloor x \rfloor$ 。将大于或等于 $x$ 的最小整数记作 $\lceil x \rceil$ 。将 $x$ 除以 $y$ 的整数部分即 $\lfloor x / y \rfloor$ 记作 $x \div y$ 。

观察不难发现，对于整数 $a, b, s$ ，对于满足 $b > a > 0$ 并且 $s > 0$ 的情形，若 $b-a$ 能被 $s$ 整除，则在 $[a, b)$ 范围内必然存在 $(b-a) \div s$ 个 $s$ 的倍数。

令 $a = i \times 2^L + (2^L \mod s)$ ， $b = (i+1) \times 2^L$ ，因此 $b-a = 2^L - (2^L \mod s)$ ，故 $b-a$ 一定能被 $s$ 整除，因而在 $[a, b)$ 即 $[i \times 2^L + (2^L \mod s), (i+1) \times 2^L)$ 范围内，存在且仅存在 $2^L \div s$ 个 $s$ 的倍数。

由于我们的目标是在每一个 $i \times 2^L \leq x \times s < (i+1) \times 2^L$ 范围内可能的 $x$ 取值映射到输出 $i$ 上，这个结论告诉我们，只要我们排除了所有位于 $[ i \times 2^L, i \times 2^L + (2^L \mod s))$ 范围内的 $x \times s$ 值，便可以获得一致的样本了。注意到此处大量出现的 $2^L$ ，我们恰好可以把这一判断写作检测该乘积的小数部分上，因为我们可以通过判断 $x \times s \mod 2^L < 2^L \mod s$ 来得到同样的结论。

不过，计算 $2^L \mod s$ 依然要用到除法。由于 $s > 2^L \mod s$ ，我们可以用 $s$ 代替 $2^L \mod s$ 先做一轮判断，即，只在 $x \times s \mod 2^L < s$ 时才去继续进一步测试是否 $x \times s \mod 2^L < 2^L \mod s$ 。

目前 FreeBSD 的 arc4random_uniform 实现如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
/*-
 * SPDX-License-Identifier: 0BSD
 *
 * Copyright (c) Robert Clausecker <fuz@FreeBSD.org>
 * Based on a publication by Daniel Lemire.
 * Public domain where applicable.
 *
 * Daniel Lemire, "Fast Random Integer Generation in an Interval",
 * Association for Computing Machinery, ACM Trans. Model. Comput. Simul.,
 * no. 1, vol. 29, pp. 1--12, New York, NY, USA, January 2019.
 */

#include <stdint.h>
#include <stdlib.h>

uint32_t
arc4random_uniform(uint32_t upper_bound)
{
        uint64_t product;

        /*
         * The paper uses these variable names:
         *
         * L -- log2(UINT32_MAX+1)
         * s -- upper_bound
         * x -- arc4random() return value
         * m -- product
         * l -- (uint32_t)product
         * t -- threshold
         */

        if (upper_bound <= 1)
                return (0);

        product = upper_bound * (uint64_t)arc4random();

        if ((uint32_t)product < upper_bound) {
                uint32_t threshold;

                /* threshold = (2**32 - upper_bound) % upper_bound */
                threshold = -upper_bound % upper_bound;
                while ((uint32_t)product < threshold)
                        product = upper_bound * (uint64_t)arc4random();
        }

        return (product >> 32);
}

以下增加一些注解来方便理解。

第32-33行：对于 upper_bound <= 1 的情形，本函数定义为返回 [0, upper_bound) 的随机数，无论 upper_bound 为 0 或 1，它只能返回 0，因而直接返回 0。

第35-46行为原论文中的算法实现。

我们希望尽量避免除法，因此首先在第35行生成一次随机数并将其乘以 upper_bound，之后做一次粗略判断（第37行），若乘积的低 32-bit 部分超过了 upper_bound，则无需计算 threshold （余数严格小于除数即 upper_bound）就已经知道它符合无偏移采样的条件，因此可以直接到第46行返回结果。

对于乘积低 32-bit 部分小于 upper_bound 的情况，此时乘积的低 32-bit 部分有机会小于 threshold，如果运气不好，后续生成的随机数可能继续出现乘积低 32-bit 部分有机会小于 threshold 的情况，因此我们一次性地在第41行把这个数计算出来。

第40-41行：由于 $2^L \equiv (2^L - upper\_bound) \mod upper\_bound$ ，我们可以把 threshold 写作 -upper_bound % upper_bound。此处不直接计算 $2^{32} \mod n$ 可以避免使用64-bit除法。

第42-43行：接下来，对于乘积低32-bit低于 upper_bound 的情况，再生成一次随机数并计算乘积，直到得到一个符合条件的乘积为止。

最终在第46行：返回乘积的高32-bit。

最后，感谢 Daniel Lemire、Frank 和 Tony Finch 关于新算法中代码快速路径部分的讨论和帮助。