github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/yuubinsya/packet.go (about)

     1  package yuubinsya
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"log/slog"
     7  	"net"
     8  
     9  	"github.com/Asutorufa/yuhaiin/pkg/log"
    10  	"github.com/Asutorufa/yuhaiin/pkg/net/nat"
    11  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    12  	"github.com/Asutorufa/yuhaiin/pkg/net/proxy/yuubinsya/types"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/pool"
    14  )
    15  
    16  type authPacketConn struct {
    17  	net.PacketConn
    18  	tcp    net.Conn
    19  	server net.Addr
    20  
    21  	auth types.Auth
    22  
    23  	prefix bool
    24  }
    25  
    26  func NewAuthPacketConn(local net.PacketConn) *authPacketConn {
    27  	return &authPacketConn{PacketConn: local}
    28  }
    29  
    30  func (s *authPacketConn) Close() error {
    31  	if s.tcp != nil {
    32  		s.tcp.Close()
    33  	}
    34  	return s.PacketConn.Close()
    35  }
    36  
    37  func (s *authPacketConn) WithTcpConn(tcp net.Conn) *authPacketConn {
    38  	s.tcp = tcp
    39  	return s
    40  }
    41  
    42  func (s *authPacketConn) WithTarget(target net.Addr) *authPacketConn {
    43  	s.server = target
    44  	return s
    45  }
    46  
    47  func (s *authPacketConn) WithAuth(auth types.Auth) *authPacketConn {
    48  	s.auth = auth
    49  
    50  	return s
    51  }
    52  
    53  func (s *authPacketConn) WithPrefix(b bool) *authPacketConn {
    54  	s.prefix = b
    55  	return s
    56  }
    57  
    58  func (s *authPacketConn) WriteTo(p []byte, addr net.Addr) (_ int, err error) {
    59  	return s.writeTo(p, addr, s.server)
    60  }
    61  
    62  func (s *authPacketConn) writeTo(p []byte, addr net.Addr, underlyingAddr net.Addr) (_ int, err error) {
    63  	buf := pool.GetBytesWriter(nat.MaxSegmentSize)
    64  	defer buf.Free()
    65  
    66  	err = types.EncodePacket(buf, addr, p, s.auth, s.prefix)
    67  	if err != nil {
    68  		return 0, fmt.Errorf("encode packet failed: %w", err)
    69  	}
    70  
    71  	_, err = s.PacketConn.WriteTo(buf.Bytes(), underlyingAddr)
    72  	if err != nil {
    73  		return 0, fmt.Errorf("write to remote failed: %w", err)
    74  	}
    75  
    76  	return len(p), nil
    77  }
    78  
    79  func (s *authPacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
    80  	n, addr, _, err := s.readFrom(p)
    81  	return n, addr, err
    82  }
    83  
    84  var errNet = errors.New("network error")
    85  
    86  func (s *authPacketConn) readFrom(p []byte) (int, netapi.Address, net.Addr, error) {
    87  	n, rawAddr, err := s.PacketConn.ReadFrom(p)
    88  	if err != nil {
    89  		return 0, nil, nil, fmt.Errorf("%w read from remote failed: %w", errNet, err)
    90  	}
    91  	buf, addr, err := types.DecodePacket(p[:n], s.auth, s.prefix)
    92  	if err != nil {
    93  		return 0, nil, nil, fmt.Errorf("decode packet failed: %w", err)
    94  	}
    95  
    96  	return copy(p[0:], buf), addr, rawAddr, nil
    97  }
    98  
    99  func StartUDPServer(packet net.PacketConn, sendPacket func(*netapi.Packet) error, auth types.Auth, prefix bool) {
   100  	p := NewAuthPacketConn(packet).WithAuth(auth).WithPrefix(prefix)
   101  	for {
   102  		buf := pool.GetBytesBuffer(nat.MaxSegmentSize)
   103  
   104  		n, dst, src, err := p.readFrom(buf.Bytes())
   105  		if err != nil {
   106  			log.Error("read udp request failed", slog.Any("err", err))
   107  
   108  			if errors.Is(err, errNet) {
   109  				return
   110  			}
   111  
   112  			continue
   113  		}
   114  
   115  		buf.Refactor(0, n)
   116  
   117  		err = sendPacket(&netapi.Packet{
   118  			Src:     src,
   119  			Dst:     dst,
   120  			Payload: buf,
   121  			WriteBack: func(b []byte, source net.Addr) (int, error) {
   122  				return p.writeTo(b, source, src)
   123  			},
   124  		})
   125  
   126  		if err != nil {
   127  			log.Error("send udp response failed", slog.Any("err", err))
   128  			break
   129  		}
   130  	}
   131  }