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

     1  package handshake
     2  
     3  import (
     4  	"crypto/rand"
     5  	"encoding/asn1"
     6  	"net"
     7  	"time"
     8  
     9  	"github.com/apernet/quic-go/internal/protocol"
    10  
    11  	. "github.com/onsi/ginkgo/v2"
    12  	. "github.com/onsi/gomega"
    13  )
    14  
    15  var _ = Describe("Token Generator", func() {
    16  	var tokenGen *TokenGenerator
    17  
    18  	BeforeEach(func() {
    19  		var key TokenProtectorKey
    20  		rand.Read(key[:])
    21  		tokenGen = NewTokenGenerator(key)
    22  	})
    23  
    24  	It("generates a token", func() {
    25  		ip := net.IPv4(127, 0, 0, 1)
    26  		token, err := tokenGen.NewRetryToken(&net.UDPAddr{IP: ip, Port: 1337}, protocol.ConnectionID{}, protocol.ConnectionID{})
    27  		Expect(err).ToNot(HaveOccurred())
    28  		Expect(token).ToNot(BeEmpty())
    29  	})
    30  
    31  	It("works with nil tokens", func() {
    32  		token, err := tokenGen.DecodeToken(nil)
    33  		Expect(err).ToNot(HaveOccurred())
    34  		Expect(token).To(BeNil())
    35  	})
    36  
    37  	It("accepts a valid token", func() {
    38  		addr := &net.UDPAddr{IP: net.IPv4(192, 168, 0, 1), Port: 1337}
    39  		tokenEnc, err := tokenGen.NewRetryToken(addr, protocol.ConnectionID{}, protocol.ConnectionID{})
    40  		Expect(err).ToNot(HaveOccurred())
    41  		token, err := tokenGen.DecodeToken(tokenEnc)
    42  		Expect(err).ToNot(HaveOccurred())
    43  		Expect(token.ValidateRemoteAddr(addr)).To(BeTrue())
    44  		Expect(token.ValidateRemoteAddr(&net.UDPAddr{IP: net.IPv4(192, 168, 0, 2), Port: 1337})).To(BeFalse())
    45  		Expect(token.SentTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
    46  		Expect(token.OriginalDestConnectionID.Len()).To(BeZero())
    47  		Expect(token.RetrySrcConnectionID.Len()).To(BeZero())
    48  	})
    49  
    50  	It("saves the connection ID", func() {
    51  		connID1 := protocol.ParseConnectionID([]byte{0xde, 0xad, 0xbe, 0xef})
    52  		connID2 := protocol.ParseConnectionID([]byte{0xde, 0xad, 0xc0, 0xde})
    53  		tokenEnc, err := tokenGen.NewRetryToken(&net.UDPAddr{}, connID1, connID2)
    54  		Expect(err).ToNot(HaveOccurred())
    55  		token, err := tokenGen.DecodeToken(tokenEnc)
    56  		Expect(err).ToNot(HaveOccurred())
    57  		Expect(token.OriginalDestConnectionID).To(Equal(connID1))
    58  		Expect(token.RetrySrcConnectionID).To(Equal(connID2))
    59  	})
    60  
    61  	It("rejects invalid tokens", func() {
    62  		_, err := tokenGen.DecodeToken([]byte("invalid token"))
    63  		Expect(err).To(HaveOccurred())
    64  	})
    65  
    66  	It("rejects tokens that cannot be decoded", func() {
    67  		token, err := tokenGen.tokenProtector.NewToken([]byte("foobar"))
    68  		Expect(err).ToNot(HaveOccurred())
    69  		_, err = tokenGen.DecodeToken(token)
    70  		Expect(err).To(HaveOccurred())
    71  	})
    72  
    73  	It("rejects tokens that can be decoded, but have additional payload", func() {
    74  		t, err := asn1.Marshal(token{RemoteAddr: []byte("foobar")})
    75  		Expect(err).ToNot(HaveOccurred())
    76  		t = append(t, []byte("rest")...)
    77  		enc, err := tokenGen.tokenProtector.NewToken(t)
    78  		Expect(err).ToNot(HaveOccurred())
    79  		_, err = tokenGen.DecodeToken(enc)
    80  		Expect(err).To(MatchError("rest when unpacking token: 4"))
    81  	})
    82  
    83  	// we don't generate tokens that have no data, but we should be able to handle them if we receive one for whatever reason
    84  	It("doesn't panic if a tokens has no data", func() {
    85  		t, err := asn1.Marshal(token{RemoteAddr: []byte("")})
    86  		Expect(err).ToNot(HaveOccurred())
    87  		enc, err := tokenGen.tokenProtector.NewToken(t)
    88  		Expect(err).ToNot(HaveOccurred())
    89  		_, err = tokenGen.DecodeToken(enc)
    90  		Expect(err).ToNot(HaveOccurred())
    91  	})
    92  
    93  	It("works with an IPv6 addresses ", func() {
    94  		addresses := []string{
    95  			"2001:db8::68",
    96  			"2001:0000:4136:e378:8000:63bf:3fff:fdd2",
    97  			"2001::1",
    98  			"ff01:0:0:0:0:0:0:2",
    99  		}
   100  		for _, addr := range addresses {
   101  			ip := net.ParseIP(addr)
   102  			Expect(ip).ToNot(BeNil())
   103  			raddr := &net.UDPAddr{IP: ip, Port: 1337}
   104  			tokenEnc, err := tokenGen.NewRetryToken(raddr, protocol.ConnectionID{}, protocol.ConnectionID{})
   105  			Expect(err).ToNot(HaveOccurred())
   106  			token, err := tokenGen.DecodeToken(tokenEnc)
   107  			Expect(err).ToNot(HaveOccurred())
   108  			Expect(token.ValidateRemoteAddr(raddr)).To(BeTrue())
   109  			Expect(token.SentTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
   110  		}
   111  	})
   112  
   113  	It("uses the string representation an address that is not a UDP address", func() {
   114  		raddr := &net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1337}
   115  		tokenEnc, err := tokenGen.NewRetryToken(raddr, protocol.ConnectionID{}, protocol.ConnectionID{})
   116  		Expect(err).ToNot(HaveOccurred())
   117  		token, err := tokenGen.DecodeToken(tokenEnc)
   118  		Expect(err).ToNot(HaveOccurred())
   119  		Expect(token.ValidateRemoteAddr(raddr)).To(BeTrue())
   120  		Expect(token.ValidateRemoteAddr(&net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1338})).To(BeFalse())
   121  		Expect(token.SentTime).To(BeTemporally("~", time.Now(), 100*time.Millisecond))
   122  	})
   123  })