github.com/apernet/quic-go@v0.43.1-0.20240515053213-5e9e635fd9f0/internal/handshake/token_protector.go (about)

     1  package handshake
     2  
     3  import (
     4  	"crypto/aes"
     5  	"crypto/cipher"
     6  	"crypto/rand"
     7  	"crypto/sha256"
     8  	"fmt"
     9  	"io"
    10  
    11  	"golang.org/x/crypto/hkdf"
    12  )
    13  
    14  // TokenProtectorKey is the key used to encrypt both Retry and session resumption tokens.
    15  type TokenProtectorKey [32]byte
    16  
    17  // TokenProtector is used to create and verify a token
    18  type tokenProtector interface {
    19  	// NewToken creates a new token
    20  	NewToken([]byte) ([]byte, error)
    21  	// DecodeToken decodes a token
    22  	DecodeToken([]byte) ([]byte, error)
    23  }
    24  
    25  const tokenNonceSize = 32
    26  
    27  // tokenProtector is used to create and verify a token
    28  type tokenProtectorImpl struct {
    29  	key TokenProtectorKey
    30  }
    31  
    32  // newTokenProtector creates a source for source address tokens
    33  func newTokenProtector(key TokenProtectorKey) tokenProtector {
    34  	return &tokenProtectorImpl{key: key}
    35  }
    36  
    37  // NewToken encodes data into a new token.
    38  func (s *tokenProtectorImpl) NewToken(data []byte) ([]byte, error) {
    39  	var nonce [tokenNonceSize]byte
    40  	if _, err := rand.Read(nonce[:]); err != nil {
    41  		return nil, err
    42  	}
    43  	aead, aeadNonce, err := s.createAEAD(nonce[:])
    44  	if err != nil {
    45  		return nil, err
    46  	}
    47  	return append(nonce[:], aead.Seal(nil, aeadNonce, data, nil)...), nil
    48  }
    49  
    50  // DecodeToken decodes a token.
    51  func (s *tokenProtectorImpl) DecodeToken(p []byte) ([]byte, error) {
    52  	if len(p) < tokenNonceSize {
    53  		return nil, fmt.Errorf("token too short: %d", len(p))
    54  	}
    55  	nonce := p[:tokenNonceSize]
    56  	aead, aeadNonce, err := s.createAEAD(nonce)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	return aead.Open(nil, aeadNonce, p[tokenNonceSize:], nil)
    61  }
    62  
    63  func (s *tokenProtectorImpl) createAEAD(nonce []byte) (cipher.AEAD, []byte, error) {
    64  	h := hkdf.New(sha256.New, s.key[:], nonce, []byte("quic-go token source"))
    65  	key := make([]byte, 32) // use a 32 byte key, in order to select AES-256
    66  	if _, err := io.ReadFull(h, key); err != nil {
    67  		return nil, nil, err
    68  	}
    69  	aeadNonce := make([]byte, 12)
    70  	if _, err := io.ReadFull(h, aeadNonce); err != nil {
    71  		return nil, nil, err
    72  	}
    73  	c, err := aes.NewCipher(key)
    74  	if err != nil {
    75  		return nil, nil, err
    76  	}
    77  	aead, err := cipher.NewGCM(c)
    78  	if err != nil {
    79  		return nil, nil, err
    80  	}
    81  	return aead, aeadNonce, nil
    82  }