github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/transport/internet/udp/dispatcher_packetaddr.go (about)

     1  package udp
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/v2fly/v2ray-core/v5/common/buf"
     7  	"github.com/v2fly/v2ray-core/v5/common/net"
     8  	"github.com/v2fly/v2ray-core/v5/common/net/packetaddr"
     9  	"github.com/v2fly/v2ray-core/v5/common/protocol/udp"
    10  	"github.com/v2fly/v2ray-core/v5/features/routing"
    11  )
    12  
    13  type PacketAddrDispatcher struct {
    14  	conn     net.PacketConn
    15  	callback ResponseCallback
    16  	ctx      context.Context
    17  }
    18  
    19  func (p PacketAddrDispatcher) Close() error {
    20  	if p.ctx.Value(DispatcherConnectionTerminationSignalReceiverMark) != nil {
    21  		p.ctx.Value(DispatcherConnectionTerminationSignalReceiverMark).(DispatcherConnectionTerminationSignalReceiver).Close()
    22  	}
    23  	return p.conn.Close()
    24  }
    25  
    26  func (p PacketAddrDispatcher) Dispatch(ctx context.Context, destination net.Destination, payload *buf.Buffer) {
    27  	if destination.Network != net.Network_UDP {
    28  		return
    29  	}
    30  	p.conn.WriteTo(payload.Bytes(), &net.UDPAddr{IP: destination.Address.IP(), Port: int(destination.Port.Value())})
    31  }
    32  
    33  func (p PacketAddrDispatcher) readWorker() {
    34  	for {
    35  		readBuf := buf.New()
    36  		n, addr, err := p.conn.ReadFrom(readBuf.Extend(2048))
    37  		if err != nil {
    38  			return
    39  		}
    40  		readBuf.Resize(0, int32(n))
    41  		p.callback(p.ctx, &udp.Packet{Payload: readBuf, Source: net.DestinationFromAddr(addr)})
    42  	}
    43  }
    44  
    45  type PacketAddrDispatcherCreator struct {
    46  	ctx context.Context
    47  }
    48  
    49  func NewPacketAddrDispatcherCreator(ctx context.Context) PacketAddrDispatcherCreator {
    50  	return PacketAddrDispatcherCreator{ctx: ctx}
    51  }
    52  
    53  func (pdc *PacketAddrDispatcherCreator) NewPacketAddrDispatcher(
    54  	dispatcher routing.Dispatcher, callback ResponseCallback,
    55  ) DispatcherI {
    56  	packetConn, _ := packetaddr.CreatePacketAddrConn(pdc.ctx, dispatcher, false)
    57  	pd := &PacketAddrDispatcher{conn: packetConn, callback: callback, ctx: pdc.ctx}
    58  	go pd.readWorker()
    59  	return pd
    60  }