From 623fcd85b019d64e7e570ca2293925a0fa84ab97 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 13 Aug 2018 11:04:47 +0700
Subject: [PATCH] move the mint cookie protector to the handshake package

It's duplicate code now, but it reduces the dependency on mint.
---
 internal/handshake/cookie_generator.go        |  6 +-
 internal/handshake/cookie_protector.go        | 86 +++++++++++++++++++
 internal/handshake/cookie_protector_test.go   | 39 +++++++++
 .../handshake/crypto_setup_server_test.go     |  4 +-
 4 files changed, 128 insertions(+), 7 deletions(-)
 create mode 100644 internal/handshake/cookie_protector.go
 create mode 100644 internal/handshake/cookie_protector_test.go

diff --git a/internal/handshake/cookie_generator.go b/internal/handshake/cookie_generator.go
index 97accb73d..00f6e7efd 100644
--- a/internal/handshake/cookie_generator.go
+++ b/internal/handshake/cookie_generator.go
@@ -5,8 +5,6 @@ import (
 	"fmt"
 	"net"
 	"time"
-
-	"github.com/bifurcation/mint"
 )
 
 const (
@@ -29,12 +27,12 @@ type token struct {
 
 // A CookieGenerator generates Cookies
 type CookieGenerator struct {
-	cookieProtector mint.CookieProtector
+	cookieProtector cookieProtector
 }
 
 // NewCookieGenerator initializes a new CookieGenerator
 func NewCookieGenerator() (*CookieGenerator, error) {
-	cookieProtector, err := mint.NewDefaultCookieProtector()
+	cookieProtector, err := newCookieProtector()
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/handshake/cookie_protector.go b/internal/handshake/cookie_protector.go
new file mode 100644
index 000000000..7ebdfa18c
--- /dev/null
+++ b/internal/handshake/cookie_protector.go
@@ -0,0 +1,86 @@
+package handshake
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"crypto/sha256"
+	"fmt"
+	"io"
+
+	"golang.org/x/crypto/hkdf"
+)
+
+// CookieProtector is used to create and verify a cookie
+type cookieProtector interface {
+	// NewToken creates a new token
+	NewToken([]byte) ([]byte, error)
+	// DecodeToken decodes a token
+	DecodeToken([]byte) ([]byte, error)
+}
+
+const (
+	cookieSecretSize = 32
+	cookieNonceSize  = 32
+)
+
+// cookieProtector is used to create and verify a cookie
+type cookieProtectorImpl struct {
+	secret []byte
+}
+
+// newCookieProtector creates a source for source address tokens
+func newCookieProtector() (cookieProtector, error) {
+	secret := make([]byte, cookieSecretSize)
+	if _, err := rand.Read(secret); err != nil {
+		return nil, err
+	}
+	return &cookieProtectorImpl{secret: secret}, nil
+}
+
+// NewToken encodes data into a new token.
+func (s *cookieProtectorImpl) NewToken(data []byte) ([]byte, error) {
+	nonce := make([]byte, cookieNonceSize)
+	if _, err := rand.Read(nonce); err != nil {
+		return nil, err
+	}
+	aead, aeadNonce, err := s.createAEAD(nonce)
+	if err != nil {
+		return nil, err
+	}
+	return append(nonce, aead.Seal(nil, aeadNonce, data, nil)...), nil
+}
+
+// DecodeToken decodes a token.
+func (s *cookieProtectorImpl) DecodeToken(p []byte) ([]byte, error) {
+	if len(p) < cookieNonceSize {
+		return nil, fmt.Errorf("Token too short: %d", len(p))
+	}
+	nonce := p[:cookieNonceSize]
+	aead, aeadNonce, err := s.createAEAD(nonce)
+	if err != nil {
+		return nil, err
+	}
+	return aead.Open(nil, aeadNonce, p[cookieNonceSize:], nil)
+}
+
+func (s *cookieProtectorImpl) createAEAD(nonce []byte) (cipher.AEAD, []byte, error) {
+	h := hkdf.New(sha256.New, s.secret, nonce, []byte("quic-go cookie source"))
+	key := make([]byte, 32) // use a 32 byte key, in order to select AES-256
+	if _, err := io.ReadFull(h, key); err != nil {
+		return nil, nil, err
+	}
+	aeadNonce := make([]byte, 12)
+	if _, err := io.ReadFull(h, aeadNonce); err != nil {
+		return nil, nil, err
+	}
+	c, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, nil, err
+	}
+	aead, err := cipher.NewGCM(c)
+	if err != nil {
+		return nil, nil, err
+	}
+	return aead, aeadNonce, nil
+}
diff --git a/internal/handshake/cookie_protector_test.go b/internal/handshake/cookie_protector_test.go
new file mode 100644
index 000000000..fbc7c27cf
--- /dev/null
+++ b/internal/handshake/cookie_protector_test.go
@@ -0,0 +1,39 @@
+package handshake
+
+import (
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Cookie Protector", func() {
+	var cp cookieProtector
+
+	BeforeEach(func() {
+		var err error
+		cp, err = newCookieProtector()
+		Expect(err).ToNot(HaveOccurred())
+	})
+
+	It("encodes and decodes tokens", func() {
+		token, err := cp.NewToken([]byte("foobar"))
+		Expect(err).ToNot(HaveOccurred())
+		Expect(token).ToNot(ContainSubstring("foobar"))
+		decoded, err := cp.DecodeToken(token)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(decoded).To(Equal([]byte("foobar")))
+	})
+
+	It("fails deconding invalid tokens", func() {
+		token, err := cp.NewToken([]byte("foobar"))
+		Expect(err).ToNot(HaveOccurred())
+		token = token[1:] // remove the first byte
+		_, err = cp.DecodeToken(token)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("message authentication failed"))
+	})
+
+	It("errors when decoding too short tokens", func() {
+		_, err := cp.DecodeToken([]byte("foobar"))
+		Expect(err).To(MatchError("Token too short: 6"))
+	})
+})
diff --git a/internal/handshake/crypto_setup_server_test.go b/internal/handshake/crypto_setup_server_test.go
index e5833edff..7b4a0e94e 100644
--- a/internal/handshake/crypto_setup_server_test.go
+++ b/internal/handshake/crypto_setup_server_test.go
@@ -8,8 +8,6 @@ import (
 	"net"
 	"time"
 
-	"github.com/bifurcation/mint"
-
 	"github.com/lucas-clemente/quic-go/internal/crypto"
 	"github.com/lucas-clemente/quic-go/internal/mocks/crypto"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
@@ -99,7 +97,7 @@ type mockCookieProtector struct {
 	decodeErr error
 }
 
-var _ mint.CookieProtector = &mockCookieProtector{}
+var _ cookieProtector = &mockCookieProtector{}
 
 func (mockCookieProtector) NewToken(sourceAddr []byte) ([]byte, error) {
 	return append([]byte("token "), sourceAddr...), nil