chore: refine code

client_test.go

@@ -261,7 +261,7 @@ func (s *testClientServer) OnTraffic(c Conn) (action Action) {
 		require.Equal(s.tester, string(pingMsg), string(ping), "bad header")
 	}
 	v := c.Context()
-	if v != nil && s.network != "udp" {
+	if v != nil {
 		v.(*sync.Once).Do(readHeader)
 	}
 
@@ -275,7 +275,7 @@ func (s *testClientServer) OnTraffic(c Conn) (action Action) {
 			_ = c.OutboundBuffered()
 			_, _ = c.Discard(1)
 		}
-		if s.network == "udp" && bytes.Equal(buf.Bytes(), pingMsg) {
+		if v == nil && bytes.Equal(buf.Bytes(), pingMsg) {
 			atomic.AddInt32(&s.udpReadHeader, 1)
 			buf.Reset()
 		}
@@ -290,7 +290,7 @@ func (s *testClientServer) OnTraffic(c Conn) (action Action) {
 	}
 
 	buf, _ := c.Next(-1)
-	if s.network == "udp" && bytes.Equal(buf, pingMsg) {
+	if v == nil && bytes.Equal(buf, pingMsg) {
 		atomic.AddInt32(&s.udpReadHeader, 1)
 		buf = nil
 	}

client_unix.go

@@ -194,7 +194,7 @@ func (cli *Client) EnrollContext(c net.Conn, ctx interface{}) (Conn, error) {
 
 	var (
 		sockAddr unix.Sockaddr
-		gc       Conn
+		gc       *conn
 	)
 	switch c.(type) {
 	case *net.UnixConn:
@@ -227,7 +227,7 @@ func (cli *Client) EnrollContext(c net.Conn, ctx interface{}) (Conn, error) {
 	default:
 		return nil, errorx.ErrUnsupportedProtocol
 	}
-	gc.SetContext(ctx)
+	gc.ctx = ctx
 
 	connOpened := make(chan struct{})
 	ccb := &connWithCallback{c: gc, cb: func() {
@@ -238,6 +238,7 @@ func (cli *Client) EnrollContext(c net.Conn, ctx interface{}) (Conn, error) {
 		gc.Close()
 		return nil, err
 	}
+
 	<-connOpened
 	return gc, nil
 }

eventloop_unix.go

@@ -62,15 +62,15 @@ func (el *eventloop) closeConns() {
 }
 
 type connWithCallback struct {
-	c  Conn
+	c  *conn
 	cb func()
 }
 
 func (el *eventloop) register(itf interface{}) error {
 	c, ok := itf.(*conn)
 	if !ok {
 		ccb := itf.(*connWithCallback)
-		c = ccb.c.(*conn)
+		c = ccb.c
 		defer ccb.cb()
 	}