decred.org/dcrwallet/v3@v3.1.0/internal/uniformprng/prng.go (about)

     1  // Package uniformprng implements a uniform, cryptographically secure
     2  // pseudo-random number generator.
     3  package uniformprng
     4  
     5  import (
     6  	"encoding/binary"
     7  	"io"
     8  	"math/bits"
     9  
    10  	"golang.org/x/crypto/chacha20"
    11  )
    12  
    13  // Source returns cryptographically-secure pseudorandom numbers with uniform
    14  // distribution.
    15  type Source struct {
    16  	buf    [8]byte
    17  	cipher *chacha20.Cipher
    18  }
    19  
    20  var nonce = make([]byte, chacha20.NonceSize)
    21  
    22  // NewSource seeds a Source from a 32-byte key.
    23  func NewSource(seed *[32]byte) *Source {
    24  	cipher, _ := chacha20.NewUnauthenticatedCipher(seed[:], nonce)
    25  	return &Source{cipher: cipher}
    26  }
    27  
    28  // RandSource creates a Source with seed randomness read from rand.
    29  func RandSource(rand io.Reader) (*Source, error) {
    30  	seed := new([32]byte)
    31  	_, err := io.ReadFull(rand, seed[:])
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  	return NewSource(seed), nil
    36  }
    37  
    38  // Uint32 returns a pseudo-random uint32.
    39  func (s *Source) Uint32() uint32 {
    40  	b := s.buf[:4]
    41  	for i := range b {
    42  		b[i] = 0
    43  	}
    44  	s.cipher.XORKeyStream(b, b)
    45  	return binary.LittleEndian.Uint32(b)
    46  }
    47  
    48  // Uint32n returns a pseudo-random uint32 in range [0,n) without modulo bias.
    49  func (s *Source) Uint32n(n uint32) uint32 {
    50  	if n < 2 {
    51  		return 0
    52  	}
    53  	n--
    54  	mask := ^uint32(0) >> bits.LeadingZeros32(n)
    55  	for {
    56  		u := s.Uint32() & mask
    57  		if u <= n {
    58  			return u
    59  		}
    60  	}
    61  }
    62  
    63  // Int63 returns a pseudo-random 63-bit positive integer as an int64 without
    64  // modulo bias.
    65  func (s *Source) Int63() int64 {
    66  	b := s.buf[:]
    67  	for i := range b {
    68  		b[i] = 0
    69  	}
    70  	s.cipher.XORKeyStream(b, b)
    71  	return int64(binary.LittleEndian.Uint64(b) &^ (1 << 63))
    72  }
    73  
    74  // Int63n returns, as an int64, a pseudo-random 63-bit positive integer in [0,n)
    75  // without modulo bias.
    76  // It panics if n <= 0.
    77  func (s *Source) Int63n(n int64) int64 {
    78  	if n <= 0 {
    79  		panic("invalid argument to Int63n")
    80  	}
    81  	n--
    82  	mask := int64(^uint64(0) >> bits.LeadingZeros64(uint64(n)))
    83  	for {
    84  		i := s.Int63() & mask
    85  		if i <= n {
    86  			return i
    87  		}
    88  	}
    89  }
    90  
    91  // Int63 returns a random non-negative int64, with randomness read from rand.
    92  func Int63(rand io.Reader) (int64, error) {
    93  	buf := make([]byte, 8)
    94  	_, err := io.ReadFull(rand, buf)
    95  	if err != nil {
    96  		return 0, err
    97  	}
    98  	return int64(binary.LittleEndian.Uint64(buf) &^ (1 << 63)), nil
    99  
   100  }
   101  
   102  // Int63n returns, as an int64, a pseudo-random 63-bit positive integer in [0,n)
   103  // without modulo bias.
   104  // Randomness is read from rand.
   105  // It panics if n <= 0.
   106  func Int63n(rand io.Reader, n int64) (int64, error) {
   107  	if n <= 0 {
   108  		panic("invalid argument to Int63n")
   109  	}
   110  	n--
   111  	mask := int64(^uint64(0) >> bits.LeadingZeros64(uint64(n)))
   112  	for {
   113  		v, err := Int63(rand)
   114  		if err != nil {
   115  			return 0, err
   116  		}
   117  		v &= mask
   118  		if v <= n {
   119  			return v, nil
   120  		}
   121  	}
   122  }