github.com/apernet/quic-go@v0.43.1-0.20240515053213-5e9e635fd9f0/frame_sorter.go (about) 1 package quic 2 3 import ( 4 "errors" 5 6 "github.com/apernet/quic-go/internal/protocol" 7 "github.com/apernet/quic-go/internal/utils" 8 "github.com/apernet/quic-go/internal/utils/tree" 9 ) 10 11 // byteInterval is an interval from one ByteCount to the other 12 type byteInterval struct { 13 Start protocol.ByteCount 14 End protocol.ByteCount 15 } 16 17 type frameSorterEntry struct { 18 Data []byte 19 DoneCb func() 20 } 21 22 type frameSorter struct { 23 queue map[protocol.ByteCount]frameSorterEntry 24 readPos protocol.ByteCount 25 gapTree *tree.Btree[utils.ByteInterval] 26 } 27 28 var errDuplicateStreamData = errors.New("duplicate stream data") 29 30 func newFrameSorter() *frameSorter { 31 s := frameSorter{ 32 gapTree: tree.New[utils.ByteInterval](), 33 queue: make(map[protocol.ByteCount]frameSorterEntry), 34 } 35 s.gapTree.Insert(utils.ByteInterval{Start: 0, End: protocol.MaxByteCount}) 36 return &s 37 } 38 39 func (s *frameSorter) Push(data []byte, offset protocol.ByteCount, doneCb func()) error { 40 err := s.push(data, offset, doneCb) 41 if err == errDuplicateStreamData { 42 if doneCb != nil { 43 doneCb() 44 } 45 return nil 46 } 47 return err 48 } 49 50 func (s *frameSorter) push(data []byte, offset protocol.ByteCount, doneCb func()) error { 51 if len(data) == 0 { 52 return errDuplicateStreamData 53 } 54 55 start := offset 56 end := offset + protocol.ByteCount(len(data)) 57 covInterval := utils.ByteInterval{Start: start, End: end} 58 59 gaps := s.gapTree.Match(covInterval) 60 61 if len(gaps) == 0 { 62 // No overlap with any existing gap 63 return errDuplicateStreamData 64 } 65 66 startGap := gaps[0] 67 endGap := gaps[len(gaps)-1] 68 startGapEqualsEndGap := len(gaps) == 1 69 70 if startGapEqualsEndGap && end <= startGap.Start { 71 return errDuplicateStreamData 72 } 73 74 startsInGap := covInterval.Start >= startGap.Start && covInterval.Start <= startGap.End 75 endsInGap := covInterval.End >= endGap.Start && covInterval.End < endGap.End 76 77 startGapEnd := startGap.End // save it, in case startGap is modified 78 endGapStart := endGap.Start // save it, in case endGap is modified 79 endGapEnd := endGap.End // save it, in case endGap is modified 80 81 var adjustedStartGapEnd bool 82 var wasCut bool 83 84 pos := start 85 var hasReplacedAtLeastOne bool 86 for { 87 oldEntry, ok := s.queue[pos] 88 if !ok { 89 break 90 } 91 oldEntryLen := protocol.ByteCount(len(oldEntry.Data)) 92 if end-pos > oldEntryLen || (hasReplacedAtLeastOne && end-pos == oldEntryLen) { 93 // The existing frame is shorter than the new frame. Replace it. 94 delete(s.queue, pos) 95 pos += oldEntryLen 96 hasReplacedAtLeastOne = true 97 if oldEntry.DoneCb != nil { 98 oldEntry.DoneCb() 99 } 100 } else { 101 if !hasReplacedAtLeastOne { 102 return errDuplicateStreamData 103 } 104 // The existing frame is longer than the new frame. 105 // Cut the new frame such that the end aligns with the start of the existing frame. 106 data = data[:pos-start] 107 end = pos 108 wasCut = true 109 break 110 } 111 } 112 113 if !startsInGap && !hasReplacedAtLeastOne { 114 // cut the frame, such that it starts at the start of the gap 115 data = data[startGap.Start-start:] 116 start = startGap.Start 117 wasCut = true 118 } 119 if start <= startGap.Start { 120 if end >= startGap.End { 121 // The frame covers the whole startGap. Delete the gap. 122 s.gapTree.Delete(startGap) 123 } else { 124 s.gapTree.Delete(startGap) 125 startGap.Start = end 126 // Re-insert the gap, but with the new start. 127 s.gapTree.Insert(startGap) 128 } 129 } else if !hasReplacedAtLeastOne { 130 s.gapTree.Delete(startGap) 131 startGap.End = start 132 // Re-insert the gap, but with the new end. 133 s.gapTree.Insert(startGap) 134 adjustedStartGapEnd = true 135 } 136 137 if !startGapEqualsEndGap { 138 s.deleteConsecutive(startGapEnd) 139 for _, g := range gaps[1:] { 140 if g.End >= endGapStart { 141 break 142 } 143 s.deleteConsecutive(g.End) 144 s.gapTree.Delete(g) 145 } 146 } 147 148 if !endsInGap && start != endGapEnd && end > endGapEnd { 149 // cut the frame, such that it ends at the end of the gap 150 data = data[:endGapEnd-start] 151 end = endGapEnd 152 wasCut = true 153 } 154 if end == endGapEnd { 155 if !startGapEqualsEndGap { 156 // The frame covers the whole endGap. Delete the gap. 157 s.gapTree.Delete(endGap) 158 } 159 } else { 160 if startGapEqualsEndGap && adjustedStartGapEnd { 161 // The frame split the existing gap into two. 162 s.gapTree.Insert(utils.ByteInterval{Start: end, End: startGapEnd}) 163 } else if !startGapEqualsEndGap { 164 s.gapTree.Delete(endGap) 165 endGap.Start = end 166 // Re-insert the gap, but with the new start. 167 s.gapTree.Insert(endGap) 168 } 169 } 170 171 if wasCut && len(data) < protocol.MinStreamFrameBufferSize { 172 newData := make([]byte, len(data)) 173 copy(newData, data) 174 data = newData 175 if doneCb != nil { 176 doneCb() 177 doneCb = nil 178 } 179 } 180 181 if s.gapTree.Len() > protocol.MaxStreamFrameSorterGaps { 182 return errors.New("too many gaps in received data") 183 } 184 185 s.queue[start] = frameSorterEntry{Data: data, DoneCb: doneCb} 186 return nil 187 } 188 189 // deleteConsecutive deletes consecutive frames from the queue, starting at pos 190 func (s *frameSorter) deleteConsecutive(pos protocol.ByteCount) { 191 for { 192 oldEntry, ok := s.queue[pos] 193 if !ok { 194 break 195 } 196 oldEntryLen := protocol.ByteCount(len(oldEntry.Data)) 197 delete(s.queue, pos) 198 if oldEntry.DoneCb != nil { 199 oldEntry.DoneCb() 200 } 201 pos += oldEntryLen 202 } 203 } 204 205 func (s *frameSorter) Pop() (protocol.ByteCount, []byte, func()) { 206 entry, ok := s.queue[s.readPos] 207 if !ok { 208 return s.readPos, nil, nil 209 } 210 delete(s.queue, s.readPos) 211 offset := s.readPos 212 s.readPos += protocol.ByteCount(len(entry.Data)) 213 return offset, entry.Data, entry.DoneCb 214 } 215 216 // HasMoreData says if there is any more data queued at *any* offset. 217 func (s *frameSorter) HasMoreData() bool { 218 return len(s.queue) > 0 219 }