nclient4: fix conn race condition

Why did I ever think this would work??

Signed-off-by: Christopher Koch <chrisko@google.com>

3 files changed, 17 insertions, 22 deletions

diff --git a/dhcpv4/nclient4/client.go b/dhcpv4/nclient4/client.go
index 099be3f..9e8d1c6 100644
--- a/dhcpv4/nclient4/client.go
+++ b/dhcpv4/nclient4/client.go
@@ -99,23 +99,16 @@ func New(iface string, opts ...ClientOpt) (*Client, error) {
 	if err != nil {
 		return nil, err
 	}
-	c := NewWithConn(nil, i.HardwareAddr, opts...)
-
-	// Do this after so that a caller can still use a WithConn to override
-	// the connection.
-	if c.conn == nil {
-		pc, err := NewRawUDPConn(iface, ClientPort)
-		if err != nil {
-			return nil, err
-		}
-		c.conn = pc
+	pc, err := NewRawUDPConn(iface, ClientPort)
+	if err != nil {
+		return nil, err
 	}
-	return c, nil
+	return NewWithConn(pc, i.HardwareAddr, opts...)
 }
 
 // NewWithConn creates a new DHCP client that sends and receives packets on the
 // given interface.
-func NewWithConn(conn net.PacketConn, ifaceHWAddr net.HardwareAddr, opts ...ClientOpt) *Client {
+func NewWithConn(conn net.PacketConn, ifaceHWAddr net.HardwareAddr, opts ...ClientOpt) (*Client, error) {
 	c := &Client{
 		ifaceHWAddr: ifaceHWAddr,
 		timeout:     defaultTimeout,
@@ -132,9 +125,12 @@ func NewWithConn(conn net.PacketConn, ifaceHWAddr net.HardwareAddr, opts ...Clie
 		opt(c)
 	}
 
+	if c.conn == nil {
+		return nil, fmt.Errorf("no connection given")
+	}
 	c.wg.Add(1)
 	go c.receiveLoop()
-	return c
+	return c, nil
 }
 
 // Close closes the underlying connection.
@@ -244,13 +240,6 @@ func WithRetry(r int) ClientOpt {
 	}
 }
 
-// WithConn configures the packet connection to use.
-func WithConn(conn net.PacketConn) ClientOpt {
-	return func(c *Client) {
-		c.conn = conn
-	}
-}
-
 // WithServerAddr configures the address to send messages to.
 func WithServerAddr(n *net.UDPAddr) ClientOpt {
 	return func(c *Client) {
diff --git a/dhcpv4/nclient4/client_test.go b/dhcpv4/nclient4/client_test.go
index d3ea68b..6d7da1f 100644
--- a/dhcpv4/nclient4/client_test.go
+++ b/dhcpv4/nclient4/client_test.go
@@ -55,7 +55,10 @@ func serveAndClient(ctx context.Context, responses [][]*dhcpv4.DHCPv4, opts ...C
 
 	o := []ClientOpt{WithRetry(1), WithTimeout(2 * time.Second)}
 	o = append(o, opts...)
-	mc := NewWithConn(clientConn, net.HardwareAddr{0xa, 0xb, 0xc, 0xd, 0xe, 0xf}, o...)
+	mc, err := NewWithConn(clientConn, net.HardwareAddr{0xa, 0xb, 0xc, 0xd, 0xe, 0xf}, o...)
+	if err != nil {
+		panic(err)
+	}
 
 	h := &handler{responses: responses}
 	s, err := server4.NewServer(nil, h.handle, server4.WithConn(serverConn))
diff --git a/dhcpv4/server4/server_test.go b/dhcpv4/server4/server_test.go
index 68d0cc3..cd46774 100644
--- a/dhcpv4/server4/server_test.go
+++ b/dhcpv4/server4/server_test.go
@@ -83,7 +83,10 @@ func setUpClientAndServer(t *testing.T, iface net.Interface, handler Handler) (*
 	if err != nil {
 		t.Fatal(err)
 	}
-	c := nclient4.NewWithConn(clientConn, iface.HardwareAddr, nclient4.WithServerAddr(&saddr))
+	c, err := nclient4.NewWithConn(clientConn, iface.HardwareAddr, nclient4.WithServerAddr(&saddr))
+	if err != nil {
+		t.Fatal(err)
+	}
 	return c, s
 }