github.com/daeuniverse/quic-go@v0.0.0-20240413031024-943f218e0810/internal/handshake/header_protector.go (about)

     1  package handshake
     2  
     3  import (
     4  	"crypto/aes"
     5  	"crypto/cipher"
     6  	"crypto/tls"
     7  	"encoding/binary"
     8  	"fmt"
     9  
    10  	"golang.org/x/crypto/chacha20"
    11  
    12  	"github.com/daeuniverse/quic-go/internal/protocol"
    13  )
    14  
    15  type headerProtector interface {
    16  	EncryptHeader(sample []byte, firstByte *byte, hdrBytes []byte)
    17  	DecryptHeader(sample []byte, firstByte *byte, hdrBytes []byte)
    18  }
    19  
    20  func hkdfHeaderProtectionLabel(v protocol.Version) string {
    21  	if v == protocol.Version2 {
    22  		return "quicv2 hp"
    23  	}
    24  	return "quic hp"
    25  }
    26  
    27  func newHeaderProtector(suite *cipherSuite, trafficSecret []byte, isLongHeader bool, v protocol.Version) headerProtector {
    28  	hkdfLabel := hkdfHeaderProtectionLabel(v)
    29  	switch suite.ID {
    30  	case tls.TLS_AES_128_GCM_SHA256, tls.TLS_AES_256_GCM_SHA384:
    31  		return newAESHeaderProtector(suite, trafficSecret, isLongHeader, hkdfLabel)
    32  	case tls.TLS_CHACHA20_POLY1305_SHA256:
    33  		return newChaChaHeaderProtector(suite, trafficSecret, isLongHeader, hkdfLabel)
    34  	default:
    35  		panic(fmt.Sprintf("Invalid cipher suite id: %d", suite.ID))
    36  	}
    37  }
    38  
    39  type aesHeaderProtector struct {
    40  	mask         [16]byte // AES always has a 16 byte block size
    41  	block        cipher.Block
    42  	isLongHeader bool
    43  }
    44  
    45  var _ headerProtector = &aesHeaderProtector{}
    46  
    47  func newAESHeaderProtector(suite *cipherSuite, trafficSecret []byte, isLongHeader bool, hkdfLabel string) headerProtector {
    48  	hpKey := hkdfExpandLabel(suite.Hash, trafficSecret, []byte{}, hkdfLabel, suite.KeyLen)
    49  	block, err := aes.NewCipher(hpKey)
    50  	if err != nil {
    51  		panic(fmt.Sprintf("error creating new AES cipher: %s", err))
    52  	}
    53  	return &aesHeaderProtector{
    54  		block:        block,
    55  		isLongHeader: isLongHeader,
    56  	}
    57  }
    58  
    59  func (p *aesHeaderProtector) DecryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
    60  	p.apply(sample, firstByte, hdrBytes)
    61  }
    62  
    63  func (p *aesHeaderProtector) EncryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
    64  	p.apply(sample, firstByte, hdrBytes)
    65  }
    66  
    67  func (p *aesHeaderProtector) apply(sample []byte, firstByte *byte, hdrBytes []byte) {
    68  	if len(sample) != len(p.mask) {
    69  		panic("invalid sample size")
    70  	}
    71  	p.block.Encrypt(p.mask[:], sample)
    72  	if p.isLongHeader {
    73  		*firstByte ^= p.mask[0] & 0xf
    74  	} else {
    75  		*firstByte ^= p.mask[0] & 0x1f
    76  	}
    77  	for i := range hdrBytes {
    78  		hdrBytes[i] ^= p.mask[i+1]
    79  	}
    80  }
    81  
    82  type chachaHeaderProtector struct {
    83  	mask [5]byte
    84  
    85  	key          [32]byte
    86  	isLongHeader bool
    87  }
    88  
    89  var _ headerProtector = &chachaHeaderProtector{}
    90  
    91  func newChaChaHeaderProtector(suite *cipherSuite, trafficSecret []byte, isLongHeader bool, hkdfLabel string) headerProtector {
    92  	hpKey := hkdfExpandLabel(suite.Hash, trafficSecret, []byte{}, hkdfLabel, suite.KeyLen)
    93  
    94  	p := &chachaHeaderProtector{
    95  		isLongHeader: isLongHeader,
    96  	}
    97  	copy(p.key[:], hpKey)
    98  	return p
    99  }
   100  
   101  func (p *chachaHeaderProtector) DecryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
   102  	p.apply(sample, firstByte, hdrBytes)
   103  }
   104  
   105  func (p *chachaHeaderProtector) EncryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
   106  	p.apply(sample, firstByte, hdrBytes)
   107  }
   108  
   109  func (p *chachaHeaderProtector) apply(sample []byte, firstByte *byte, hdrBytes []byte) {
   110  	if len(sample) != 16 {
   111  		panic("invalid sample size")
   112  	}
   113  	for i := 0; i < 5; i++ {
   114  		p.mask[i] = 0
   115  	}
   116  	cipher, err := chacha20.NewUnauthenticatedCipher(p.key[:], sample[4:])
   117  	if err != nil {
   118  		panic(err)
   119  	}
   120  	cipher.SetCounter(binary.LittleEndian.Uint32(sample[:4]))
   121  	cipher.XORKeyStream(p.mask[:], p.mask[:])
   122  	p.applyMask(firstByte, hdrBytes)
   123  }
   124  
   125  func (p *chachaHeaderProtector) applyMask(firstByte *byte, hdrBytes []byte) {
   126  	if p.isLongHeader {
   127  		*firstByte ^= p.mask[0] & 0xf
   128  	} else {
   129  		*firstByte ^= p.mask[0] & 0x1f
   130  	}
   131  	for i := range hdrBytes {
   132  		hdrBytes[i] ^= p.mask[i+1]
   133  	}
   134  }