From 1760bc32279f623dc68a2097fc8b787bda0ac7d2 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Thu, 14 Apr 2016 20:39:51 +0200
Subject: [PATCH] add key derivation and change AEAD in session, not working
 right now :(

---
 crypto/key_derivation.go | 43 ++++++++++++++++++++++++++++++++++++++++
 handshake/tags.go        |  3 +++
 session.go               | 34 ++++++++++++++++++++++++-------
 3 files changed, 73 insertions(+), 7 deletions(-)
 create mode 100644 crypto/key_derivation.go

diff --git a/crypto/key_derivation.go b/crypto/key_derivation.go
new file mode 100644
index 00000000..0268b82d
--- /dev/null
+++ b/crypto/key_derivation.go
@@ -0,0 +1,43 @@
+package crypto
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"io"
+
+	"github.com/lucas-clemente/quic-go/protocol"
+	"github.com/lucas-clemente/quic-go/utils"
+
+	"golang.org/x/crypto/hkdf"
+)
+
+// DeriveKeysAESGCM derives the client and server keys and creates a matching AES-GCM instance
+func DeriveKeysAESGCM(sharedSecret, nonces []byte, connID protocol.ConnectionID, chlo []byte, scfg []byte) (AEAD, error) {
+	var info bytes.Buffer
+	info.Write([]byte("QUIC key expansion\x00"))
+	utils.WriteUint64(&info, uint64(connID))
+	info.Write(chlo)
+	info.Write(scfg)
+
+	r := hkdf.New(sha256.New, sharedSecret, nonces, info.Bytes())
+
+	otherKey := make([]byte, 16)
+	myKey := make([]byte, 16)
+	otherIV := make([]byte, 4)
+	myIV := make([]byte, 4)
+
+	if _, err := io.ReadFull(r, otherKey); err != nil {
+		return nil, err
+	}
+	if _, err := io.ReadFull(r, myKey); err != nil {
+		return nil, err
+	}
+	if _, err := io.ReadFull(r, otherIV); err != nil {
+		return nil, err
+	}
+	if _, err := io.ReadFull(r, myIV); err != nil {
+		return nil, err
+	}
+
+	return NewAEADAESGCM(otherKey, myKey, otherIV, myIV)
+}
diff --git a/handshake/tags.go b/handshake/tags.go
index 1dbc821f..38daebc3 100644
--- a/handshake/tags.go
+++ b/handshake/tags.go
@@ -51,6 +51,9 @@ const (
 	// TagPROF is the server proof
 	TagPROF Tag = 'P' + 'R'<<8 + 'O'<<16 + 'F'<<24
 
+	// TagNONC is the client nonce
+	TagNONC Tag = 'N' + 'O'<<8 + 'N'<<16 + 'C'<<24
+
 	// TagSCID is the server config ID
 	TagSCID Tag = 'S' + 'C'<<8 + 'I'<<16 + 'D'<<24
 	// TagKEXS is the list of key exchange algos
diff --git a/session.go b/session.go
index ac6ee4bd..da42bb5f 100644
--- a/session.go
+++ b/session.go
@@ -19,7 +19,11 @@ type Session struct {
 	Connection        *net.UDPConn
 	CurrentRemoteAddr *net.UDPAddr
 
+	aead crypto.AEAD
+
 	Entropy EntropyAccumulator
+
+	lastSentPacketNumber protocol.PacketNumber
 }
 
 // NewSession makes a new session
@@ -28,6 +32,7 @@ func NewSession(conn *net.UDPConn, connectionID protocol.ConnectionID, sCfg *Ser
 		Connection:   conn,
 		ConnectionID: connectionID,
 		ServerConfig: sCfg,
+		aead:         &crypto.NullAEAD{},
 	}
 }
 
@@ -38,8 +43,7 @@ func (s *Session) HandlePacket(addr *net.UDPAddr, publicHeaderBinary []byte, pub
 		s.CurrentRemoteAddr = addr
 	}
 
-	nullAEAD := &crypto.NullAEAD{}
-	r, err := nullAEAD.Open(0, publicHeaderBinary, r)
+	r, err := s.aead.Open(publicHeader.PacketNumber, publicHeaderBinary, r)
 	if err != nil {
 		return err
 	}
@@ -72,8 +76,22 @@ func (s *Session) HandlePacket(addr *net.UDPAddr, publicHeaderBinary []byte, pub
 		return errors.New("Session: expected CHLO")
 	}
 
-	if _, ok := cryptoData[handshake.TagSCFG]; ok {
-		return errors.New("Session: received CHLO with PUBS")
+	if _, ok := cryptoData[handshake.TagSCID]; ok {
+		var sharedSecret []byte
+		sharedSecret, err = s.ServerConfig.kex.CalculateSharedKey(cryptoData[handshake.TagPUBS])
+		if err != nil {
+			return err
+		}
+		s.aead, err = crypto.DeriveKeysAESGCM(sharedSecret, cryptoData[handshake.TagNONC], s.ConnectionID, frame.Data, s.ServerConfig.Get())
+		if err != nil {
+			return err
+		}
+		fmt.Println("Got common secret")
+		s.SendFrames([]Frame{&AckFrame{
+			Entropy:         s.Entropy.Get(),
+			LargestObserved: 2,
+		}})
+		return nil
 	}
 
 	proof, err := s.ServerConfig.Sign(frame.Data)
@@ -111,14 +129,16 @@ func (s *Session) SendFrames(frames []Frame) error {
 		}
 	}
 
+	s.lastSentPacketNumber++
+
 	var fullReply bytes.Buffer
-	responsePublicHeader := PublicHeader{ConnectionID: s.ConnectionID, PacketNumber: 1}
-	fmt.Printf("%#v\n", responsePublicHeader)
+	responsePublicHeader := PublicHeader{ConnectionID: s.ConnectionID, PacketNumber: s.lastSentPacketNumber}
+	fmt.Printf("Sending packet # %d\n", responsePublicHeader.PacketNumber)
 	if err := responsePublicHeader.WritePublicHeader(&fullReply); err != nil {
 		return err
 	}
 
-	(&crypto.NullAEAD{}).Seal(0, &fullReply, fullReply.Bytes(), framesData.Bytes())
+	s.aead.Seal(s.lastSentPacketNumber, &fullReply, fullReply.Bytes(), framesData.Bytes())
 
 	_, err := s.Connection.WriteToUDP(fullReply.Bytes(), s.CurrentRemoteAddr)
 	return err