github.com/metacubex/mihomo@v1.18.5/transport/tuic/v5/frag.go (about)

     1  package v5
     2  
     3  import (
     4  	"bytes"
     5  	"sync"
     6  
     7  	"github.com/metacubex/mihomo/common/lru"
     8  
     9  	"github.com/metacubex/quic-go"
    10  )
    11  
    12  // MaxFragSize is a safe udp relay packet size
    13  // because tuicv5 support udp fragment so we unneeded to do a magic modify for quic-go to increase MaxDatagramFrameSize
    14  // it may not work fine in some platform
    15  // "1200" from quic-go's MaxDatagramSize
    16  // "-3" from quic-go's DatagramFrame.MaxDataLen
    17  var MaxFragSize = 1200 - PacketOverHead - 3
    18  
    19  func fragWriteNative(quicConn quic.Connection, packet Packet, buf *bytes.Buffer, fragSize int) (err error) {
    20  	fullPayload := packet.DATA
    21  	off := 0
    22  	fragID := uint8(0)
    23  	fragCount := uint8((len(fullPayload) + fragSize - 1) / fragSize) // round up
    24  	packet.FRAG_TOTAL = fragCount
    25  	for off < len(fullPayload) {
    26  		payloadSize := len(fullPayload) - off
    27  		if payloadSize > fragSize {
    28  			payloadSize = fragSize
    29  		}
    30  		frag := packet
    31  		frag.FRAG_ID = fragID
    32  		frag.SIZE = uint16(payloadSize)
    33  		frag.DATA = fullPayload[off : off+payloadSize]
    34  		off += payloadSize
    35  		fragID++
    36  		buf.Reset()
    37  		err = frag.WriteTo(buf)
    38  		if err != nil {
    39  			return
    40  		}
    41  		data := buf.Bytes()
    42  		err = quicConn.SendDatagram(data)
    43  		if err != nil {
    44  			return
    45  		}
    46  		packet.ADDR.TYPE = AtypNone // avoid "fragment 2/2: address in non-first fragment"
    47  	}
    48  	return
    49  }
    50  
    51  type deFragger struct {
    52  	lru  *lru.LruCache[uint16, *packetBag]
    53  	once sync.Once
    54  }
    55  
    56  type packetBag struct {
    57  	frags []*Packet
    58  	count uint8
    59  	mutex sync.Mutex
    60  }
    61  
    62  func newPacketBag() *packetBag {
    63  	return new(packetBag)
    64  }
    65  
    66  func (d *deFragger) init() {
    67  	if d.lru == nil {
    68  		d.lru = lru.New(
    69  			lru.WithAge[uint16, *packetBag](10),
    70  			lru.WithUpdateAgeOnGet[uint16, *packetBag](),
    71  		)
    72  	}
    73  }
    74  
    75  func (d *deFragger) Feed(m *Packet) *Packet {
    76  	if m.FRAG_TOTAL <= 1 {
    77  		return m
    78  	}
    79  	if m.FRAG_ID >= m.FRAG_TOTAL {
    80  		// wtf is this?
    81  		return nil
    82  	}
    83  	d.once.Do(d.init) // lazy init
    84  	bag, _ := d.lru.GetOrStore(m.PKT_ID, newPacketBag)
    85  	bag.mutex.Lock()
    86  	defer bag.mutex.Unlock()
    87  	if int(m.FRAG_TOTAL) != len(bag.frags) {
    88  		// new message, clear previous state
    89  		bag.frags = make([]*Packet, m.FRAG_TOTAL)
    90  		bag.count = 1
    91  		bag.frags[m.FRAG_ID] = m
    92  		return nil
    93  	}
    94  	if bag.frags[m.FRAG_ID] != nil {
    95  		return nil
    96  	}
    97  	bag.frags[m.FRAG_ID] = m
    98  	bag.count++
    99  	if int(bag.count) != len(bag.frags) {
   100  		return nil
   101  	}
   102  
   103  	// all fragments received, assemble
   104  	var data []byte
   105  	for _, frag := range bag.frags {
   106  		data = append(data, frag.DATA...)
   107  	}
   108  	p := *bag.frags[0] // recover from first fragment
   109  	p.SIZE = uint16(len(data))
   110  	p.DATA = data
   111  	p.FRAG_ID = 0
   112  	p.FRAG_TOTAL = 1
   113  	bag.frags = nil
   114  	d.lru.Delete(m.PKT_ID)
   115  	return &p
   116  }