From 3d28f281fb6551d2c39250abcebbdbc8eef3a571 Mon Sep 17 00:00:00 2001
From: Justin Li <jli.justinli@gmail.com>
Date: Sun, 22 Feb 2015 16:58:43 -0500
Subject: [PATCH] udp: Add constructor for ConnectionIDGenerator, return buffer
 to pool on socket timeout

---
 udp/connection.go      | 18 +++++++++++-------
 udp/connection_test.go |  9 +++------
 udp/udp.go             |  5 +++--
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/udp/connection.go b/udp/connection.go
index 49b2c1a..27b497a 100644
--- a/udp/connection.go
+++ b/udp/connection.go
@@ -18,20 +18,24 @@ type ConnectionIDGenerator struct {
 	block   cipher.Block
 }
 
-// Init generates the AES key and sets up the first initialization vector.
-func (g *ConnectionIDGenerator) Init() error {
+// NewConnectionIDGenerator creates a ConnectionIDGenerator and generates its
+// AES key and first initialization vector.
+func NewConnectionIDGenerator() (gen *ConnectionIDGenerator, err error) {
+	gen = &ConnectionIDGenerator{}
 	key := make([]byte, 16)
-	_, err := rand.Read(key)
+
+	_, err = rand.Read(key)
 	if err != nil {
-		return err
+		return
 	}
 
-	g.block, err = aes.NewCipher(key)
+	gen.block, err = aes.NewCipher(key)
 	if err != nil {
-		return err
+		return
 	}
 
-	return g.NewIV()
+	err = gen.NewIV()
+	return
 }
 
 // Generate returns the 64-bit connection ID for an IP
diff --git a/udp/connection_test.go b/udp/connection_test.go
index aebccbc..e4ad12b 100644
--- a/udp/connection_test.go
+++ b/udp/connection_test.go
@@ -11,15 +11,13 @@ import (
 )
 
 func TestInitReturnsNoError(t *testing.T) {
-	gen := &ConnectionIDGenerator{}
-	if err := gen.Init(); err != nil {
+	if _, err := NewConnectionIDGenerator(); err != nil {
 		t.Error("Init returned", err)
 	}
 }
 
 func testGenerateConnectionID(t *testing.T, ip net.IP) {
-	gen := &ConnectionIDGenerator{}
-	gen.Init()
+	gen, _ := NewConnectionIDGenerator()
 
 	id1 := gen.Generate(ip)
 	id2 := gen.Generate(ip)
@@ -46,8 +44,7 @@ func TestGenerateConnectionIDIPv6(t *testing.T) {
 }
 
 func TestMatchesWorksWithPreviousIV(t *testing.T) {
-	gen := &ConnectionIDGenerator{}
-	gen.Init()
+	gen, _ := NewConnectionIDGenerator()
 	ip := net.ParseIP("192.168.1.123").To4()
 
 	id1 := gen.Generate(ip)
diff --git a/udp/udp.go b/udp/udp.go
index 8609eaa..890dc71 100644
--- a/udp/udp.go
+++ b/udp/udp.go
@@ -62,6 +62,7 @@ func (s *Server) serve(listenAddr string) error {
 
 		if err != nil {
 			if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
+				pool.GiveSlice(buffer)
 				continue
 			}
 			return err
@@ -113,8 +114,8 @@ func (s *Server) Stop() {
 
 // NewServer returns a new UDP server for a given configuration and tracker.
 func NewServer(cfg *config.Config, tkr *tracker.Tracker) *Server {
-	gen := &ConnectionIDGenerator{}
-	if err := gen.Init(); err != nil {
+	gen, err := NewConnectionIDGenerator()
+	if err != nil {
 		panic(err)
 	}