rng.go raw

   1  // Copyright 2023 The gVisor Authors.
   2  //
   3  // Licensed under the Apache License, Version 2.0 (the "License");
   4  // you may not use this file except in compliance with the License.
   5  // You may obtain a copy of the License at
   6  //
   7  //     http://www.apache.org/licenses/LICENSE-2.0
   8  //
   9  // Unless required by applicable law or agreed to in writing, software
  10  // distributed under the License is distributed on an "AS IS" BASIS,
  11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12  // See the License for the specific language governing permissions and
  13  // limitations under the License.
  14  
  15  // Package rand implements a cryptographically secure pseudorandom number
  16  // generator.
  17  package rand
  18  
  19  import (
  20  	"encoding/binary"
  21  	"fmt"
  22  	"io"
  23  )
  24  
  25  // RNG exposes convenience functions based on a cryptographically secure
  26  // io.Reader.
  27  type RNG struct {
  28  	Reader io.Reader
  29  }
  30  
  31  // RNGFrom returns a new RNG. r must be a cryptographically secure io.Reader.
  32  func RNGFrom(r io.Reader) RNG {
  33  	return RNG{Reader: r}
  34  }
  35  
  36  // Uint16 is analogous to the standard library's math/rand.Uint16.
  37  func (rg *RNG) Uint16() uint16 {
  38  	var data [2]byte
  39  	if _, err := rg.Reader.Read(data[:]); err != nil {
  40  		panic(fmt.Sprintf("Read() failed: %v", err))
  41  	}
  42  	return binary.NativeEndian.Uint16(data[:])
  43  }
  44  
  45  // Uint32 is analogous to the standard library's math/rand.Uint32.
  46  func (rg *RNG) Uint32() uint32 {
  47  	var data [4]byte
  48  	if _, err := rg.Reader.Read(data[:]); err != nil {
  49  		panic(fmt.Sprintf("Read() failed: %v", err))
  50  	}
  51  	return binary.NativeEndian.Uint32(data[:])
  52  }
  53  
  54  // Int63n is analogous to the standard library's math/rand.Int63n.
  55  func (rg *RNG) Int63n(n int64) int64 {
  56  	// Based on Go's rand package implementation, but using
  57  	// cryptographically secure random numbers.
  58  	if n <= 0 {
  59  		panic(fmt.Sprintf("n must be positive, but got %d", n))
  60  	}
  61  
  62  	// This can be done quickly when n is a power of 2.
  63  	if n&(n-1) == 0 {
  64  		return int64(rg.Uint64()) & (n - 1)
  65  	}
  66  
  67  	// The naive approach would be to return rg.Int63()%n, but we need the
  68  	// random number to be fair. It shouldn't be biased towards certain
  69  	// results, but simple modular math can be very biased. For example, if
  70  	// n is 40% of the maximum int64, then the output values of rg.Int63
  71  	// map to return values as follows:
  72  	//
  73  	//  - The first 40% of values map to themselves.
  74  	//  - The second 40% map to themselves - maximum int64.
  75  	//  - The remaining 20% map to the themselves - 2 * (maximum int64),
  76  	//    i.e. the first half of possible output values.
  77  	//
  78  	// And thus 60% of results map the first half of possible output
  79  	// values, and 40% map the second half. Oops!
  80  	//
  81  	// We use the same trick as Go to deal with this: shave off the last
  82  	// segment (the 20% in our example) to make the RNG more fair.
  83  	//
  84  	// In the worst case, n is just over half of maximum int64, meaning
  85  	// that the upper half of rg.Int63 return values are bad. So each call
  86  	// to rg.Int63 has, at worst, a 50% chance of needing a retry.
  87  	maximum := int64((1 << 63) - 1 - (1<<63)%uint64(n))
  88  	ret := rg.Int63()
  89  	for ret > maximum {
  90  		ret = rg.Int63()
  91  	}
  92  	return ret % n
  93  }
  94  
  95  // Int63 is analogous to the standard library's math/rand.Int63.
  96  func (rg *RNG) Int63() int64 {
  97  	return ((1 << 63) - 1) & int64(rg.Uint64())
  98  }
  99  
 100  // Uint64 is analogous to the standard library's math/rand.Uint64.
 101  func (rg *RNG) Uint64() uint64 {
 102  	var data [8]byte
 103  	if _, err := rg.Reader.Read(data[:]); err != nil {
 104  		panic(fmt.Sprintf("Read() failed: %v", err))
 105  	}
 106  	return binary.NativeEndian.Uint64(data[:])
 107  }
 108  
 109  // Uint32 is analogous to the standard library's math/rand.Uint32.
 110  func Uint32() uint32 {
 111  	rng := RNG{Reader: Reader}
 112  	return rng.Uint32()
 113  }
 114  
 115  // Int63n is analogous to the standard library's math/rand.Int63n.
 116  func Int63n(n int64) int64 {
 117  	rng := RNG{Reader: Reader}
 118  	return rng.Int63n(n)
 119  }
 120  
 121  // Int63 is analogous to the standard library's math/rand.Int63.
 122  func Int63() int64 {
 123  	rng := RNG{Reader: Reader}
 124  	return rng.Int63()
 125  }
 126  
 127  // Uint64 is analogous to the standard library's math/rand.Uint64.
 128  func Uint64() uint64 {
 129  	rng := RNG{Reader: Reader}
 130  	return rng.Uint64()
 131  }
 132