add interface xnet.SrcAddr/DstAddr

internal/net/net.go

@@ -26,8 +26,12 @@ type RemoteAddr interface {
 	RemoteAddr() net.Addr
 }
 
-type ClientAddr interface {
-	ClientAddr() net.Addr
+type SrcAddr interface {
+	SrcAddr() net.Addr
+}
+
+type DstAddr interface {
+	DstAddr() net.Addr
 }
 
 // tcpraw.TCPConn

internal/net/pipe.go

@@ -69,13 +69,13 @@ func pipeBuffer(dst io.ReadWriteCloser, src io.ReadWriteCloser, bufferSize int)
 	if cw, ok := dst.(xio.CloseWrite); ok {
 		if e := cw.CloseWrite(); e == xio.ErrUnsupported {
 			dst.Close()
+		} else {
+			// Set TCP half-close timeout.
+			xio.SetReadDeadline(dst, time.Now().Add(tcpWaitTimeout))
 		}
 	} else {
 		dst.Close()
 	}
 
-	// Set TCP half-close timeout.
-	xio.SetReadDeadline(dst, time.Now().Add(tcpWaitTimeout))
-
 	return err
 }

internal/net/proxyproto/conn.go

@@ -2,6 +2,7 @@ package proxyproto
 
 import (
 	"net"
+	"strconv"
 
 	xio "github.com/go-gost/x/internal/io"
 	proxyproto "github.com/pires/go-proxyproto"
@@ -11,6 +12,28 @@ type serverConn struct {
 	net.Conn
 }
 
+func (c *serverConn) RemoteAddr() net.Addr {
+	if conn, ok := c.Conn.(*proxyproto.Conn); ok {
+		return conn.Raw().RemoteAddr()
+	}
+	return c.Conn.RemoteAddr()
+}
+
+func (c *serverConn) LocalAddr() net.Addr {
+	if conn, ok := c.Conn.(*proxyproto.Conn); ok {
+		return conn.Raw().LocalAddr()
+	}
+	return c.Conn.LocalAddr()
+}
+
+func (c *serverConn) SrcAddr() net.Addr {
+	return c.Conn.RemoteAddr()
+}
+
+func (c *serverConn) DstAddr() net.Addr {
+	return c.Conn.LocalAddr()
+}
+
 func (c *serverConn) CloseRead() error {
 	if sc, ok := c.Conn.(xio.CloseRead); ok {
 		return sc.CloseRead()
@@ -36,7 +59,14 @@ func (c *serverConn) CloseWrite() error {
 }
 
 func WrapClientConn(ppv int, src, dst net.Addr, c net.Conn) net.Conn {
-	if ppv <= 0 {
+	if ppv <= 0 || c == nil {
+		return c
+	}
+
+	if src = convertAddr(src); src == nil {
+		return c
+	}
+	if dst = convertAddr(dst); dst == nil {
 		return c
 	}
 
@@ -44,3 +74,31 @@ func WrapClientConn(ppv int, src, dst net.Addr, c net.Conn) net.Conn {
 	header.WriteTo(c)
 	return c
 }
+
+func convertAddr(addr net.Addr) net.Addr {
+	if addr == nil {
+		return nil
+	}
+
+	host, sp, _ := net.SplitHostPort(addr.String())
+	ip := net.ParseIP(host)
+	port, _ := strconv.Atoi(sp)
+
+	if ip == nil || ip.Equal(net.IPv6zero) {
+		ip = net.IPv4zero
+	}
+
+	switch addr.Network() {
+	case "tcp", "tcp4", "tcp6":
+		return &net.TCPAddr{
+			IP:   ip,
+			Port: port,
+		}
+
+	default:
+		return &net.UDPAddr{
+			IP:   ip,
+			Port: port,
+		}
+	}
+}

internal/util/ws/ws.go

@@ -5,6 +5,7 @@ import (
 	"sync"
 	"time"
 
+	xio "github.com/go-gost/x/internal/io"
 	xnet "github.com/go-gost/x/internal/net"
 	"github.com/gorilla/websocket"
 )
@@ -13,14 +14,16 @@ type WebsocketConn interface {
 	net.Conn
 	WriteMessage(int, []byte) error
 	ReadMessage() (int, []byte, error)
-	xnet.ClientAddr
+	xnet.SrcAddr
+	xio.CloseRead
+	xio.CloseWrite
 }
 
 type websocketConn struct {
 	*websocket.Conn
-	rb         []byte
-	clientAddr net.Addr
-	mux        sync.Mutex
+	rb      []byte
+	srcAddr net.Addr
+	mux     sync.Mutex
 }
 
 func Conn(conn *websocket.Conn) WebsocketConn {
@@ -29,10 +32,10 @@ func Conn(conn *websocket.Conn) WebsocketConn {
 	}
 }
 
-func ConnWithClientAddr(conn *websocket.Conn, clientAddr net.Addr) WebsocketConn {
+func ConnWithSrcAddr(conn *websocket.Conn, srcAddr net.Addr) WebsocketConn {
 	return &websocketConn{
-		Conn:       conn,
-		clientAddr: clientAddr,
+		Conn:    conn,
+		srcAddr: srcAddr,
 	}
 }
 
@@ -77,6 +80,20 @@ func (c *websocketConn) SetWriteDeadline(t time.Time) error {
 	return c.Conn.SetWriteDeadline(t)
 }
 
-func (c *websocketConn) ClientAddr() net.Addr {
-	return c.clientAddr
+func (c *websocketConn) SrcAddr() net.Addr {
+	return c.srcAddr
+}
+
+func (c *websocketConn) CloseRead() error {
+	if sc, ok := c.Conn.NetConn().(xio.CloseRead); ok {
+		return sc.CloseRead()
+	}
+	return xio.ErrUnsupported
+}
+
+func (c *websocketConn) CloseWrite() error {
+	if sc, ok := c.Conn.NetConn().(xio.CloseWrite); ok {
+		return sc.CloseWrite()
+	}
+	return xio.ErrUnsupported
 }

internal/util/wt/conn.go

@@ -8,9 +8,9 @@ import (
 )
 
 type conn struct {
-	session    *wt.Session
-	stream     *wt.Stream
-	clientAddr net.Addr
+	session *wt.Session
+	stream  *wt.Stream
+	srcAddr net.Addr
 }
 
 func Conn(session *wt.Session, stream *wt.Stream) net.Conn {
@@ -20,11 +20,11 @@ func Conn(session *wt.Session, stream *wt.Stream) net.Conn {
 	}
 }
 
-func ConnWithClientAddr(session *wt.Session, stream *wt.Stream, clientAddr net.Addr) net.Conn {
+func ConnWithSrcAddr(session *wt.Session, stream *wt.Stream, srcAddr net.Addr) net.Conn {
 	return &conn{
-		session:    session,
-		stream:     stream,
-		clientAddr: clientAddr,
+		session: session,
+		stream:  stream,
+		srcAddr: srcAddr,
 	}
 }
 
@@ -48,8 +48,8 @@ func (c *conn) RemoteAddr() net.Addr {
 	return c.session.RemoteAddr()
 }
 
-func (c *conn) ClientAddr() net.Addr {
-	return c.clientAddr
+func (c *conn) SrcAddr() net.Addr {
+	return c.srcAddr
 }
 
 func (c *conn) SetDeadline(t time.Time) error {