nebula/noiseutil/cipher_state_test.go

package noiseutil

import (
	"math"
	"testing"

	"github.com/flynn/noise"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestCipherStateAESGCMRoundtrip(t *testing.T) {
	enc, dec := buildCipherStates(t, CipherAESGCM)
	roundtrip(t, NewCipherStateAESGCM(enc), NewCipherStateAESGCM(dec))
}

func TestCipherStateChaChaPolyRoundtrip(t *testing.T) {
	enc, dec := buildCipherStates(t, noise.CipherChaChaPoly)
	roundtrip(t, NewCipherStateChaChaPoly(enc), NewCipherStateChaChaPoly(dec))
}

func TestNewCipherStateDispatch(t *testing.T) {
	encA, _ := buildCipherStates(t, CipherAESGCM)
	encC, _ := buildCipherStates(t, noise.CipherChaChaPoly)

	assert.IsType(t, &CipherStateAESGCM{}, NewCipherState(encA, CipherAESGCM))
	assert.IsType(t, &CipherStateChaChaPoly{}, NewCipherState(encC, noise.CipherChaChaPoly))
}

func TestNewCipherStateUnsupportedPanics(t *testing.T) {
	enc, _ := buildCipherStates(t, CipherAESGCM)
	assert.Panics(t, func() {
		NewCipherState(enc, fakeCipher{})
	})
}

type fakeCipher struct{}

func (fakeCipher) Cipher(k [32]byte) noise.Cipher { return nil }
func (fakeCipher) CipherName() string             { return "Fake" }

// buildCipherStates runs an in-memory NN handshake with the requested cipher
// to produce a pair of post-handshake CipherStates that share keys.
func buildCipherStates(t *testing.T, c noise.CipherFunc) (*noise.CipherState, *noise.CipherState) {
	t.Helper()
	suite := noise.NewCipherSuite(noise.DH25519, c, noise.HashSHA256)
	cfg := noise.Config{CipherSuite: suite, Pattern: noise.HandshakeNN}
	cfg.Initiator = true
	hsI, err := noise.NewHandshakeState(cfg)
	require.NoError(t, err)
	cfg.Initiator = false
	hsR, err := noise.NewHandshakeState(cfg)
	require.NoError(t, err)

	msg, _, _, err := hsI.WriteMessage(nil, nil)
	require.NoError(t, err)
	_, _, _, err = hsR.ReadMessage(nil, msg)
	require.NoError(t, err)

	msg, dR, _, err := hsR.WriteMessage(nil, nil)
	require.NoError(t, err)
	_, eI, _, err := hsI.ReadMessage(nil, msg)
	require.NoError(t, err)
	require.NotNil(t, eI)
	require.NotNil(t, dR)

	// noise returns (cs1, cs2) where cs1 is the initiator->responder cipher.
	return eI, dR
}

func roundtrip(t *testing.T, enc, dec CipherState) {
	t.Helper()
	plaintext := []byte("nebula cipher state roundtrip")
	ad := []byte("aad")
	nb := make([]byte, 12)

	ct, err := enc.EncryptDanger(nil, ad, plaintext, 1, nb)
	require.NoError(t, err)
	assert.NotEqual(t, plaintext, ct)

	pt, err := dec.DecryptDanger(nil, ad, ct, 1, nb)
	require.NoError(t, err)
	assert.Equal(t, plaintext, pt)

	// Wrong nonce must fail authentication.
	_, err = dec.DecryptDanger(nil, ad, ct, 2, nb)
	require.Error(t, err)

	assert.Equal(t, enc.Overhead(), dec.Overhead())
	assert.Equal(t, 16, enc.Overhead())
}

func TestEncryptRejectsExhaustedCounter(t *testing.T) {
	// Pin the headroom below the uint64 wrap so a typo can't silently move the ceiling.
	require.Equal(t, uint64(1)<<40, RejectHeadroom)
	require.Equal(t, math.MaxUint64-RejectHeadroom, RejectAfterMessages)

	encA, _ := buildCipherStates(t, CipherAESGCM)
	encC, _ := buildCipherStates(t, noise.CipherChaChaPoly)
	nb := make([]byte, 12)

	for _, cs := range []CipherState{NewCipherStateAESGCM(encA), NewCipherStateChaChaPoly(encC)} {
		_, err := cs.EncryptDanger(nil, nil, []byte("x"), RejectAfterMessages-1, nb)
		require.NoError(t, err)

		_, err = cs.EncryptDanger(nil, nil, []byte("x"), RejectAfterMessages, nb)
		require.ErrorIs(t, err, ErrMessageCounterExhausted)
	}
}

func BenchmarkCipherStateEncryptAESGCM(b *testing.B) {
	enc, _ := buildCipherStatesB(b, CipherAESGCM)
	benchEncryptCipherState(b, NewCipherState(enc, CipherAESGCM))
}

func BenchmarkCipherStateEncryptChaChaPoly(b *testing.B) {
	enc, _ := buildCipherStatesB(b, noise.CipherChaChaPoly)
	benchEncryptCipherState(b, NewCipherState(enc, noise.CipherChaChaPoly))
}

func benchEncryptCipherState(b *testing.B, cs CipherState) {
	plaintext := make([]byte, 1280)
	ad := make([]byte, 16)
	nb := make([]byte, 12)
	out := make([]byte, 0, len(plaintext)+cs.Overhead())
	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		var err error
		out, err = cs.EncryptDanger(out[:0], ad, plaintext, uint64(i+1), nb)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func buildCipherStatesB(b *testing.B, c noise.CipherFunc) (*noise.CipherState, *noise.CipherState) {
	b.Helper()
	suite := noise.NewCipherSuite(noise.DH25519, c, noise.HashSHA256)
	cfg := noise.Config{CipherSuite: suite, Pattern: noise.HandshakeNN}
	cfg.Initiator = true
	hsI, err := noise.NewHandshakeState(cfg)
	if err != nil {
		b.Fatal(err)
	}
	cfg.Initiator = false
	hsR, err := noise.NewHandshakeState(cfg)
	if err != nil {
		b.Fatal(err)
	}
	msg, _, _, err := hsI.WriteMessage(nil, nil)
	if err != nil {
		b.Fatal(err)
	}
	if _, _, _, err := hsR.ReadMessage(nil, msg); err != nil {
		b.Fatal(err)
	}
	msg, dR, _, err := hsR.WriteMessage(nil, nil)
	if err != nil {
		b.Fatal(err)
	}
	_, eI, _, err := hsI.ReadMessage(nil, msg)
	if err != nil {
		b.Fatal(err)
	}
	return eI, dR
}

func TestCipherStateNilSafety(t *testing.T) {
	var aes *CipherStateAESGCM
	_, err := aes.EncryptDanger(nil, nil, nil, 0, make([]byte, 12))
	require.Error(t, err)
	out, err := aes.DecryptDanger(nil, nil, nil, 0, make([]byte, 12))
	require.NoError(t, err)
	assert.Empty(t, out)
	assert.Equal(t, 0, aes.Overhead())

	var cc *CipherStateChaChaPoly
	_, err = cc.EncryptDanger(nil, nil, nil, 0, make([]byte, 12))
	require.Error(t, err)
	out, err = cc.DecryptDanger(nil, nil, nil, 0, make([]byte, 12))
	require.NoError(t, err)
	assert.Empty(t, out)
	assert.Equal(t, 0, cc.Overhead())
}

func TestCipherStateAESGCMInPlaceDecrypt(t *testing.T) {
	enc, dec := buildCipherStates(t, CipherAESGCM)
	inPlaceDecrypt(t, NewCipherStateAESGCM(enc), NewCipherStateAESGCM(dec))
}

func TestCipherStateChaChaPolyInPlaceDecrypt(t *testing.T) {
	enc, dec := buildCipherStates(t, noise.CipherChaChaPoly)
	inPlaceDecrypt(t, NewCipherStateChaChaPoly(enc), NewCipherStateChaChaPoly(dec))
}

func inPlaceDecrypt(t *testing.T, enc, dec CipherState) {
	t.Helper()
	const hdrLen = 16
	plaintext := []byte("in-place decrypt should replace the ciphertext bytes")
	nb := make([]byte, 12)

	// packet = [16-byte header | ciphertext+tag], like a nebula Message.
	packet := make([]byte, hdrLen, hdrLen+len(plaintext)+enc.Overhead())
	for i := range packet {
		packet[i] = byte(i)
	}
	packet, err := enc.EncryptDanger(packet, packet[:hdrLen], plaintext, 1, nb)
	require.NoError(t, err)

	// Simulate a GRO row: [packet | next segment]. A failed auth on packet
	// may zero packet's plaintext region but must not touch the header, the
	// tag, or the neighboring segment.
	neighbor := []byte("next coalesced segment, must stay intact")
	row := append(append([]byte(nil), packet...), neighbor...)
	tampered := row[:len(packet)]
	tampered[hdrLen] ^= 0x01
	_, err = dec.DecryptDanger(tampered[hdrLen:hdrLen], tampered[:hdrLen], tampered[hdrLen:], 1, nb)
	require.Error(t, err)
	assert.Equal(t, packet[:hdrLen], tampered[:hdrLen], "failed auth must not touch the header")
	assert.Equal(t, packet[len(packet)-dec.Overhead():], tampered[len(tampered)-dec.Overhead():],
		"failed auth must not touch the tag")
	assert.Equal(t, neighbor, row[len(packet):], "failed auth must not touch the next segment")

	out, err := dec.DecryptDanger(packet[hdrLen:hdrLen], packet[:hdrLen], packet[hdrLen:], 1, nb)
	require.NoError(t, err)
	assert.Equal(t, plaintext, out)
	// The plaintext must be IN the packet buffer, not a fresh allocation.
	assert.Equal(t, &packet[hdrLen], &out[0], "plaintext must alias the packet buffer")
}