github.com/eagleql/xray-core@v1.4.4/common/xudp/xudp.go (about)

     1  package xudp
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/eagleql/xray-core/common/buf"
     7  	"github.com/eagleql/xray-core/common/net"
     8  	"github.com/eagleql/xray-core/common/protocol"
     9  )
    10  
    11  var addrParser = protocol.NewAddressParser(
    12  	protocol.AddressFamilyByte(byte(protocol.AddressTypeIPv4), net.AddressFamilyIPv4),
    13  	protocol.AddressFamilyByte(byte(protocol.AddressTypeDomain), net.AddressFamilyDomain),
    14  	protocol.AddressFamilyByte(byte(protocol.AddressTypeIPv6), net.AddressFamilyIPv6),
    15  	protocol.PortThenAddress(),
    16  )
    17  
    18  func NewPacketWriter(writer buf.Writer, dest net.Destination) *PacketWriter {
    19  	return &PacketWriter{
    20  		Writer: writer,
    21  		Dest:   dest,
    22  	}
    23  }
    24  
    25  type PacketWriter struct {
    26  	Writer buf.Writer
    27  	Dest   net.Destination
    28  }
    29  
    30  func (w *PacketWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
    31  	defer buf.ReleaseMulti(mb)
    32  	mb2Write := make(buf.MultiBuffer, 0, len(mb))
    33  	for _, b := range mb {
    34  		length := b.Len()
    35  		if length == 0 || length+666 > buf.Size {
    36  			continue
    37  		}
    38  
    39  		eb := buf.New()
    40  		eb.Write([]byte{0, 0, 0, 0})
    41  		if w.Dest.Network == net.Network_UDP {
    42  			eb.WriteByte(1) // New
    43  			eb.WriteByte(1) // Opt
    44  			eb.WriteByte(2) // UDP
    45  			addrParser.WriteAddressPort(eb, w.Dest.Address, w.Dest.Port)
    46  			w.Dest.Network = net.Network_Unknown
    47  		} else {
    48  			eb.WriteByte(2) // Keep
    49  			eb.WriteByte(1)
    50  			if b.UDP != nil {
    51  				eb.WriteByte(2)
    52  				addrParser.WriteAddressPort(eb, b.UDP.Address, b.UDP.Port)
    53  			}
    54  		}
    55  		l := eb.Len() - 2
    56  		eb.SetByte(0, byte(l>>8))
    57  		eb.SetByte(1, byte(l))
    58  		eb.WriteByte(byte(length >> 8))
    59  		eb.WriteByte(byte(length))
    60  		eb.Write(b.Bytes())
    61  
    62  		mb2Write = append(mb2Write, eb)
    63  	}
    64  	if mb2Write.IsEmpty() {
    65  		return nil
    66  	}
    67  	return w.Writer.WriteMultiBuffer(mb2Write)
    68  }
    69  
    70  func NewPacketReader(reader io.Reader) *PacketReader {
    71  	return &PacketReader{
    72  		Reader: reader,
    73  		cache:  make([]byte, 2),
    74  	}
    75  }
    76  
    77  type PacketReader struct {
    78  	Reader io.Reader
    79  	cache  []byte
    80  }
    81  
    82  func (r *PacketReader) ReadMultiBuffer() (buf.MultiBuffer, error) {
    83  	for {
    84  		if _, err := io.ReadFull(r.Reader, r.cache); err != nil {
    85  			return nil, err
    86  		}
    87  		l := int32(r.cache[0])<<8 | int32(r.cache[1])
    88  		if l < 4 {
    89  			return nil, io.EOF
    90  		}
    91  		b := buf.New()
    92  		if _, err := b.ReadFullFrom(r.Reader, l); err != nil {
    93  			b.Release()
    94  			return nil, err
    95  		}
    96  		discard := false
    97  		switch b.Byte(2) {
    98  		case 2:
    99  			if l != 4 {
   100  				b.Advance(5)
   101  				addr, port, err := addrParser.ReadAddressPort(nil, b)
   102  				if err != nil {
   103  					b.Release()
   104  					return nil, err
   105  				}
   106  				b.UDP = &net.Destination{
   107  					Network: net.Network_UDP,
   108  					Address: addr,
   109  					Port:    port,
   110  				}
   111  			}
   112  		case 4:
   113  			discard = true
   114  		default:
   115  			b.Release()
   116  			return nil, io.EOF
   117  		}
   118  		if b.Byte(3) == 1 {
   119  			if _, err := io.ReadFull(r.Reader, r.cache); err != nil {
   120  				b.Release()
   121  				return nil, err
   122  			}
   123  			length := int32(r.cache[0])<<8 | int32(r.cache[1])
   124  			if length > 0 {
   125  				b.Clear()
   126  				if _, err := b.ReadFullFrom(r.Reader, length); err != nil {
   127  					b.Release()
   128  					return nil, err
   129  				}
   130  				if !discard {
   131  					return buf.MultiBuffer{b}, nil
   132  				}
   133  			}
   134  		}
   135  		b.Release()
   136  	}
   137  }