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

     1  package handshake
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/asn1"
     6  	"fmt"
     7  	"net"
     8  	"time"
     9  
    10  	"github.com/daeuniverse/quic-go/internal/protocol"
    11  )
    12  
    13  const (
    14  	tokenPrefixIP byte = iota
    15  	tokenPrefixString
    16  )
    17  
    18  // A Token is derived from the client address and can be used to verify the ownership of this address.
    19  type Token struct {
    20  	IsRetryToken      bool
    21  	SentTime          time.Time
    22  	encodedRemoteAddr []byte
    23  	// only set for retry tokens
    24  	OriginalDestConnectionID protocol.ConnectionID
    25  	RetrySrcConnectionID     protocol.ConnectionID
    26  }
    27  
    28  // ValidateRemoteAddr validates the address, but does not check expiration
    29  func (t *Token) ValidateRemoteAddr(addr net.Addr) bool {
    30  	return bytes.Equal(encodeRemoteAddr(addr), t.encodedRemoteAddr)
    31  }
    32  
    33  // token is the struct that is used for ASN1 serialization and deserialization
    34  type token struct {
    35  	IsRetryToken             bool
    36  	RemoteAddr               []byte
    37  	Timestamp                int64
    38  	OriginalDestConnectionID []byte
    39  	RetrySrcConnectionID     []byte
    40  }
    41  
    42  // A TokenGenerator generates tokens
    43  type TokenGenerator struct {
    44  	tokenProtector tokenProtector
    45  }
    46  
    47  // NewTokenGenerator initializes a new TokenGenerator
    48  func NewTokenGenerator(key TokenProtectorKey) *TokenGenerator {
    49  	return &TokenGenerator{tokenProtector: newTokenProtector(key)}
    50  }
    51  
    52  // NewRetryToken generates a new token for a Retry for a given source address
    53  func (g *TokenGenerator) NewRetryToken(
    54  	raddr net.Addr,
    55  	origDestConnID protocol.ConnectionID,
    56  	retrySrcConnID protocol.ConnectionID,
    57  ) ([]byte, error) {
    58  	data, err := asn1.Marshal(token{
    59  		IsRetryToken:             true,
    60  		RemoteAddr:               encodeRemoteAddr(raddr),
    61  		OriginalDestConnectionID: origDestConnID.Bytes(),
    62  		RetrySrcConnectionID:     retrySrcConnID.Bytes(),
    63  		Timestamp:                time.Now().UnixNano(),
    64  	})
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  	return g.tokenProtector.NewToken(data)
    69  }
    70  
    71  // NewToken generates a new token to be sent in a NEW_TOKEN frame
    72  func (g *TokenGenerator) NewToken(raddr net.Addr) ([]byte, error) {
    73  	data, err := asn1.Marshal(token{
    74  		RemoteAddr: encodeRemoteAddr(raddr),
    75  		Timestamp:  time.Now().UnixNano(),
    76  	})
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  	return g.tokenProtector.NewToken(data)
    81  }
    82  
    83  // DecodeToken decodes a token
    84  func (g *TokenGenerator) DecodeToken(encrypted []byte) (*Token, error) {
    85  	// if the client didn't send any token, DecodeToken will be called with a nil-slice
    86  	if len(encrypted) == 0 {
    87  		return nil, nil
    88  	}
    89  
    90  	data, err := g.tokenProtector.DecodeToken(encrypted)
    91  	if err != nil {
    92  		return nil, err
    93  	}
    94  	t := &token{}
    95  	rest, err := asn1.Unmarshal(data, t)
    96  	if err != nil {
    97  		return nil, err
    98  	}
    99  	if len(rest) != 0 {
   100  		return nil, fmt.Errorf("rest when unpacking token: %d", len(rest))
   101  	}
   102  	token := &Token{
   103  		IsRetryToken:      t.IsRetryToken,
   104  		SentTime:          time.Unix(0, t.Timestamp),
   105  		encodedRemoteAddr: t.RemoteAddr,
   106  	}
   107  	if t.IsRetryToken {
   108  		token.OriginalDestConnectionID = protocol.ParseConnectionID(t.OriginalDestConnectionID)
   109  		token.RetrySrcConnectionID = protocol.ParseConnectionID(t.RetrySrcConnectionID)
   110  	}
   111  	return token, nil
   112  }
   113  
   114  // encodeRemoteAddr encodes a remote address such that it can be saved in the token
   115  func encodeRemoteAddr(remoteAddr net.Addr) []byte {
   116  	if udpAddr, ok := remoteAddr.(*net.UDPAddr); ok {
   117  		return append([]byte{tokenPrefixIP}, udpAddr.IP...)
   118  	}
   119  	return append([]byte{tokenPrefixString}, []byte(remoteAddr.String())...)
   120  }