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