add interface xnet.SrcAddr/DstAddr
This commit is contained in:
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user