github.com/mikelsr/quic-go@v0.36.1-0.20230701132136-1d9415b66898/internal/handshake/token_generator.go (about)

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