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