internal/socket: use interface not type assertions to infer a wrapped net.Conn's kind

When passing a wrapped connection to ipv{4,6}.NewPacketConn, calls to
ReadBatch will fail with an invalid connection error. This is because
NewConn performs a type assertion to determine the type (TCP, UDP, IP)
of the net.Conn.
In order to allow for wrapped connections to be passed to NewConn, we
can use an interface assertion that checks for methods on the net.Conn
that are unique to TCP, UDP and IP.

Fixes golang/go#42444.

Change-Id: I07cdb037829643f5cf4cd77105a772fcff8c4b2a
Reviewed-on: https://go-review.googlesource.com/c/net/+/271927
Reviewed-by: Emmanuel Odeke <[email protected]>
Reviewed-by: Ian Lance Taylor <[email protected]>
Run-TryBot: Emmanuel Odeke <[email protected]>
TryBot-Result: Go Bot <[email protected]>
Trust: Daniel Martí <[email protected]>

Author: marten-seemann

internal/socket/rawconn.go

@@ -17,18 +17,45 @@ type Conn struct {
 	c       syscall.RawConn
 }
 
+// tcpConn is an interface implemented by net.TCPConn.
+// It can be used for interface assertions to check if a net.Conn is a TCP connection.
+type tcpConn interface {
+	SyscallConn() (syscall.RawConn, error)
+	SetLinger(int) error
+}
+
+var _ tcpConn = (*net.TCPConn)(nil)
+
+// udpConn is an interface implemented by net.UDPConn.
+// It can be used for interface assertions to check if a net.Conn is a UDP connection.
+type udpConn interface {
+	SyscallConn() (syscall.RawConn, error)
+	ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
+}
+
+var _ udpConn = (*net.UDPConn)(nil)
+
+// ipConn is an interface implemented by net.IPConn.
+// It can be used for interface assertions to check if a net.Conn is an IP connection.
+type ipConn interface {
+	SyscallConn() (syscall.RawConn, error)
+	ReadMsgIP(b, oob []byte) (n, oobn, flags int, addr *net.IPAddr, err error)
+}
+
+var _ ipConn = (*net.IPConn)(nil)
+
 // NewConn returns a new raw connection.
 func NewConn(c net.Conn) (*Conn, error) {
 	var err error
 	var cc Conn
 	switch c := c.(type) {
-	case *net.TCPConn:
+	case tcpConn:
 		cc.network = "tcp"
 		cc.c, err = c.SyscallConn()
-	case *net.UDPConn:
+	case udpConn:
 		cc.network = "udp"
 		cc.c, err = c.SyscallConn()
-	case *net.IPConn:
+	case ipConn:
 		cc.network = "ip"
 		cc.c, err = c.SyscallConn()
 	default:

internal/socket/socket_test.go

@@ -164,7 +164,9 @@ func TestUDP(t *testing.T) {
 		t.Skipf("not supported on %s/%s: %v", runtime.GOOS, runtime.GOARCH, err)
 	}
 	defer c.Close()
-	cc, err := socket.NewConn(c.(net.Conn))
+	// test that wrapped connections work with NewConn too
+	type wrappedConn struct{ *net.UDPConn }
+	cc, err := socket.NewConn(&wrappedConn{c.(*net.UDPConn)})
 	if err != nil {
 		t.Fatal(err)
 	}