1package quic 2 3import ( 4 "context" 5 "sync" 6 7 "github.com/ooni/psiphon/oopsi/github.com/Psiphon-Labs/quic-go/internal/protocol" 8 "github.com/ooni/psiphon/oopsi/github.com/Psiphon-Labs/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 openend 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 // streamNumToID func(protocol.StreamNum) protocol.StreamID // only used for generating errors 30 31 closeErr error 32} 33 34func newIncomingItemsMap( 35 newStream func(protocol.StreamNum) item, 36 maxStreams uint64, 37 queueControlFrame func(wire.Frame), 38) *incomingItemsMap { 39 return &incomingItemsMap{ 40 newStreamChan: make(chan struct{}), 41 streams: make(map[protocol.StreamNum]item), 42 streamsToDelete: make(map[protocol.StreamNum]struct{}), 43 maxStream: protocol.StreamNum(maxStreams), 44 maxNumStreams: maxStreams, 45 newStream: newStream, 46 nextStreamToOpen: 1, 47 nextStreamToAccept: 1, 48 queueMaxStreamID: func(f *wire.MaxStreamsFrame) { queueControlFrame(f) }, 49 } 50} 51 52func (m *incomingItemsMap) AcceptStream(ctx context.Context) (item, error) { 53 m.mutex.Lock() 54 55 var num protocol.StreamNum 56 var str item 57 for { 58 num = m.nextStreamToAccept 59 if m.closeErr != nil { 60 m.mutex.Unlock() 61 return nil, m.closeErr 62 } 63 var ok bool 64 str, ok = m.streams[num] 65 if ok { 66 break 67 } 68 m.mutex.Unlock() 69 select { 70 case <-ctx.Done(): 71 return nil, ctx.Err() 72 case <-m.newStreamChan: 73 } 74 m.mutex.Lock() 75 } 76 m.nextStreamToAccept++ 77 // If this stream was completed before being accepted, we can delete it now. 78 if _, ok := m.streamsToDelete[num]; ok { 79 delete(m.streamsToDelete, num) 80 if err := m.deleteStream(num); err != nil { 81 m.mutex.Unlock() 82 return nil, err 83 } 84 } 85 m.mutex.Unlock() 86 return str, nil 87} 88 89func (m *incomingItemsMap) GetOrOpenStream(num protocol.StreamNum) (item, error) { 90 m.mutex.RLock() 91 if num > m.maxStream { 92 m.mutex.RUnlock() 93 return nil, streamError{ 94 message: "peer tried to open stream %d (current limit: %d)", 95 nums: []protocol.StreamNum{num, m.maxStream}, 96 } 97 } 98 // if the num is smaller than the highest we accepted 99 // * this stream exists in the map, and we can return it, or 100 // * this stream was already closed, then we can return the nil 101 if num < m.nextStreamToOpen { 102 var s item 103 // If the stream was already queued for deletion, and is just waiting to be accepted, don't return it. 104 if _, ok := m.streamsToDelete[num]; !ok { 105 s = m.streams[num] 106 } 107 m.mutex.RUnlock() 108 return s, nil 109 } 110 m.mutex.RUnlock() 111 112 m.mutex.Lock() 113 // no need to check the two error conditions from above again 114 // * maxStream can only increase, so if the id was valid before, it definitely is valid now 115 // * highestStream is only modified by this function 116 for newNum := m.nextStreamToOpen; newNum <= num; newNum++ { 117 m.streams[newNum] = m.newStream(newNum) 118 select { 119 case m.newStreamChan <- struct{}{}: 120 default: 121 } 122 } 123 m.nextStreamToOpen = num + 1 124 s := m.streams[num] 125 m.mutex.Unlock() 126 return s, nil 127} 128 129func (m *incomingItemsMap) DeleteStream(num protocol.StreamNum) error { 130 m.mutex.Lock() 131 defer m.mutex.Unlock() 132 133 return m.deleteStream(num) 134} 135 136func (m *incomingItemsMap) deleteStream(num protocol.StreamNum) error { 137 if _, ok := m.streams[num]; !ok { 138 return streamError{ 139 message: "Tried to delete unknown incoming stream %d", 140 nums: []protocol.StreamNum{num}, 141 } 142 } 143 144 // Don't delete this stream yet, if it was not yet accepted. 145 // Just save it to streamsToDelete map, to make sure it is deleted as soon as it gets accepted. 146 if num >= m.nextStreamToAccept { 147 if _, ok := m.streamsToDelete[num]; ok { 148 return streamError{ 149 message: "Tried to delete incoming stream %d multiple times", 150 nums: []protocol.StreamNum{num}, 151 } 152 } 153 m.streamsToDelete[num] = struct{}{} 154 return nil 155 } 156 157 delete(m.streams, num) 158 // queue a MAX_STREAM_ID frame, giving the peer the option to open a new stream 159 if m.maxNumStreams > uint64(len(m.streams)) { 160 numNewStreams := m.maxNumStreams - uint64(len(m.streams)) 161 m.maxStream = m.nextStreamToOpen + protocol.StreamNum(numNewStreams) - 1 162 m.queueMaxStreamID(&wire.MaxStreamsFrame{ 163 Type: streamTypeGeneric, 164 MaxStreamNum: m.maxStream, 165 }) 166 } 167 return nil 168} 169 170func (m *incomingItemsMap) CloseWithError(err error) { 171 m.mutex.Lock() 172 m.closeErr = err 173 for _, str := range m.streams { 174 str.closeForShutdown(err) 175 } 176 m.mutex.Unlock() 177 close(m.newStreamChan) 178} 179