From 5c8a598955cbf32f89e2dfb442ad5e8864c14d97 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Tue, 23 Feb 2021 10:55:22 +0800
Subject: [PATCH] allow using a wrapped net.TCPConn in NewConn

---
 conn_test.go | 15 +++++++++++++++
 rawconn.go   |  7 ++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/conn_test.go b/conn_test.go
index a3d6fa1..7a75ac7 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -46,3 +46,18 @@ func TestConn(t *testing.T) {
 		return tc, p.Conn, func() { tc.Close(); p.Conn.Close() }, nil
 	})
 }
+
+func TestWrappedConn(t *testing.T) {
+	type wrappedConn struct{ *net.TCPConn }
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	conn, err := net.DialTCP("tcp", nil, ln.Addr().(*net.TCPAddr))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err = tcp.NewConn(&wrappedConn{conn}); err != nil {
+		t.Fatalf("couldn't initialize from a wrapped conn: %v", err)
+	}
+}
diff --git a/rawconn.go b/rawconn.go
index 3ffd154..8e053fb 100644
--- a/rawconn.go
+++ b/rawconn.go
@@ -97,9 +97,14 @@ func (c *Conn) available() int {
 
 // NewConn returns a new end point.
 func NewConn(c net.Conn) (*Conn, error) {
+	type tcpConn interface {
+		SyscallConn() (syscall.RawConn, error)
+		SetLinger(int) error
+	}
+	var _ tcpConn = &net.TCPConn{}
 	cc := &Conn{Conn: c}
 	switch c := c.(type) {
-	case *net.TCPConn:
+	case tcpConn:
 		var err error
 		cc.c, err = c.SyscallConn()
 		if err != nil {