1// Copyright 2023 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// Package chacha8rand implements a pseudorandom generator
6// based on ChaCha8. It is used by both runtime and math/rand/v2
7// and must have minimal dependencies.
8package chacha8rand
9
10import "internal/byteorder"
11
12const (
13	ctrInc = 4  // increment counter by 4 between block calls
14	ctrMax = 16 // reseed when counter reaches 16
15	chunk  = 32 // each chunk produced by block is 32 uint64s
16	reseed = 4  // reseed with 4 words
17)
18
19// block is the chacha8rand block function.
20func block(seed *[4]uint64, blocks *[32]uint64, counter uint32)
21
22// A State holds the state for a single random generator.
23// It must be used from one goroutine at a time.
24// If used by multiple goroutines at a time, the goroutines
25// may see the same random values, but the code will not
26// crash or cause out-of-bounds memory accesses.
27type State struct {
28	buf  [32]uint64
29	seed [4]uint64
30	i    uint32
31	n    uint32
32	c    uint32
33}
34
35// Next returns the next random value, along with a boolean
36// indicating whether one was available.
37// If one is not available, the caller should call Refill
38// and then repeat the call to Next.
39//
40// Next is //go:nosplit to allow its use in the runtime
41// with per-m data without holding the per-m lock.
42//
43//go:nosplit
44func (s *State) Next() (uint64, bool) {
45	i := s.i
46	if i >= s.n {
47		return 0, false
48	}
49	s.i = i + 1
50	return s.buf[i&31], true // i&31 eliminates bounds check
51}
52
53// Init seeds the State with the given seed value.
54func (s *State) Init(seed [32]byte) {
55	s.Init64([4]uint64{
56		byteorder.LeUint64(seed[0*8:]),
57		byteorder.LeUint64(seed[1*8:]),
58		byteorder.LeUint64(seed[2*8:]),
59		byteorder.LeUint64(seed[3*8:]),
60	})
61}
62
63// Init64 seeds the state with the given seed value.
64func (s *State) Init64(seed [4]uint64) {
65	s.seed = seed
66	block(&s.seed, &s.buf, 0)
67	s.c = 0
68	s.i = 0
69	s.n = chunk
70}
71
72// Refill refills the state with more random values.
73// After a call to Refill, an immediate call to Next will succeed
74// (unless multiple goroutines are incorrectly sharing a state).
75func (s *State) Refill() {
76	s.c += ctrInc
77	if s.c == ctrMax {
78		// Reseed with generated uint64s for forward secrecy.
79		// Normally this is done immediately after computing a block,
80		// but we do it immediately before computing the next block,
81		// to allow a much smaller serialized state (just the seed plus offset).
82		// This gives a delayed benefit for the forward secrecy
83		// (you can reconstruct the recent past given a memory dump),
84		// which we deem acceptable in exchange for the reduced size.
85		s.seed[0] = s.buf[len(s.buf)-reseed+0]
86		s.seed[1] = s.buf[len(s.buf)-reseed+1]
87		s.seed[2] = s.buf[len(s.buf)-reseed+2]
88		s.seed[3] = s.buf[len(s.buf)-reseed+3]
89		s.c = 0
90	}
91	block(&s.seed, &s.buf, s.c)
92	s.i = 0
93	s.n = uint32(len(s.buf))
94	if s.c == ctrMax-ctrInc {
95		s.n = uint32(len(s.buf)) - reseed
96	}
97}
98
99// Reseed reseeds the state with new random values.
100// After a call to Reseed, any previously returned random values
101// have been erased from the memory of the state and cannot be
102// recovered.
103func (s *State) Reseed() {
104	var seed [4]uint64
105	for i := range seed {
106		for {
107			x, ok := s.Next()
108			if ok {
109				seed[i] = x
110				break
111			}
112			s.Refill()
113		}
114	}
115	s.Init64(seed)
116}
117
118// Marshal marshals the state into a byte slice.
119// Marshal and Unmarshal are functions, not methods,
120// so that they will not be linked into the runtime
121// when it uses the State struct, since the runtime
122// does not need these.
123func Marshal(s *State) []byte {
124	data := make([]byte, 6*8)
125	copy(data, "chacha8:")
126	used := (s.c/ctrInc)*chunk + s.i
127	byteorder.BePutUint64(data[1*8:], uint64(used))
128	for i, seed := range s.seed {
129		byteorder.LePutUint64(data[(2+i)*8:], seed)
130	}
131	return data
132}
133
134type errUnmarshalChaCha8 struct{}
135
136func (*errUnmarshalChaCha8) Error() string {
137	return "invalid ChaCha8 encoding"
138}
139
140// Unmarshal unmarshals the state from a byte slice.
141func Unmarshal(s *State, data []byte) error {
142	if len(data) != 6*8 || string(data[:8]) != "chacha8:" {
143		return new(errUnmarshalChaCha8)
144	}
145	used := byteorder.BeUint64(data[1*8:])
146	if used > (ctrMax/ctrInc)*chunk-reseed {
147		return new(errUnmarshalChaCha8)
148	}
149	for i := range s.seed {
150		s.seed[i] = byteorder.LeUint64(data[(2+i)*8:])
151	}
152	s.c = ctrInc * (uint32(used) / chunk)
153	block(&s.seed, &s.buf, s.c)
154	s.i = uint32(used) % chunk
155	s.n = chunk
156	if s.c == ctrMax-ctrInc {
157		s.n = chunk - reseed
158	}
159	return nil
160}
161