udp: Add constructor for ConnectionIDGenerator, return buffer to pool on socket timeout

This commit is contained in:
Justin Li 2015-02-22 16:58:43 -05:00
parent 7512f50731
commit 3d28f281fb
3 changed files with 17 additions and 15 deletions

View file

@ -18,20 +18,24 @@ type ConnectionIDGenerator struct {
block cipher.Block block cipher.Block
} }
// Init generates the AES key and sets up the first initialization vector. // NewConnectionIDGenerator creates a ConnectionIDGenerator and generates its
func (g *ConnectionIDGenerator) Init() error { // AES key and first initialization vector.
func NewConnectionIDGenerator() (gen *ConnectionIDGenerator, err error) {
gen = &ConnectionIDGenerator{}
key := make([]byte, 16) key := make([]byte, 16)
_, err := rand.Read(key)
_, err = rand.Read(key)
if err != nil { if err != nil {
return err return
} }
g.block, err = aes.NewCipher(key) gen.block, err = aes.NewCipher(key)
if err != nil { if err != nil {
return err return
} }
return g.NewIV() err = gen.NewIV()
return
} }
// Generate returns the 64-bit connection ID for an IP // Generate returns the 64-bit connection ID for an IP

View file

@ -11,15 +11,13 @@ import (
) )
func TestInitReturnsNoError(t *testing.T) { func TestInitReturnsNoError(t *testing.T) {
gen := &ConnectionIDGenerator{} if _, err := NewConnectionIDGenerator(); err != nil {
if err := gen.Init(); err != nil {
t.Error("Init returned", err) t.Error("Init returned", err)
} }
} }
func testGenerateConnectionID(t *testing.T, ip net.IP) { func testGenerateConnectionID(t *testing.T, ip net.IP) {
gen := &ConnectionIDGenerator{} gen, _ := NewConnectionIDGenerator()
gen.Init()
id1 := gen.Generate(ip) id1 := gen.Generate(ip)
id2 := gen.Generate(ip) id2 := gen.Generate(ip)
@ -46,8 +44,7 @@ func TestGenerateConnectionIDIPv6(t *testing.T) {
} }
func TestMatchesWorksWithPreviousIV(t *testing.T) { func TestMatchesWorksWithPreviousIV(t *testing.T) {
gen := &ConnectionIDGenerator{} gen, _ := NewConnectionIDGenerator()
gen.Init()
ip := net.ParseIP("192.168.1.123").To4() ip := net.ParseIP("192.168.1.123").To4()
id1 := gen.Generate(ip) id1 := gen.Generate(ip)

View file

@ -62,6 +62,7 @@ func (s *Server) serve(listenAddr string) error {
if err != nil { if err != nil {
if netErr, ok := err.(net.Error); ok && netErr.Temporary() { if netErr, ok := err.(net.Error); ok && netErr.Temporary() {
pool.GiveSlice(buffer)
continue continue
} }
return err return err
@ -113,8 +114,8 @@ func (s *Server) Stop() {
// NewServer returns a new UDP server for a given configuration and tracker. // NewServer returns a new UDP server for a given configuration and tracker.
func NewServer(cfg *config.Config, tkr *tracker.Tracker) *Server { func NewServer(cfg *config.Config, tkr *tracker.Tracker) *Server {
gen := &ConnectionIDGenerator{} gen, err := NewConnectionIDGenerator()
if err := gen.Init(); err != nil { if err != nil {
panic(err) panic(err)
} }