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

View file

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

View file

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