encode the original connection ID into the cookie

internal/handshake/cookie_generator.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"net"
 	"time"
+
+	"github.com/lucas-clemente/quic-go/internal/protocol"
 )
 
 const (
@@ -14,14 +16,17 @@ const (
 
 // A Cookie is derived from the client address and can be used to verify the ownership of this address.
 type Cookie struct {
-	RemoteAddr string
-	// The time that the STK was issued (resolution 1 second)
+	RemoteAddr               string
+	OriginalDestConnectionID protocol.ConnectionID
+	// The time that the Cookie was issued (resolution 1 second)
 	SentTime time.Time
 }
 
 // token is the struct that is used for ASN1 serialization and deserialization
 type token struct {
-	Data      []byte
+	RemoteAddr               []byte
+	OriginalDestConnectionID []byte
+
 	Timestamp int64
 }
 
@@ -42,10 +47,11 @@ func NewCookieGenerator() (*CookieGenerator, error) {
 }
 
 // NewToken generates a new Cookie for a given source address
-func (g *CookieGenerator) NewToken(raddr net.Addr) ([]byte, error) {
+func (g *CookieGenerator) NewToken(raddr net.Addr, origConnID protocol.ConnectionID) ([]byte, error) {
 	data, err := asn1.Marshal(token{
-		Data:      encodeRemoteAddr(raddr),
-		Timestamp: time.Now().Unix(),
+		RemoteAddr:               encodeRemoteAddr(raddr),
+		OriginalDestConnectionID: origConnID,
+		Timestamp:                time.Now().Unix(),
 	})
 	if err != nil {
 		return nil, err
@@ -72,10 +78,14 @@ func (g *CookieGenerator) DecodeToken(encrypted []byte) (*Cookie, error) {
 	if len(rest) != 0 {
 		return nil, fmt.Errorf("rest when unpacking token: %d", len(rest))
 	}
-	return &Cookie{
-		RemoteAddr: decodeRemoteAddr(t.Data),
+	cookie := &Cookie{
+		RemoteAddr: decodeRemoteAddr(t.RemoteAddr),
 		SentTime:   time.Unix(t.Timestamp, 0),
-	}, nil
+	}
+	if len(t.OriginalDestConnectionID) > 0 {
+		cookie.OriginalDestConnectionID = protocol.ConnectionID(t.OriginalDestConnectionID)
+	}
+	return cookie, nil
 }
 
 // encodeRemoteAddr encodes a remote address such that it can be saved in the Cookie

internal/handshake/cookie_generator_test.go

@@ -5,6 +5,8 @@ import (
 	"net"
 	"time"
 
+	"github.com/lucas-clemente/quic-go/internal/protocol"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
@@ -20,7 +22,7 @@ var _ = Describe("Cookie Generator", func() {
 
 	It("generates a Cookie", func() {
 		ip := net.IPv4(127, 0, 0, 1)
-		token, err := cookieGen.NewToken(&net.UDPAddr{IP: ip, Port: 1337})
+		token, err := cookieGen.NewToken(&net.UDPAddr{IP: ip, Port: 1337}, nil)
 		Expect(err).ToNot(HaveOccurred())
 		Expect(token).ToNot(BeEmpty())
 	})
@@ -33,7 +35,10 @@ var _ = Describe("Cookie Generator", func() {
 
 	It("accepts a valid cookie", func() {
 		ip := net.IPv4(192, 168, 0, 1)
-		token, err := cookieGen.NewToken(&net.UDPAddr{IP: ip, Port: 1337})
+		token, err := cookieGen.NewToken(
+			&net.UDPAddr{IP: ip, Port: 1337},
+			nil,
+		)
 		Expect(err).ToNot(HaveOccurred())
 		cookie, err := cookieGen.DecodeToken(token)
 		Expect(err).ToNot(HaveOccurred())
@@ -41,6 +46,18 @@ var _ = Describe("Cookie Generator", func() {
 		// the time resolution of the Cookie is just 1 second
 		// if Cookie generation and this check happen in "different seconds", the difference will be between 1 and 2 seconds
 		Expect(cookie.SentTime).To(BeTemporally("~", time.Now(), 2*time.Second))
+		Expect(cookie.OriginalDestConnectionID).To(BeNil())
+	})
+
+	It("saves the connection ID", func() {
+		token, err := cookieGen.NewToken(
+			&net.UDPAddr{},
+			protocol.ConnectionID{0xde, 0xad, 0xbe, 0xef},
+		)
+		Expect(err).ToNot(HaveOccurred())
+		cookie, err := cookieGen.DecodeToken(token)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(cookie.OriginalDestConnectionID).To(Equal(protocol.ConnectionID{0xde, 0xad, 0xbe, 0xef}))
 	})
 
 	It("rejects invalid tokens", func() {
@@ -56,7 +73,7 @@ var _ = Describe("Cookie Generator", func() {
 	})
 
 	It("rejects tokens that can be decoded, but have additional payload", func() {
-		t, err := asn1.Marshal(token{Data: []byte("foobar")})
+		t, err := asn1.Marshal(token{RemoteAddr: []byte("foobar")})
 		Expect(err).ToNot(HaveOccurred())
 		t = append(t, []byte("rest")...)
 		enc, err := cookieGen.cookieProtector.NewToken(t)
@@ -67,7 +84,7 @@ var _ = Describe("Cookie Generator", func() {
 
 	// we don't generate tokens that have no data, but we should be able to handle them if we receive one for whatever reason
 	It("doesn't panic if a tokens has no data", func() {
-		t, err := asn1.Marshal(token{Data: []byte("")})
+		t, err := asn1.Marshal(token{RemoteAddr: []byte("")})
 		Expect(err).ToNot(HaveOccurred())
 		enc, err := cookieGen.cookieProtector.NewToken(t)
 		Expect(err).ToNot(HaveOccurred())
@@ -86,7 +103,7 @@ var _ = Describe("Cookie Generator", func() {
 			ip := net.ParseIP(addr)
 			Expect(ip).ToNot(BeNil())
 			raddr := &net.UDPAddr{IP: ip, Port: 1337}
-			token, err := cookieGen.NewToken(raddr)
+			token, err := cookieGen.NewToken(raddr, nil)
 			Expect(err).ToNot(HaveOccurred())
 			cookie, err := cookieGen.DecodeToken(token)
 			Expect(err).ToNot(HaveOccurred())
@@ -99,7 +116,7 @@ var _ = Describe("Cookie Generator", func() {
 
 	It("uses the string representation an address that is not a UDP address", func() {
 		raddr := &net.TCPAddr{IP: net.IPv4(192, 168, 13, 37), Port: 1337}
-		token, err := cookieGen.NewToken(raddr)
+		token, err := cookieGen.NewToken(raddr, nil)
 		Expect(err).ToNot(HaveOccurred())
 		cookie, err := cookieGen.DecodeToken(token)
 		Expect(err).ToNot(HaveOccurred())