github.com/inazumav/sing-box@v0.0.0-20230926072359-ab51429a14f1/transport/hysteria/frag.go (about)

     1  package hysteria
     2  
     3  func FragUDPMessage(m UDPMessage, maxSize int) []UDPMessage {
     4  	if m.Size() <= maxSize {
     5  		return []UDPMessage{m}
     6  	}
     7  	fullPayload := m.Data
     8  	maxPayloadSize := maxSize - m.HeaderSize()
     9  	off := 0
    10  	fragID := uint8(0)
    11  	fragCount := uint8((len(fullPayload) + maxPayloadSize - 1) / maxPayloadSize) // round up
    12  	var frags []UDPMessage
    13  	for off < len(fullPayload) {
    14  		payloadSize := len(fullPayload) - off
    15  		if payloadSize > maxPayloadSize {
    16  			payloadSize = maxPayloadSize
    17  		}
    18  		frag := m
    19  		frag.FragID = fragID
    20  		frag.FragCount = fragCount
    21  		frag.Data = fullPayload[off : off+payloadSize]
    22  		frags = append(frags, frag)
    23  		off += payloadSize
    24  		fragID++
    25  	}
    26  	return frags
    27  }
    28  
    29  type Defragger struct {
    30  	msgID uint16
    31  	frags []*UDPMessage
    32  	count uint8
    33  }
    34  
    35  func (d *Defragger) Feed(m UDPMessage) *UDPMessage {
    36  	if m.FragCount <= 1 {
    37  		return &m
    38  	}
    39  	if m.FragID >= m.FragCount {
    40  		// wtf is this?
    41  		return nil
    42  	}
    43  	if m.MsgID != d.msgID {
    44  		// new message, clear previous state
    45  		d.msgID = m.MsgID
    46  		d.frags = make([]*UDPMessage, m.FragCount)
    47  		d.count = 1
    48  		d.frags[m.FragID] = &m
    49  	} else if d.frags[m.FragID] == nil {
    50  		d.frags[m.FragID] = &m
    51  		d.count++
    52  		if int(d.count) == len(d.frags) {
    53  			// all fragments received, assemble
    54  			var data []byte
    55  			for _, frag := range d.frags {
    56  				data = append(data, frag.Data...)
    57  			}
    58  			m.Data = data
    59  			m.FragID = 0
    60  			m.FragCount = 1
    61  			return &m
    62  		}
    63  	}
    64  	return nil
    65  }