From 04921c29af2f8b8359245d5676c7ad4990e2a994 Mon Sep 17 00:00:00 2001
From: Lucas Clemente <lucas@clemente.io>
Date: Thu, 14 Apr 2016 19:39:49 +0200
Subject: [PATCH] introduce ConnectionID type

---
 example/main.go       | 3 ++-
 frame.go              | 1 +
 protocol/protocol.go  | 3 +++
 public_header.go      | 7 ++++---
 public_header_test.go | 6 +++---
 session.go            | 5 +++--
 6 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/example/main.go b/example/main.go
index eb1262b40..9a8a05254 100644
--- a/example/main.go
+++ b/example/main.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/lucas-clemente/quic-go"
 	"github.com/lucas-clemente/quic-go/crypto"
+	"github.com/lucas-clemente/quic-go/protocol"
 	"github.com/lucas-clemente/quic-go/utils"
 )
 
@@ -28,7 +29,7 @@ func main() {
 	serverConfig := quic.NewServerConfig(crypto.NewCurve25519KEX(), keyData)
 
 	// TODO: When should a session be created?
-	sessions := map[uint64]*quic.Session{}
+	sessions := map[protocol.ConnectionID]*quic.Session{}
 
 	addr, err := net.ResolveUDPAddr("udp", "localhost:6121")
 	if err != nil {
diff --git a/frame.go b/frame.go
index 2f6f0d230..4579b21ab 100644
--- a/frame.go
+++ b/frame.go
@@ -7,6 +7,7 @@ import (
 	"github.com/lucas-clemente/quic-go/utils"
 )
 
+// A Frame in QUIC
 type Frame interface {
 	Write(b *bytes.Buffer) error
 }
diff --git a/protocol/protocol.go b/protocol/protocol.go
index e869b8bef..429c3d067 100644
--- a/protocol/protocol.go
+++ b/protocol/protocol.go
@@ -2,3 +2,6 @@ package protocol
 
 // A PacketNumber in QUIC
 type PacketNumber uint64
+
+// A ConnectionID in QUIC
+type ConnectionID uint64
diff --git a/public_header.go b/public_header.go
index 416f2e6b5..4333c14a1 100644
--- a/public_header.go
+++ b/public_header.go
@@ -13,7 +13,7 @@ import (
 type PublicHeader struct {
 	VersionFlag  bool
 	ResetFlag    bool
-	ConnectionID uint64
+	ConnectionID protocol.ConnectionID
 	QuicVersion  uint32
 	PacketNumber protocol.PacketNumber
 	// packetNumberLen uint8
@@ -33,7 +33,7 @@ func (h *PublicHeader) WritePublicHeader(b *bytes.Buffer) error {
 	}
 
 	b.WriteByte(publicFlagByte)
-	utils.WriteUint64(b, h.ConnectionID)         // TODO: Send shorter connection id if possible
+	utils.WriteUint64(b, uint64(h.ConnectionID)) // TODO: Send shorter connection id if possible
 	utils.WriteUint32(b, uint32(h.PacketNumber)) // TODO: Send shorter packet number if possible
 	return nil
 }
@@ -71,10 +71,11 @@ func ParsePublicHeader(b io.ByteReader) (*PublicHeader, error) {
 	}
 
 	// Connection ID
-	header.ConnectionID, err = utils.ReadUintN(b, connectionIDLen)
+	connID, err := utils.ReadUintN(b, connectionIDLen)
 	if err != nil {
 		return nil, err
 	}
+	header.ConnectionID = protocol.ConnectionID(connID)
 	if header.ConnectionID == 0 {
 		return nil, errors.New("PublicHeader: connection ID cannot be 0")
 	}
diff --git a/public_header_test.go b/public_header_test.go
index 601121853..8defb9ebd 100644
--- a/public_header_test.go
+++ b/public_header_test.go
@@ -17,7 +17,7 @@ var _ = Describe("Public Header", func() {
 			Expect(err).ToNot(HaveOccurred())
 			Expect(publicHeader.VersionFlag).To(BeTrue())
 			Expect(publicHeader.ResetFlag).To(BeFalse())
-			Expect(publicHeader.ConnectionID).To(Equal(uint64(0x4cfa9f9b668619f6)))
+			Expect(publicHeader.ConnectionID).To(Equal(protocol.ConnectionID(0x4cfa9f9b668619f6)))
 			Expect(publicHeader.QuicVersion).To(Equal(binary.BigEndian.Uint32([]byte("Q030"))))
 			Expect(publicHeader.PacketNumber).To(Equal(protocol.PacketNumber(1)))
 			Expect(b.Len()).To(BeZero())
@@ -28,7 +28,7 @@ var _ = Describe("Public Header", func() {
 			publicHeader, err := ParsePublicHeader(b)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(publicHeader.VersionFlag).To(BeFalse())
-			Expect(publicHeader.ConnectionID).To(Equal(uint64(0x4cfa9f9b)))
+			Expect(publicHeader.ConnectionID).To(Equal(protocol.ConnectionID(0x4cfa9f9b)))
 			Expect(b.Len()).To(BeZero())
 		})
 
@@ -37,7 +37,7 @@ var _ = Describe("Public Header", func() {
 			publicHeader, err := ParsePublicHeader(b)
 			Expect(err).ToNot(HaveOccurred())
 			Expect(publicHeader.VersionFlag).To(BeFalse())
-			Expect(publicHeader.ConnectionID).To(Equal(uint64(0x4c)))
+			Expect(publicHeader.ConnectionID).To(Equal(protocol.ConnectionID(0x4c)))
 			Expect(b.Len()).To(BeZero())
 		})
 
diff --git a/session.go b/session.go
index 67a77abd9..02e8eb50f 100644
--- a/session.go
+++ b/session.go
@@ -6,11 +6,12 @@ import (
 	"net"
 
 	"github.com/lucas-clemente/quic-go/crypto"
+	"github.com/lucas-clemente/quic-go/protocol"
 )
 
 // A Session is a QUIC session
 type Session struct {
-	ConnectionID uint64
+	ConnectionID protocol.ConnectionID
 	ServerConfig *ServerConfig
 
 	Connection        *net.UDPConn
@@ -20,7 +21,7 @@ type Session struct {
 }
 
 // NewSession makes a new session
-func NewSession(conn *net.UDPConn, connectionID uint64, sCfg *ServerConfig) *Session {
+func NewSession(conn *net.UDPConn, connectionID protocol.ConnectionID, sCfg *ServerConfig) *Session {
 	return &Session{
 		Connection:   conn,
 		ConnectionID: connectionID,