| 1 | package testutils |
| 2 | |
| 3 | import ( |
| 4 | "crypto/sha256" |
| 5 | "io" |
| 6 | |
| 7 | "github.com/dustin/go-humanize" |
| 8 | "golang.org/x/crypto/chacha20" |
| 9 | ) |
| 10 | |
| 11 | type randomReader struct { |
| 12 | cipher *chacha20.Cipher |
| 13 | remaining int64 |
| 14 | } |
| 15 | |
| 16 | func (r *randomReader) Read(p []byte) (int, error) { |
| 17 | if r.remaining <= 0 { |
| 18 | return 0, io.EOF |
| 19 | } |
| 20 | n := min(int64(len(p)), r.remaining) |
| 21 | // Generate random bytes directly into the provided buffer |
| 22 | r.cipher.XORKeyStream(p[:n], make([]byte, n)) |
| 23 | r.remaining -= n |
| 24 | return int(n), nil |
| 25 | } |
| 26 | |
| 27 | // DeterministicRandomReader produces specified number of pseudo-random bytes |
| 28 | // from a seed. Size can be specified as a humanize string (e.g., "256KiB", "1MiB"). |
| 29 | func DeterministicRandomReader(sizeStr string, seed string) (io.Reader, error) { |
| 30 | size, err := humanize.ParseBytes(sizeStr) |
| 31 | if err != nil { |
| 32 | return nil, err |
| 33 | } |
| 34 | return DeterministicRandomReaderBytes(int64(size), seed) |
| 35 | } |
| 36 | |
| 37 | // DeterministicRandomReaderBytes produces exactly `size` pseudo-random bytes |
| 38 | // from a seed. Use this when exact byte precision is needed. |
| 39 | func DeterministicRandomReaderBytes(size int64, seed string) (io.Reader, error) { |
| 40 | // Hash the seed string to a 32-byte key for ChaCha20 |
| 41 | key := sha256.Sum256([]byte(seed)) |
| 42 | // Use ChaCha20 for deterministic random bytes |
| 43 | var nonce [chacha20.NonceSize]byte // Zero nonce for simplicity |
| 44 | cipher, err := chacha20.NewUnauthenticatedCipher(key[:chacha20.KeySize], nonce[:]) |
| 45 | if err != nil { |
| 46 | return nil, err |
| 47 | } |
| 48 | return &randomReader{cipher: cipher, remaining: size}, nil |
| 49 | } |