1package quic 2 3import ( 4 "context" 5 "sync" 6 7 "github.com/lucas-clemente/quic-go/internal/protocol" 8 "github.com/lucas-clemente/quic-go/internal/wire" 9) 10 11//go:generate genny -in $GOFILE -out streams_map_incoming_bidi.go gen "item=streamI Item=BidiStream streamTypeGeneric=protocol.StreamTypeBidi" 12//go:generate genny -in $GOFILE -out streams_map_incoming_uni.go gen "item=receiveStreamI Item=UniStream streamTypeGeneric=protocol.StreamTypeUni" 13type incomingItemsMap struct { 14 mutex sync.RWMutex 15 newStreamChan chan struct{} 16 17 streams map[protocol.StreamNum]item 18 // When a stream is deleted before it was accepted, we can't delete it immediately. 19 // We need to wait until the application accepts it, and delete it immediately then. 20 streamsToDelete map[protocol.StreamNum]struct{} // used as a set 21 22 nextStreamToAccept protocol.StreamNum // the next stream that will be returned by AcceptStream() 23 nextStreamToOpen protocol.StreamNum // the highest stream that the peer opened 24 maxStream protocol.StreamNum // the highest stream that the peer is allowed to open 25 maxNumStreams uint64 // maximum number of streams 26 27 newStream func(protocol.StreamNum) item 28 queueMaxStreamID func(*wire.MaxStreamsFrame) 29 30 closeErr error 31} 32 33func newIncomingItemsMap( 34 newStream func(protocol.StreamNum) item, 35 maxStreams uint64, 36 queueControlFrame func(wire.Frame), 37) *incomingItemsMap { 38 return &incomingItemsMap{ 39 newStreamChan: make(chan struct{}, 1), 40 streams: make(map[protocol.StreamNum]item), 41 streamsToDelete: make(map[protocol.StreamNum]struct{}), 42 maxStream: protocol.StreamNum(maxStreams), 43 maxNumStreams: maxStreams, 44 newStream: newStream, 45 nextStreamToOpen: 1, 46 nextStreamToAccept: 1, 47 queueMaxStreamID: func(f *wire.MaxStreamsFrame) { queueControlFrame(f) }, 48 } 49} 50 51func (m *incomingItemsMap) AcceptStream(ctx context.Context) (item, error) { 52 // drain the newStreamChan, so we don't check the map twice if the stream doesn't exist 53 select { 54 case <-m.newStreamChan: 55 default: 56 } 57 58 m.mutex.Lock() 59 60 var num protocol.StreamNum 61 var str item 62 for { 63 num = m.nextStreamToAccept 64 if m.closeErr != nil { 65 m.mutex.Unlock() 66 return nil, m.closeErr 67 } 68 var ok bool 69 str, ok = m.streams[num] 70 if ok { 71 break 72 } 73 m.mutex.Unlock() 74 select { 75 case <-ctx.Done(): 76 return nil, ctx.Err() 77 case <-m.newStreamChan: 78 } 79 m.mutex.Lock() 80 } 81 m.nextStreamToAccept++ 82 // If this stream was completed before being accepted, we can delete it now. 83 if _, ok := m.streamsToDelete[num]; ok { 84 delete(m.streamsToDelete, num) 85 if err := m.deleteStream(num); err != nil { 86 m.mutex.Unlock() 87 return nil, err 88 } 89 } 90 m.mutex.Unlock() 91 return str, nil 92} 93 94func (m *incomingItemsMap) GetOrOpenStream(num protocol.StreamNum) (item, error) { 95 m.mutex.RLock() 96 if num > m.maxStream { 97 m.mutex.RUnlock() 98 return nil, streamError{ 99 message: "peer tried to open stream %d (current limit: %d)", 100 nums: []protocol.StreamNum{num, m.maxStream}, 101 } 102 } 103 // if the num is smaller than the highest we accepted 104 // * this stream exists in the map, and we can return it, or 105 // * this stream was already closed, then we can return the nil 106 if num < m.nextStreamToOpen { 107 var s item 108 // If the stream was already queued for deletion, and is just waiting to be accepted, don't return it. 109 if _, ok := m.streamsToDelete[num]; !ok { 110 s = m.streams[num] 111 } 112 m.mutex.RUnlock() 113 return s, nil 114 } 115 m.mutex.RUnlock() 116 117 m.mutex.Lock() 118 // no need to check the two error conditions from above again 119 // * maxStream can only increase, so if the id was valid before, it definitely is valid now 120 // * highestStream is only modified by this function 121 for newNum := m.nextStreamToOpen; newNum <= num; newNum++ { 122 m.streams[newNum] = m.newStream(newNum) 123 select { 124 case m.newStreamChan <- struct{}{}: 125 default: 126 } 127 } 128 m.nextStreamToOpen = num + 1 129 s := m.streams[num] 130 m.mutex.Unlock() 131 return s, nil 132} 133 134func (m *incomingItemsMap) DeleteStream(num protocol.StreamNum) error { 135 m.mutex.Lock() 136 defer m.mutex.Unlock() 137 138 return m.deleteStream(num) 139} 140 141func (m *incomingItemsMap) deleteStream(num protocol.StreamNum) error { 142 if _, ok := m.streams[num]; !ok { 143 return streamError{ 144 message: "Tried to delete unknown incoming stream %d", 145 nums: []protocol.StreamNum{num}, 146 } 147 } 148 149 // Don't delete this stream yet, if it was not yet accepted. 150 // Just save it to streamsToDelete map, to make sure it is deleted as soon as it gets accepted. 151 if num >= m.nextStreamToAccept { 152 if _, ok := m.streamsToDelete[num]; ok { 153 return streamError{ 154 message: "Tried to delete incoming stream %d multiple times", 155 nums: []protocol.StreamNum{num}, 156 } 157 } 158 m.streamsToDelete[num] = struct{}{} 159 return nil 160 } 161 162 delete(m.streams, num) 163 // queue a MAX_STREAM_ID frame, giving the peer the option to open a new stream 164 if m.maxNumStreams > uint64(len(m.streams)) { 165 maxStream := m.nextStreamToOpen + protocol.StreamNum(m.maxNumStreams-uint64(len(m.streams))) - 1 166 // Never send a value larger than protocol.MaxStreamCount. 167 if maxStream <= protocol.MaxStreamCount { 168 m.maxStream = maxStream 169 m.queueMaxStreamID(&wire.MaxStreamsFrame{ 170 Type: streamTypeGeneric, 171 MaxStreamNum: m.maxStream, 172 }) 173 } 174 } 175 return nil 176} 177 178func (m *incomingItemsMap) CloseWithError(err error) { 179 m.mutex.Lock() 180 m.closeErr = err 181 for _, str := range m.streams { 182 str.closeForShutdown(err) 183 } 184 m.mutex.Unlock() 185 close(m.newStreamChan) 186} 187