frand.go raw

   1  package frand
   2  
   3  import (
   4  	crand "crypto/rand"
   5  	"encoding/binary"
   6  	"math"
   7  	"math/big"
   8  	"math/rand/v2"
   9  	"strconv"
  10  	"sync"
  11  )
  12  
  13  // An RNG is a cryptographically-strong RNG constructed from the ChaCha stream
  14  // cipher.
  15  type RNG struct {
  16  	r *rand.ChaCha8
  17  }
  18  
  19  // Read fills b with random data. It always returns len(b), nil.
  20  func (r *RNG) Read(b []byte) (int, error) {
  21  	return r.r.Read(b)
  22  }
  23  
  24  // Bytes is a helper function that allocates and returns n bytes of random data.
  25  func (r *RNG) Bytes(n int) []byte {
  26  	b := make([]byte, n)
  27  	r.Read(b)
  28  	return b
  29  }
  30  
  31  // Entropy256 is a helper function that returns 256 bits of random data.
  32  func (r *RNG) Entropy256() (entropy [32]byte) {
  33  	r.Read(entropy[:])
  34  	return
  35  }
  36  
  37  // Entropy192 is a helper function that returns 192 bits of random data.
  38  func (r *RNG) Entropy192() (entropy [24]byte) {
  39  	r.Read(entropy[:])
  40  	return
  41  }
  42  
  43  // Entropy128 is a helper function that returns 128 bits of random data.
  44  func (r *RNG) Entropy128() (entropy [16]byte) {
  45  	r.Read(entropy[:])
  46  	return
  47  }
  48  
  49  // Uint64n returns a uniform random uint64 in [0,n). It panics if n == 0.
  50  func (r *RNG) Uint64n(n uint64) uint64 {
  51  	if n == 0 {
  52  		panic("frand: argument to Uint64n is 0")
  53  	}
  54  	// To eliminate modulo bias, keep selecting at random until we fall within
  55  	// a range that is evenly divisible by n.
  56  	// NOTE: since n is at most math.MaxUint64, max is minimized when:
  57  	//    n = math.MaxUint64/2 + 1 -> max = math.MaxUint64 - math.MaxUint64/2
  58  	// This gives an expected 2 tries before choosing a value < max.
  59  	max := math.MaxUint64 - math.MaxUint64%n
  60  again:
  61  	i := r.r.Uint64()
  62  	if i >= max {
  63  		goto again
  64  	}
  65  	return i % n
  66  }
  67  
  68  // Intn returns a uniform random int in [0,n). It panics if n <= 0.
  69  func (r *RNG) Intn(n int) int {
  70  	if n <= 0 {
  71  		panic("frand: argument to Intn is <= 0: " + strconv.Itoa(n))
  72  	}
  73  	// NOTE: since n is at most math.MaxUint64/2, max is minimized when:
  74  	//    n = math.MaxUint64/4 + 1 -> max = math.MaxUint64 - math.MaxUint64/4
  75  	// This gives an expected 1.333 tries before choosing a value < max.
  76  	return int(r.Uint64n(uint64(n)))
  77  }
  78  
  79  // BigIntn returns a uniform random *big.Int in [0,n). It panics if n <= 0.
  80  func (r *RNG) BigIntn(n *big.Int) *big.Int {
  81  	i, _ := crand.Int(r, n)
  82  	return i
  83  }
  84  
  85  // Float64 returns a random float64 in [0,1).
  86  func (r *RNG) Float64() float64 {
  87  	return float64(r.Uint64n(1<<53)) / (1 << 53)
  88  }
  89  
  90  // Perm returns a random permutation of the integers [0,n). It panics if n < 0.
  91  func (r *RNG) Perm(n int) []int {
  92  	m := make([]int, n)
  93  	for i := 1; i < n; i++ {
  94  		j := r.Intn(i + 1)
  95  		m[i] = m[j]
  96  		m[j] = i
  97  	}
  98  	return m
  99  }
 100  
 101  // Shuffle randomly permutes n elements by repeatedly calling swap in the range
 102  // [0,n). It panics if n < 0.
 103  func (r *RNG) Shuffle(n int, swap func(i, j int)) {
 104  	for i := n - 1; i > 0; i-- {
 105  		swap(i, r.Intn(i+1))
 106  	}
 107  }
 108  
 109  // NewCustom returns a new RNG instance seeded with the provided entropy. It
 110  // panics if len(seed) != 32. The bufsize and rounds parameters are ignored.
 111  func NewCustom(seed []byte, bufsize int, rounds int) *RNG {
 112  	if len(seed) != 32 {
 113  		panic("frand: invalid seed size")
 114  	}
 115  	return &RNG{rand.NewChaCha8(([32]byte)(seed))}
 116  }
 117  
 118  // "master" RNG, seeded from crypto/rand; RNGs returned by New derive their seed
 119  // from this RNG. This means we only ever need to read system entropy a single
 120  // time, at startup.
 121  var masterRNG = func() *RNG {
 122  	var seed [32]byte
 123  	if _, err := crand.Read(seed[:]); err != nil {
 124  		panic("not enough system entropy to seed master RNG")
 125  	}
 126  	return &RNG{rand.NewChaCha8(seed)}
 127  }()
 128  var masterMu sync.Mutex
 129  
 130  // New returns a new RNG instance. The instance is seeded with entropy from
 131  // crypto/rand, albeit indirectly; its seed is generated by a global "master"
 132  // RNG, which itself is seeded from crypto/rand. This means the frand package
 133  // only reads system entropy once, at startup.
 134  func New() *RNG {
 135  	masterMu.Lock()
 136  	defer masterMu.Unlock()
 137  	return &RNG{rand.NewChaCha8(masterRNG.Entropy256())}
 138  }
 139  
 140  // Global versions of each RNG method, leveraging a pool of RNGs.
 141  
 142  var rngpool = sync.Pool{
 143  	New: func() interface{} {
 144  		return New()
 145  	},
 146  }
 147  
 148  // Read fills b with random data. It always returns len(b), nil.
 149  func Read(b []byte) (int, error) {
 150  	r := rngpool.Get().(*RNG)
 151  	defer rngpool.Put(r)
 152  	return r.Read(b)
 153  }
 154  
 155  // Bytes is a helper function that allocates and returns n bytes of random data.
 156  func Bytes(n int) []byte {
 157  	r := rngpool.Get().(*RNG)
 158  	defer rngpool.Put(r)
 159  	return r.Bytes(n)
 160  }
 161  
 162  // Entropy256 is a helper function that returns 256 bits of random data.
 163  func Entropy256() [32]byte {
 164  	r := rngpool.Get().(*RNG)
 165  	defer rngpool.Put(r)
 166  	return r.Entropy256()
 167  }
 168  
 169  // Entropy192 is a helper function that returns 192 bits of random data.
 170  func Entropy192() [24]byte {
 171  	r := rngpool.Get().(*RNG)
 172  	defer rngpool.Put(r)
 173  	return r.Entropy192()
 174  }
 175  
 176  // Entropy128 is a helper function that returns 128 bits of random data.
 177  func Entropy128() [16]byte {
 178  	r := rngpool.Get().(*RNG)
 179  	defer rngpool.Put(r)
 180  	return r.Entropy128()
 181  }
 182  
 183  // Uint64n returns a uniform random uint64 in [0,n). It panics if n == 0.
 184  func Uint64n(n uint64) uint64 {
 185  	r := rngpool.Get().(*RNG)
 186  	defer rngpool.Put(r)
 187  	return r.Uint64n(n)
 188  }
 189  
 190  // Intn returns a uniform random int in [0,n). It panics if n <= 0.
 191  func Intn(n int) int {
 192  	r := rngpool.Get().(*RNG)
 193  	defer rngpool.Put(r)
 194  	return r.Intn(n)
 195  }
 196  
 197  // BigIntn returns a uniform random *big.Int in [0,n). It panics if n <= 0.
 198  func BigIntn(n *big.Int) *big.Int {
 199  	r := rngpool.Get().(*RNG)
 200  	defer rngpool.Put(r)
 201  	return r.BigIntn(n)
 202  }
 203  
 204  // Float64 returns a random float64 in [0,1).
 205  func Float64() float64 {
 206  	r := rngpool.Get().(*RNG)
 207  	defer rngpool.Put(r)
 208  	return r.Float64()
 209  }
 210  
 211  // Perm returns a random permutation of the integers [0,n). It panics if n < 0.
 212  func Perm(n int) []int {
 213  	r := rngpool.Get().(*RNG)
 214  	defer rngpool.Put(r)
 215  	return r.Perm(n)
 216  }
 217  
 218  // Shuffle randomly permutes n elements by repeatedly calling swap in the range
 219  // [0,n). It panics if n < 0.
 220  func Shuffle(n int, swap func(i, j int)) {
 221  	r := rngpool.Get().(*RNG)
 222  	defer rngpool.Put(r)
 223  	r.Shuffle(n, swap)
 224  }
 225  
 226  // Reader is a global, shared instance of a cryptographically strong pseudo-
 227  // random generator. Reader is safe for concurrent use by multiple goroutines.
 228  var Reader rngReader
 229  
 230  type rngReader struct{}
 231  
 232  func (rngReader) Read(b []byte) (int, error) { return Read(b) }
 233  
 234  // A Source is a math/rand-compatible source of entropy. It is safe for
 235  // concurrent use by multiple goroutines.
 236  type Source struct {
 237  	rng *RNG
 238  	mu  sync.Mutex
 239  }
 240  
 241  // Seed uses the provided seed value to initialize the Source to a
 242  // deterministic state.
 243  func (s *Source) Seed(i int64) {
 244  	s.mu.Lock()
 245  	defer s.mu.Unlock()
 246  	var seed [32]byte
 247  	binary.LittleEndian.PutUint64(seed[:], uint64(i))
 248  	s.rng.r.Seed(seed)
 249  }
 250  
 251  // Int63 returns a non-negative random 63-bit integer as an int64.
 252  func (s *Source) Int63() int64 {
 253  	s.mu.Lock()
 254  	defer s.mu.Unlock()
 255  	return int64(s.rng.Uint64n(1 << 63))
 256  }
 257  
 258  // Uint64 returns a random 64-bit integer.
 259  func (s *Source) Uint64() uint64 {
 260  	s.mu.Lock()
 261  	defer s.mu.Unlock()
 262  	return s.rng.Uint64n(math.MaxUint64)
 263  }
 264  
 265  // NewSource returns a source for use with the math/rand package.
 266  func NewSource() *Source { return &Source{rng: New()} }
 267