github.com/TugasAkhir-QUIC/quic-go@v0.0.2-0.20240215011318-d20e25a9054c/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/TugasAkhir-QUIC/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         []byte
    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  		mask:         make([]byte, block.BlockSize()),
    56  		isLongHeader: isLongHeader,
    57  	}
    58  }
    59  
    60  func (p *aesHeaderProtector) DecryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
    61  	p.apply(sample, firstByte, hdrBytes)
    62  }
    63  
    64  func (p *aesHeaderProtector) EncryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
    65  	p.apply(sample, firstByte, hdrBytes)
    66  }
    67  
    68  func (p *aesHeaderProtector) apply(sample []byte, firstByte *byte, hdrBytes []byte) {
    69  	if len(sample) != len(p.mask) {
    70  		panic("invalid sample size")
    71  	}
    72  	p.block.Encrypt(p.mask, sample)
    73  	if p.isLongHeader {
    74  		*firstByte ^= p.mask[0] & 0xf
    75  	} else {
    76  		*firstByte ^= p.mask[0] & 0x1f
    77  	}
    78  	for i := range hdrBytes {
    79  		hdrBytes[i] ^= p.mask[i+1]
    80  	}
    81  }
    82  
    83  type chachaHeaderProtector struct {
    84  	mask [5]byte
    85  
    86  	key          [32]byte
    87  	isLongHeader bool
    88  }
    89  
    90  var _ headerProtector = &chachaHeaderProtector{}
    91  
    92  func newChaChaHeaderProtector(suite *cipherSuite, trafficSecret []byte, isLongHeader bool, hkdfLabel string) headerProtector {
    93  	hpKey := hkdfExpandLabel(suite.Hash, trafficSecret, []byte{}, hkdfLabel, suite.KeyLen)
    94  
    95  	p := &chachaHeaderProtector{
    96  		isLongHeader: isLongHeader,
    97  	}
    98  	copy(p.key[:], hpKey)
    99  	return p
   100  }
   101  
   102  func (p *chachaHeaderProtector) DecryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
   103  	p.apply(sample, firstByte, hdrBytes)
   104  }
   105  
   106  func (p *chachaHeaderProtector) EncryptHeader(sample []byte, firstByte *byte, hdrBytes []byte) {
   107  	p.apply(sample, firstByte, hdrBytes)
   108  }
   109  
   110  func (p *chachaHeaderProtector) apply(sample []byte, firstByte *byte, hdrBytes []byte) {
   111  	if len(sample) != 16 {
   112  		panic("invalid sample size")
   113  	}
   114  	for i := 0; i < 5; i++ {
   115  		p.mask[i] = 0
   116  	}
   117  	cipher, err := chacha20.NewUnauthenticatedCipher(p.key[:], sample[4:])
   118  	if err != nil {
   119  		panic(err)
   120  	}
   121  	cipher.SetCounter(binary.LittleEndian.Uint32(sample[:4]))
   122  	cipher.XORKeyStream(p.mask[:], p.mask[:])
   123  	p.applyMask(firstByte, hdrBytes)
   124  }
   125  
   126  func (p *chachaHeaderProtector) applyMask(firstByte *byte, hdrBytes []byte) {
   127  	if p.isLongHeader {
   128  		*firstByte ^= p.mask[0] & 0xf
   129  	} else {
   130  		*firstByte ^= p.mask[0] & 0x1f
   131  	}
   132  	for i := range hdrBytes {
   133  		hdrBytes[i] ^= p.mask[i+1]
   134  	}
   135  }