<commit_before>package handshake

import (
	"crypto/cipher"
	"encoding/binary"

	"github.com/lucas-clemente/quic-go/internal/protocol"
)

type shortHeaderOpener struct {
	aead        cipher.AEAD
	hpDecrypter cipher.Block

	// use a single slice to avoid allocations
	nonceBuf []byte
	hpMask   []byte
}

var _ ShortHeaderOpener = &shortHeaderOpener{}

func newShortHeaderOpener(aead cipher.AEAD, hpDecrypter cipher.Block) ShortHeaderOpener {
	return &shortHeaderOpener{
		aead:        aead,
		nonceBuf:    make([]byte, aead.NonceSize()),
		hpDecrypter: hpDecrypter,
		hpMask:      make([]byte, hpDecrypter.BlockSize()),
	}
}

func (o *shortHeaderOpener) Open(dst, src []byte, pn protocol.PacketNumber, _ protocol.KeyPhase, ad []byte) ([]byte, error) {
	binary.BigEndian.PutUint64(o.nonceBuf[len(o.nonceBuf)-8:], uint64(pn))
	// The AEAD we're using here will be the qtls.aeadAESGCM13.
	// It uses the nonce provided here and XOR it with the IV.
	return o.aead.Open(dst, o.nonceBuf, src, ad)
}

func (o *shortHeaderOpener) DecryptHeader(sample []byte, firstByte *byte, pnBytes []byte) {
	if len(sample) != o.hpDecrypter.BlockSize() {
		panic("invalid sample size")
	}
	o.hpDecrypter.Encrypt(o.hpMask, sample)
	*firstByte ^= o.hpMask[0] & 0x1f
	for i := range pnBytes {
		pnBytes[i] ^= o.hpMask[i+1]
	}
}

type shortHeaderSealer struct {
	sealer
}

func newShortHeaderSealer(aead cipher.AEAD, hpEncrypter cipher.Block) ShortHeaderSealer {
	return &shortHeaderSealer{
		sealer: sealer{
			aead:        aead,
			nonceBuf:    make([]byte, aead.NonceSize()),
			hpEncrypter: hpEncrypter,
			hpMask:      make([]byte, hpEncrypter.BlockSize()),
		},
	}
}

func (s *shortHeaderSealer) EncryptHeader(sample []byte, firstByte *byte, pnBytes []byte) {
	if len(sample) != s.hpEncrypter.BlockSize() {
		panic("invalid sample size")
	}
	s.hpEncrypter.Encrypt(s.hpMask, sample)
	*firstByte ^= s.hpMask[0] & 0x1f
	for i := range pnBytes {
		pnBytes[i] ^= s.hpMask[i+1]
	}
}

func (s *shortHeaderSealer) KeyPhase() protocol.KeyPhase {
	return protocol.KeyPhaseOne
}

type updatableAEAD struct {
	ShortHeaderOpener
	ShortHeaderSealer
}

func newUpdatableAEAD() *updatableAEAD {
	return &updatableAEAD{}
}

func (a *updatableAEAD) SetReadKey(suite cipherSuite, trafficSecret []byte) {
	a.ShortHeaderOpener = newShortHeaderOpener(createAEAD(suite, trafficSecret))
}

func (a *updatableAEAD) SetWriteKey(suite cipherSuite, trafficSecret []byte) {
	a.ShortHeaderSealer = newShortHeaderSealer(createAEAD(suite, trafficSecret))
}<commit_msg>move header encryption / decryption to the updatable AEAD<commit_after>package handshake

import (
	"crypto/cipher"
	"encoding/binary"

	"github.com/lucas-clemente/quic-go/internal/protocol"
)

type shortHeaderOpener struct {
	aead cipher.AEAD

	// use a single slice to avoid allocations
	nonceBuf []byte
}

func newShortHeaderOpener(aead cipher.AEAD) *shortHeaderOpener {
	return &shortHeaderOpener{
		aead:     aead,
		nonceBuf: make([]byte, aead.NonceSize()),
	}
}

func (o *shortHeaderOpener) Open(dst, src []byte, pn protocol.PacketNumber, _ protocol.KeyPhase, ad []byte) ([]byte, error) {
	binary.BigEndian.PutUint64(o.nonceBuf[len(o.nonceBuf)-8:], uint64(pn))
	// The AEAD we're using here will be the qtls.aeadAESGCM13.
	// It uses the nonce provided here and XOR it with the IV.
	return o.aead.Open(dst, o.nonceBuf, src, ad)
}

type shortHeaderSealer struct {
	sealer
}

func newShortHeaderSealer(aead cipher.AEAD) *shortHeaderSealer {
	return &shortHeaderSealer{
		sealer: sealer{
			aead:     aead,
			nonceBuf: make([]byte, aead.NonceSize()),
		},
	}
}

func (s *shortHeaderSealer) KeyPhase() protocol.KeyPhase {
	return protocol.KeyPhaseOne
}

type updatableAEAD struct {
	*shortHeaderSealer
	*shortHeaderOpener

	hpDecrypter cipher.Block
	hpEncrypter cipher.Block

	// use a single slice to avoid allocations
	hpMask []byte
}

var _ ShortHeaderOpener = &updatableAEAD{}
var _ ShortHeaderSealer = &updatableAEAD{}

func newUpdatableAEAD() *updatableAEAD {
	return &updatableAEAD{}
}

func (a *updatableAEAD) SetReadKey(suite cipherSuite, trafficSecret []byte) {
	aead, hpDecrypter := createAEAD(suite, trafficSecret)
	a.shortHeaderOpener = newShortHeaderOpener(aead)
	if len(a.hpMask) == 0 {
		a.hpMask = make([]byte, hpDecrypter.BlockSize())
	} else if len(a.hpMask) != hpDecrypter.BlockSize() {
		panic("invalid header protection block size")
	}
	a.hpDecrypter = hpDecrypter
}

func (a *updatableAEAD) SetWriteKey(suite cipherSuite, trafficSecret []byte) {
	aead, hpEncrypter := createAEAD(suite, trafficSecret)
	a.shortHeaderSealer = newShortHeaderSealer(aead)
	if len(a.hpMask) == 0 {
		a.hpMask = make([]byte, hpEncrypter.BlockSize())
	} else if len(a.hpMask) != hpEncrypter.BlockSize() {
		panic("invalid header protection block size")
	}
	a.hpEncrypter = hpEncrypter
}

func (a *updatableAEAD) EncryptHeader(sample []byte, firstByte *byte, pnBytes []byte) {
	if len(sample) != a.hpEncrypter.BlockSize() {
		panic("invalid sample size")
	}
	a.hpEncrypter.Encrypt(a.hpMask, sample)
	*firstByte ^= a.hpMask[0] & 0x1f
	for i := range pnBytes {
		pnBytes[i] ^= a.hpMask[i+1]
	}
}

func (a *updatableAEAD) DecryptHeader(sample []byte, firstByte *byte, pnBytes []byte) {
	if len(sample) != a.hpDecrypter.BlockSize() {
		panic("invalid sample size")
	}
	a.hpDecrypter.Encrypt(a.hpMask, sample)
	*firstByte ^= a.hpMask[0] & 0x1f
	for i := range pnBytes {
		pnBytes[i] ^= a.hpMask[i+1]
	}
}
{'fraction_non_alphanumeric': 0.12850812407680945, 'fraction_numerical': 0.005908419497784343, 'mean_word_length': 4.868905742145179, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6522944', 'n_tokens_mistral': 2164, 'n_tokens_neox': 2081, 'n_words': 443}