Improve bufio.NATPacketConn

This commit is contained in:
dyhkwong 2024-04-10 20:15:01 +08:00 committed by 世界
parent 8b68fc4d7a
commit f9c59e9940
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
2 changed files with 96 additions and 18 deletions

View file

@ -17,16 +17,16 @@ type NATPacketConn interface {
func NewUnidirectionalNATPacketConn(conn N.NetPacketConn, origin M.Socksaddr, destination M.Socksaddr) NATPacketConn {
return &unidirectionalNATPacketConn{
NetPacketConn: conn,
origin: origin,
destination: destination,
origin: socksaddrWithoutPort(origin),
destination: socksaddrWithoutPort(destination),
}
}
func NewNATPacketConn(conn N.NetPacketConn, origin M.Socksaddr, destination M.Socksaddr) NATPacketConn {
return &bidirectionalNATPacketConn{
NetPacketConn: conn,
origin: origin,
destination: destination,
origin: socksaddrWithoutPort(origin),
destination: socksaddrWithoutPort(destination),
}
}
@ -37,15 +37,24 @@ type unidirectionalNATPacketConn struct {
}
func (c *unidirectionalNATPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
if M.SocksaddrFromNet(addr) == c.destination {
addr = c.origin.UDPAddr()
destination := M.SocksaddrFromNet(addr)
if socksaddrWithoutPort(destination) == c.destination {
destination = M.Socksaddr{
Addr: c.origin.Addr,
Fqdn: c.origin.Fqdn,
Port: destination.Port,
}
}
return c.NetPacketConn.WriteTo(p, addr)
return c.NetPacketConn.WriteTo(p, destination.UDPAddr())
}
func (c *unidirectionalNATPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
if destination == c.destination {
destination = c.origin
if socksaddrWithoutPort(destination) == c.destination {
destination = M.Socksaddr{
Addr: c.origin.Addr,
Fqdn: c.origin.Fqdn,
Port: destination.Port,
}
}
return c.NetPacketConn.WritePacket(buffer, destination)
}
@ -66,30 +75,55 @@ type bidirectionalNATPacketConn struct {
func (c *bidirectionalNATPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
n, addr, err = c.NetPacketConn.ReadFrom(p)
if err == nil && M.SocksaddrFromNet(addr) == c.origin {
addr = c.destination.UDPAddr()
if err != nil {
return
}
destination := M.SocksaddrFromNet(addr)
if socksaddrWithoutPort(destination) == c.origin {
destination = M.Socksaddr{
Addr: c.destination.Addr,
Fqdn: c.destination.Fqdn,
Port: destination.Port,
}
}
addr = destination.UDPAddr()
return
}
func (c *bidirectionalNATPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
if M.SocksaddrFromNet(addr) == c.destination {
addr = c.origin.UDPAddr()
destination := M.SocksaddrFromNet(addr)
if socksaddrWithoutPort(destination) == c.destination {
destination = M.Socksaddr{
Addr: c.origin.Addr,
Fqdn: c.origin.Fqdn,
Port: destination.Port,
}
}
return c.NetPacketConn.WriteTo(p, addr)
return c.NetPacketConn.WriteTo(p, destination.UDPAddr())
}
func (c *bidirectionalNATPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
destination, err = c.NetPacketConn.ReadPacket(buffer)
if destination == c.origin {
destination = c.destination
if err != nil {
return
}
if socksaddrWithoutPort(destination) == c.origin {
destination = M.Socksaddr{
Addr: c.destination.Addr,
Fqdn: c.destination.Fqdn,
Port: destination.Port,
}
}
return
}
func (c *bidirectionalNATPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
if destination == c.destination {
destination = c.origin
if socksaddrWithoutPort(destination) == c.destination {
destination = M.Socksaddr{
Addr: c.origin.Addr,
Fqdn: c.origin.Fqdn,
Port: destination.Port,
}
}
return c.NetPacketConn.WritePacket(buffer, destination)
}
@ -101,3 +135,8 @@ func (c *bidirectionalNATPacketConn) UpdateDestination(destinationAddress netip.
func (c *bidirectionalNATPacketConn) Upstream() any {
return c.NetPacketConn
}
func socksaddrWithoutPort(destination M.Socksaddr) M.Socksaddr {
destination.Port = 0
return destination
}

39
common/bufio/nat_wait.go Normal file
View file

@ -0,0 +1,39 @@
package bufio
import (
"github.com/sagernet/sing/common/buf"
M "github.com/sagernet/sing/common/metadata"
N "github.com/sagernet/sing/common/network"
)
func (c *bidirectionalNATPacketConn) CreatePacketReadWaiter() (N.PacketReadWaiter, bool) {
waiter, created := CreatePacketReadWaiter(c.NetPacketConn)
if !created {
return nil, false
}
return &waitBidirectionalNATPacketConn{c, waiter}, true
}
type waitBidirectionalNATPacketConn struct {
*bidirectionalNATPacketConn
readWaiter N.PacketReadWaiter
}
func (c *waitBidirectionalNATPacketConn) InitializeReadWaiter(options N.ReadWaitOptions) (needCopy bool) {
return c.readWaiter.InitializeReadWaiter(options)
}
func (c *waitBidirectionalNATPacketConn) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
buffer, destination, err = c.readWaiter.WaitReadPacket()
if err != nil {
return
}
if socksaddrWithoutPort(destination) == c.origin {
destination = M.Socksaddr{
Addr: c.destination.Addr,
Fqdn: c.destination.Fqdn,
Port: destination.Port,
}
}
return
}