chacha20_test.go

// chacha20_test.go - test ChaCha20 implementation.
// By Ron Charlton, public domain, 2022-08-28.
// $Id: chacha20_test.go,v 1.163 2025-02-12 12:44:54-05 ron Exp $
//
// Requires randIn.dat and randOut.dat files to run to completion.

package chacha20

import (
	"bytes"
	"crypto/cipher"
	crand "crypto/rand"
	"io"
	"log"
	"os"
	"testing"
)

var _ io.Reader = &ChaCha20_ctx{}
var _ cipher.Stream = &ChaCha20_ctx{}

func TestChaCha20(t *testing.T) {
	// IETF test vector for 20 rounds with a zero key and iv. Key length: 32
	// See https://datatracker.ietf.org/doc/html/draft-strombergson-chacha-test-vectors-00
	want := []byte{
		// block 1
		0x76, 0xb8, 0xe0, 0xad, 0xa0, 0xf1, 0x3d, 0x90,
		0x40, 0x5d, 0x6a, 0xe5, 0x53, 0x86, 0xbd, 0x28,
		0xbd, 0xd2, 0x19, 0xb8, 0xa0, 0x8d, 0xed, 0x1a,
		0xa8, 0x36, 0xef, 0xcc, 0x8b, 0x77, 0x0d, 0xc7,
		0xda, 0x41, 0x59, 0x7c, 0x51, 0x57, 0x48, 0x8d,
		0x77, 0x24, 0xe0, 0x3f, 0xb8, 0xd8, 0x4a, 0x37,
		0x6a, 0x43, 0xb8, 0xf4, 0x15, 0x18, 0xa1, 0x1c,
		0xc3, 0x87, 0xb6, 0x69, 0xb2, 0xee, 0x65, 0x86,
		// block 2
		0x9f, 0x07, 0xe7, 0xbe, 0x55, 0x51, 0x38, 0x7a,
		0x98, 0xba, 0x97, 0x7c, 0x73, 0x2d, 0x08, 0x0d,
		0xcb, 0x0f, 0x29, 0xa0, 0x48, 0xe3, 0x65, 0x69,
		0x12, 0xc6, 0x53, 0x3e, 0x32, 0xee, 0x7a, 0xed,
		0x29, 0xb7, 0x21, 0x76, 0x9c, 0xe6, 0x4e, 0x43,
		0xd5, 0x71, 0x33, 0xb0, 0x74, 0xd8, 0x39, 0xd5,
		0x31, 0xed, 0x1f, 0x28, 0x51, 0x0a, 0xfb, 0x45,
		0xac, 0xe1, 0x0a, 0x1f, 0x4b, 0x79, 0x4d, 0x6f,
	}

	key := make([]byte, 32)
	iv := make([]byte, 8)
	blockLen := 64 // constant from chacha20; also in chacha20.go.

	const (
		randInFileName  = "testdata/randIn.dat"  // random message
		randOutFileName = "testdata/randOut.dat" // encrypted random message
	)

	// 'if false {' is for normal testing.
	// ONLY USE true IF YOU CAN RUN WITH A KNOWN-GOOD Encrypt METHOD, LIKE
	// chacha20.go v4.48.
	// 'if true {' generates two files: non-zero plaintext and its ciphertext
	// and exits with t.Fatalf on purpose that says "files x and y were created".
	if false {
		// Create random non-zero randIn and randOut files for checking
		// encryption of non-zero input and output.
		gotRandIn := make([]byte, 300_030) // 4,687.9 blocks)
		n, err := crand.Read(gotRandIn)
		if err != nil || n != len(gotRandIn) {
			t.Fatalf("Error creating randIn data: n: %d, err: %v\n", n, err)
		}
		err = os.WriteFile(randInFileName, gotRandIn, 0644)
		if err == nil {
			gotRandOut := make([]byte, len(gotRandIn))
			ctx := New(key, iv)
			ctx.Encrypt(gotRandIn, gotRandOut)
			err = os.WriteFile(randOutFileName, gotRandOut, 0644)
		}
		if err != nil {
			t.Fatalf("Error creating %s or %s: %v\n",
				randInFileName, randOutFileName, err)
		}
		t.Fatalf("files %s and %s were created",
			randInFileName, randOutFileName)
	}

	// Get the very long plaintext and ciphertext of a known-good encryption
	// into two variables: nonZeroIn and nonZeroOut.
	var err error
	var nonZeroIn, nonZeroOut []byte

	nonZeroIn, err = os.ReadFile(randInFileName)
	if err == nil {
		nonZeroOut, err = os.ReadFile(randOutFileName)
	}
	if err != nil {
		t.Fatalf("Error reading random file input or output, err=%v", err)
	}
	if len(nonZeroOut) != len(nonZeroIn) {
		t.Fatalf("assert: len(nonZeroOut) == len(nonZeroIn) failed")
	}

	// test encrypt with key, iv and input block of all zeroes with IETF data
	got := make([]byte, len(want))
	ctx := New(key, iv)
	ctx.Encrypt(got, got)

	if !bytes.Equal(got, want) {
		for i := 0; i < 2*blockLen; i++ {
			if got[i] != want[i] {
				t.Errorf("Encrypt() No 1: got[%d]=%d, want[%d]=%d", i, got[i], i, want[i])
			}
		}
		t.Errorf("Encrypt() No. 1:\n got %v\nwant %v", got, want)
	}

	// test encrypt with long non-Zero input and expected non-zero output
	gotNonZeroOut := make([]byte, len(nonZeroOut))
	ctx = New(key, iv)
	if n, err := ctx.Encrypt(nonZeroIn, gotNonZeroOut); err != nil ||
		n != len(gotNonZeroOut) {
		t.Errorf("Encrypt() NonZero: n=%d  err=%v", n, err)
	}
	if !bytes.Equal(gotNonZeroOut, nonZeroOut) {
		for i := 0; i < len(nonZeroIn); i++ {
			if gotNonZeroOut[i] != nonZeroOut[i] {
				t.Errorf("Encrypt() NonZero: got[%d]=%d, want[%d]=%d", i, gotNonZeroOut[i], i, nonZeroOut[i])
			}
		}
	}

	// test piecewise encryption (use 'nonZeroIn' as input; expect nonZeroOut
	// as output)
	piecewiseWant := nonZeroOut // randomly generated bytes
	temp := make([]byte, len(nonZeroIn))
	part1Size := 5
	part2Size := len(nonZeroIn) - part1Size
	got = []byte{}
	ctx = New(key, iv)
	ctx.Encrypt(nonZeroIn[:part1Size], temp)
	got = append(got, temp[:part1Size]...)
	ctx.Encrypt(nonZeroIn[part1Size:], temp)
	got = append(got, temp[:part2Size]...)
	if !bytes.Equal(got, piecewiseWant) {
		for i := 0; i < len(nonZeroIn); i++ {
			if got[i] != piecewiseWant[i] {
				t.Errorf("Encrypt() piecewise: got[%d]=%d, want[%d]=%d", i, got[i], i, piecewiseWant[i])
			}
		}
	}

	// Test sequential chunk encryption of long non-Zero input using two halves
	// of data from randIn.dat.
	ctx = New(key, iv)
	half := len(nonZeroIn) / 2
	if n, err := ctx.Encrypt(nonZeroIn[:half], gotNonZeroOut); err != nil ||
		n != half {
		t.Errorf("Encrypt() NonZero first half: n=%d  err=%v", n, err)
	}
	if n, err := ctx.Encrypt(nonZeroIn[half:], gotNonZeroOut[half:]); err != nil ||
		n != half {
		t.Errorf("Encrypt() NonZero second half: n=%d  err=%v", n, err)
	}
	if !bytes.Equal(gotNonZeroOut, nonZeroOut) {
		for i := 0; i < len(nonZeroIn); i++ {
			if gotNonZeroOut[i] != nonZeroOut[i] {
				t.Errorf("Encrypt() NonZero by halves: got[%d]=%d, want[%d]=%d", i, gotNonZeroOut[i], i, nonZeroOut[i])
			}
		}
	}

	// Keystream should yield same result as Encrypt given the same input
	got = make([]byte, len(want))
	ctx = New(key, iv)
	ctx.Keystream(got)
	if !bytes.Equal(got, want) {
		t.Errorf("Keystream():\n got %v\nwant %v", got, want)
	}

	// Seek to block 0 should yield same result with Keystream
	ctx.Seek(0)
	ctx.Keystream(got)
	if !bytes.Equal(got, want) {
		t.Errorf("Keystream() after Seek(0):\n got %v\nwant %v", got, want)
	}

	// Test Seek(1) for correct endian-ness.
	ctx.Seek(1)
	block2 := make([]byte, blockLen)
	ctx.Keystream(block2)
	if !bytes.Equal(block2, want[blockLen:]) {
		t.Errorf("Seek(1) wrong endian-ness:\n got %v\nwant %v", block2, want[blockLen:])
	}

	// Test Read
	ctx.Seek(0)
	n, err := ctx.Read(got)
	if err != nil {
		t.Errorf("Read() err:\n got %v\nwant %v", err, nil)
	}
	if n != len(want) {
		t.Errorf("Read() n:\n got %d\nwant %d", n, len(want))
	}
	if !bytes.Equal(got, want) {
		t.Errorf("Read() after Seek(0):\n got %v\nwant %v", got, want)
	}

	// Test Read for io.EOF when keystream is exhausted, then test for panic.
	got = make([]byte, blockLen)
	ctx.Seek(0xffffffffffffffff)
	n, err = ctx.Read(got)
	if err != io.EOF {
		t.Errorf("Read() EOF test: got %v want %v", err, io.EOF)
	}
	if n != blockLen {
		t.Errorf("Read() return length at EOF: got %d\nwant %d", n, blockLen)
	}
	func() {
		defer func() {
			if r := recover(); r == nil {
				t.Errorf("ChaCha20 did not panic for Read after EOF")
			}
		}()
		ctx.Read(got[:1])
	}()

	// Test chunk processing for proper err and n returns when keystream is
	// exhausted, then test for panic.  Assumes chacha20:blocksPerChunk=100.
	const blocksPerChunk = 100 // MUST MATCH ITS VALUE IN chacha20.go.
	var offsets = []int{-2, -1, 0, 1, 2}
	for k := 0; k < len(offsets); k++ {
		bcOffset := blocksPerChunk + offsets[k]
		got = make([]byte, blockLen*bcOffset)
		ctx.Seek(0 - uint64(bcOffset))
		n, err = ctx.Read(got)
		if err != io.EOF {
			t.Errorf("chunking EOF test: offset %d got %v want %v",
				offsets[k], err, io.EOF)
		}
		if n != blockLen*bcOffset {
			t.Errorf("chunking Read() return length at EOF:\n offset %d got %d, want %d",
				offsets[k], n, blockLen*bcOffset)
		}
		func() {
			defer func() {
				if r := recover(); r == nil {
					t.Errorf("chunking ChaCha20 with offset %d did not panic for Read after EOF",
						offsets[k])
				}
			}()
			ctx.Read(got[:1])
		}()
	}

	// Seek to block 0 should yield same result with XORKeyStream
	ctx.Seek(0)
	got = make([]byte, len(want))
	ctx.XORKeyStream(got, got)
	if !bytes.Equal(got, want) {
		t.Errorf("XORKeyStream() after Seek(0):\n got %v\nwant %v", got, want)
	}

	// Do a simple encrypt/decrypt with random key and iv, and verify
	// encrypt/decrypt are complementary.
	if _, err = crand.Read(key); err != nil {
		log.Fatalf("error from crypto/rand.Read: %v", err)
	}
	if _, err = crand.Read(iv); err != nil {
		log.Fatalf("error from crypto/rand.Read: %v", err)
	}
	// encrypt
	ctx = New(key, iv)
	m := make([]byte, 5_000_000)
	crand.Read(m)
	c := make([]byte, len(m))
	n, err = ctx.Encrypt(m, c)
	if err != nil {
		t.Errorf("enc/dec (encryption): got %v want %v (n=%d)", err, nil, n)
	}
	// decrypt with the same key and iv as used to encrypt
	ctx = New(key, iv)
	got = make([]byte, len(c))
	ctx.Decrypt(c, got)
	if !bytes.Equal(got, m) {
		for i := 0; i < len(m); i++ {
			if got[i] != m[i] {
				t.Errorf("simple enc+dec: got[%d]=%d, m[%d]=%d", i, got[i], i, m[i])
			}
		}
		t.Errorf("simple enc+dec - got[:5]: %v; want[:5]: %v", got[:5], m[:5])
	}

	// Test encrypting a []byte longer than 2^32 (verify use of variable types).
	// This test took 17 seconds to 'go test -race', but it passed.
	// With 'go test' it took about 1.2 seconds and passed.
	/*
		got = make([]byte, 4.4e9)
		ctx.Seek(0)
		ctx.Encrypt(got, got)
	*/
	ctx.Seek(0)
}

// setup for benchmarks:
var key, iv []byte
var ctx *ChaCha20_ctx
var m5e6 = make([]byte, 5_000_000)

func init() {
	_, err := crand.Read(m5e6)
	if err != nil {
		log.Fatalf("error from crypto/rand.Read: %v", err)
	}
	key = make([]byte, 32)
	crand.Read(key)
	iv = make([]byte, 8)
	crand.Read(iv)
	ctx = New(key, iv)
}

func BenchmarkChaCha_8rnds(b *testing.B) {
	b.SetBytes(int64(len(m5e6)))
	ctx.SetRounds(8)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		n, err := ctx.Encrypt(m5e6, m5e6)
		if n != len(m5e6) || err != nil {
			b.Fatalf("8rnds failed: n=%d, err=%v", n, err)
		}
	}
}

func BenchmarkChaCha_12rnds(b *testing.B) {
	b.SetBytes(int64(len(m5e6)))
	ctx.SetRounds(12)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		n, err := ctx.Encrypt(m5e6, m5e6)
		if n != len(m5e6) || err != nil {
			b.Fatalf("12rnds failed: n=%d, err=%v", n, err)
		}
	}
}

func BenchmarkChaCha_20rnds(b *testing.B) {
	b.SetBytes(int64(len(m5e6)))
	ctx.SetRounds(20)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		n, err := ctx.Encrypt(m5e6, m5e6)
		if n != len(m5e6) || err != nil {
			b.Fatalf("20rnds failed: n=%d, err=%v", n, err)
		}
	}
}

func BenchmarkChaCha_Read(b *testing.B) {
	m5e6 := make([]byte, 5_000_001)
	b.SetBytes(int64(len(m5e6)))
	ctx.SetRounds(20)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ctx.Read(m5e6)
	}
}
func BenchmarkChaCha_SmallMemoryRead(b *testing.B) {
	m5e6 := make([]byte, 5_000_001)
	b.SetBytes(int64(len(m5e6)))
	ctx = NewSmallMemory(key, iv)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		ctx.Read(m5e6)
	}
}

// This is the test used by skeeto's chacha-go implementation on GitHub.
func BenchmarkOneBlockXORStream(b *testing.B) {
	var key [32]byte
	var iv [8]byte
	var buf [blockLen]byte
	c := New(key[:], iv[:])
	c.SetRounds(20)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.XORKeyStream(buf[:], buf[:])
	}
}