From 1097698c4b52b4b78143b3148ca70bf5347c753c Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sun, 17 Apr 2016 14:21:36 +0700
Subject: [PATCH] send entropy in packets, validate entropy in received ACKs

---
 session.go          | 41 ++++++++++++++++++++++++++++++++++-------
 utils/utils.go      | 15 +++++++++++++++
 utils/utils_test.go |  8 ++++++++
 3 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/session.go b/session.go
index e5774662..eaa4d3ad 100644
--- a/session.go
+++ b/session.go
@@ -9,6 +9,7 @@ import (
 	"github.com/lucas-clemente/quic-go/frames"
 	"github.com/lucas-clemente/quic-go/handshake"
 	"github.com/lucas-clemente/quic-go/protocol"
+	"github.com/lucas-clemente/quic-go/utils"
 )
 
 // StreamCallback gets a stream frame and returns a reply frame
@@ -25,7 +26,9 @@ type Session struct {
 	ServerConfig *handshake.ServerConfig
 	cryptoSetup  *handshake.CryptoSetup
 
-	Entropy EntropyAccumulator
+	EntropyReceived EntropyAccumulator
+	EntropySent     EntropyAccumulator
+	EntropyHistory  map[protocol.PacketNumber]EntropyAccumulator // ToDo: store this with the packet itself
 
 	lastSentPacketNumber     protocol.PacketNumber
 	lastObservedPacketNumber protocol.PacketNumber
@@ -48,6 +51,7 @@ func NewSession(conn *net.UDPConn, v protocol.VersionNumber, connectionID protoc
 		streamCallback:           streamCallback,
 		lastObservedPacketNumber: 0,
 		Streams:                  make(map[protocol.StreamID]*Stream),
+		EntropyHistory:           make(map[protocol.PacketNumber]EntropyAccumulator),
 	}
 }
 
@@ -77,11 +81,11 @@ func (s *Session) HandlePacket(addr *net.UDPAddr, publicHeaderBinary []byte, pub
 	if err != nil {
 		return err
 	}
-	s.Entropy.Add(publicHeader.PacketNumber, privateFlag&0x01 > 0)
+	s.EntropyReceived.Add(publicHeader.PacketNumber, privateFlag&0x01 > 0)
 
 	s.SendFrames([]frames.Frame{&frames.AckFrame{
 		LargestObserved: publicHeader.PacketNumber,
-		Entropy:         s.Entropy.Get(),
+		Entropy:         s.EntropyReceived.Get(),
 	}})
 
 	// read all frames in the packet
@@ -173,10 +177,21 @@ func (s *Session) handleStreamFrame(r *bytes.Reader) error {
 }
 
 func (s *Session) handleAckFrame(r *bytes.Reader) error {
-	_, err := frames.ParseAckFrame(r)
+	frame, err := frames.ParseAckFrame(r)
 	if err != nil {
 		return err
 	}
+
+	expectedEntropy, ok := s.EntropyHistory[frame.LargestObserved]
+	if !ok {
+		return errors.New("No entropy value saved for received ACK packet")
+	}
+	delete(s.EntropyHistory, frame.LargestObserved)
+
+	if byte(expectedEntropy) != frame.Entropy {
+		return errors.New("Incorrect entropy value in ACK package")
+	}
+
 	return nil
 }
 
@@ -210,7 +225,16 @@ func (s *Session) handleRstStreamFrame(r *bytes.Reader) error {
 // SendFrames sends a number of frames to the client
 func (s *Session) SendFrames(frames []frames.Frame) error {
 	var framesData bytes.Buffer
-	framesData.WriteByte(0) // TODO: entropy
+	entropyBit, err := utils.RandomBit()
+	if err != nil {
+		return err
+	}
+	if entropyBit {
+		framesData.WriteByte(1)
+	} else {
+		framesData.WriteByte(0)
+	}
+
 	for _, f := range frames {
 		if err := f.Write(&framesData); err != nil {
 			return err
@@ -220,14 +244,17 @@ func (s *Session) SendFrames(frames []frames.Frame) error {
 	s.lastSentPacketNumber++
 
 	var fullReply bytes.Buffer
-	responsePublicHeader := PublicHeader{ConnectionID: s.ConnectionID, PacketNumber: s.lastSentPacketNumber}
+	packetNumber := s.lastSentPacketNumber
+	responsePublicHeader := PublicHeader{ConnectionID: s.ConnectionID, PacketNumber: packetNumber}
 	if err := responsePublicHeader.WritePublicHeader(&fullReply); err != nil {
 		return err
 	}
+	s.EntropySent.Add(packetNumber, entropyBit)
+	s.EntropyHistory[packetNumber] = s.EntropySent
 
 	s.cryptoSetup.Seal(s.lastSentPacketNumber, &fullReply, fullReply.Bytes(), framesData.Bytes())
 
 	fmt.Printf("-> Sending packet %d (%d bytes) to %v\n", responsePublicHeader.PacketNumber, len(fullReply.Bytes()), s.CurrentRemoteAddr)
-	_, err := s.Connection.WriteToUDP(fullReply.Bytes(), s.CurrentRemoteAddr)
+	_, err = s.Connection.WriteToUDP(fullReply.Bytes(), s.CurrentRemoteAddr)
 	return err
 }
diff --git a/utils/utils.go b/utils/utils.go
index d96d8370..4b988193 100644
--- a/utils/utils.go
+++ b/utils/utils.go
@@ -2,6 +2,7 @@ package utils
 
 import (
 	"bytes"
+	"crypto/rand"
 	"io"
 )
 
@@ -123,6 +124,20 @@ func Min(a, b int) int {
 	return b
 }
 
+// RandomBit returns a cryptographically secure random bit (encoded as true / false)
+func RandomBit() (bool, error) {
+	// ToDo: it's probably more efficient to read a bigger slice of random numbers at once and to cache them somewhere
+	b := make([]byte, 1)
+	_, err := rand.Read(b)
+	if err != nil {
+		return false, err
+	}
+	if uint8(b[0])%2 == 0 {
+		return false, nil
+	}
+	return true, nil
+}
+
 // Uint32Slice attaches the methods of sort.Interface to []uint32, sorting in increasing order.
 type Uint32Slice []uint32
 
diff --git a/utils/utils_test.go b/utils/utils_test.go
index d105e747..ac20e929 100644
--- a/utils/utils_test.go
+++ b/utils/utils_test.go
@@ -110,6 +110,14 @@ var _ = Describe("Utils", func() {
 		})
 	})
 
+	Context("Rand", func() {
+		It("returns either true or false", func() {
+			val, err := RandomBit()
+			Expect(err).ToNot(HaveOccurred())
+			Expect(val).To(SatisfyAny(Equal(true), Equal(false)))
+		})
+	})
+
 	Context("ReadUintN", func() {
 		It("reads n bytes", func() {
 			m := map[uint8]uint64{