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) }