package yamux import ( "io" "math" "sync" "sync/atomic" "time" ) type streamState int const ( streamInit streamState = iota streamSYNSent streamSYNReceived streamEstablished streamFinished ) type halfStreamState int const ( halfOpen halfStreamState = iota halfClosed halfReset ) // Stream is used to represent a logical stream // within a session. type Stream struct { sendWindow uint32 memorySpan MemoryManager id uint32 session *Session recvWindow uint32 epochStart time.Time state streamState writeState, readState halfStreamState stateLock sync.Mutex recvBuf segmentedBuffer recvNotifyCh chan struct{} sendNotifyCh chan struct{} readDeadline, writeDeadline pipeDeadline } // newStream is used to construct a new stream within a given session for an ID. // It assumes that a memory allocation has been obtained for the initialWindow. func newStream(session *Session, id uint32, state streamState, initialWindow uint32, memorySpan MemoryManager) *Stream { s := &Stream{ id: id, session: session, state: state, sendWindow: initialStreamWindow, readDeadline: makePipeDeadline(), writeDeadline: makePipeDeadline(), memorySpan: memorySpan, // Initialize the recvBuf with initialStreamWindow, not config.InitialStreamWindowSize. // The peer isn't allowed to send more data than initialStreamWindow until we've sent // the first window update (which will grant it up to config.InitialStreamWindowSize). recvBuf: newSegmentedBuffer(initialWindow), recvWindow: session.config.InitialStreamWindowSize, epochStart: time.Now(), recvNotifyCh: make(chan struct{}, 1), sendNotifyCh: make(chan struct{}, 1), } return s } // Session returns the associated stream session func (s *Stream) Session() *Session { return s.session } // StreamID returns the ID of this stream func (s *Stream) StreamID() uint32 { return s.id } // Read is used to read from the stream func (s *Stream) Read(b []byte) (n int, err error) { START: s.stateLock.Lock() state := s.readState s.stateLock.Unlock() switch state { case halfOpen: // Open -> read case halfClosed: empty := s.recvBuf.Len() == 0 if empty { return 0, io.EOF } // Closed, but we have data pending -> read. case halfReset: return 0, ErrStreamReset default: panic("unknown state") } // If there is no data available, block if s.recvBuf.Len() == 0 { select { case <-s.recvNotifyCh: goto START case <-s.readDeadline.wait(): return 0, ErrTimeout } } // Read any bytes n, _ = s.recvBuf.Read(b) // Send a window update potentially err = s.sendWindowUpdate(s.readDeadline.wait()) return n, err } // Write is used to write to the stream func (s *Stream) Write(b []byte) (int, error) { var total int for total < len(b) { n, err := s.write(b[total:]) total += n if err != nil { return total, err } } return total, nil } // write is used to write to the stream, may return on // a short write. func (s *Stream) write(b []byte) (n int, err error) { var flags uint16 var max uint32 var hdr header START: s.stateLock.Lock() state := s.writeState s.stateLock.Unlock() switch state { case halfOpen: // Open for writing -> write case halfClosed: return 0, ErrStreamClosed case halfReset: return 0, ErrStreamReset default: panic("unknown state") } // If there is no data available, block window := atomic.LoadUint32(&s.sendWindow) if window == 0 { select { case <-s.sendNotifyCh: goto START case <-s.writeDeadline.wait(): return 0, ErrTimeout } } // Determine the flags if any flags = s.sendFlags() // Send up to min(message, window max = min(window, s.session.config.MaxMessageSize-headerSize, uint32(len(b))) // Send the header hdr = encode(typeData, flags, s.id, max) if err = s.session.sendMsg(hdr, b[:max], s.writeDeadline.wait()); err != nil { return 0, err } // Reduce our send window atomic.AddUint32(&s.sendWindow, ^uint32(max-1)) // Unlock return int(max), err } // sendFlags determines any flags that are appropriate // based on the current stream state func (s *Stream) sendFlags() uint16 { s.stateLock.Lock() defer s.stateLock.Unlock() var flags uint16 switch s.state { case streamInit: flags |= flagSYN s.state = streamSYNSent case streamSYNReceived: flags |= flagACK s.state = streamEstablished } return flags } // sendWindowUpdate potentially sends a window update enabling // further writes to take place. Must be invoked with the lock. func (s *Stream) sendWindowUpdate(deadline <-chan struct{}) error { // Determine the flags if any flags := s.sendFlags() // Update the receive window. needed, delta := s.recvBuf.GrowTo(s.recvWindow, flags != 0) if !needed { return nil } now := time.Now() if rtt := s.session.getRTT(); flags == 0 && rtt > 0 && now.Sub(s.epochStart) < rtt*4 { var recvWindow uint32 if s.recvWindow > math.MaxUint32/2 { recvWindow = min(math.MaxUint32, s.session.config.MaxStreamWindowSize) } else { recvWindow = min(s.recvWindow*2, s.session.config.MaxStreamWindowSize) } if recvWindow > s.recvWindow { grow := recvWindow - s.recvWindow if err := s.memorySpan.ReserveMemory(int(grow), 128); err == nil { s.recvWindow = recvWindow _, delta = s.recvBuf.GrowTo(s.recvWindow, true) } } } s.epochStart = now hdr := encode(typeWindowUpdate, flags, s.id, delta) return s.session.sendMsg(hdr, nil, deadline) } // sendClose is used to send a FIN func (s *Stream) sendClose() error { flags := s.sendFlags() flags |= flagFIN hdr := encode(typeWindowUpdate, flags, s.id, 0) return s.session.sendMsg(hdr, nil, nil) } // sendReset is used to send a RST func (s *Stream) sendReset() error { hdr := encode(typeWindowUpdate, flagRST, s.id, 0) return s.session.sendMsg(hdr, nil, nil) } // Reset resets the stream (forcibly closes the stream) func (s *Stream) Reset() error { sendReset := false s.stateLock.Lock() switch s.state { case streamFinished: s.stateLock.Unlock() return nil case streamInit: // we haven't sent anything, so we don't need to send a reset. case streamSYNSent, streamSYNReceived, streamEstablished: sendReset = true default: panic("unhandled state") } // at least one direction is open, we need to reset. // If we've already sent/received an EOF, no need to reset that side. if s.writeState == halfOpen { s.writeState = halfReset } if s.readState == halfOpen { s.readState = halfReset } s.state = streamFinished s.notifyWaiting() s.stateLock.Unlock() if sendReset { _ = s.sendReset() } s.cleanup() return nil } // CloseWrite is used to close the stream for writing. func (s *Stream) CloseWrite() error { s.stateLock.Lock() switch s.writeState { case halfOpen: // Open for writing -> close write case halfClosed: s.stateLock.Unlock() return nil case halfReset: s.stateLock.Unlock() return ErrStreamReset default: panic("invalid state") } s.writeState = halfClosed cleanup := s.readState != halfOpen if cleanup { s.state = streamFinished } s.stateLock.Unlock() s.notifyWaiting() err := s.sendClose() if cleanup { // we're fully closed, might as well be nice to the user and // free everything early. s.cleanup() } return err } // CloseRead is used to close the stream for writing. func (s *Stream) CloseRead() error { cleanup := false s.stateLock.Lock() switch s.readState { case halfOpen: // Open for reading -> close read case halfClosed, halfReset: s.stateLock.Unlock() return nil default: panic("invalid state") } s.readState = halfReset cleanup = s.writeState != halfOpen if cleanup { s.state = streamFinished } s.stateLock.Unlock() s.notifyWaiting() if cleanup { // we're fully closed, might as well be nice to the user and // free everything early. s.cleanup() } return nil } // Close is used to close the stream. func (s *Stream) Close() error { _ = s.CloseRead() // can't fail. return s.CloseWrite() } // forceClose is used for when the session is exiting func (s *Stream) forceClose() { s.stateLock.Lock() if s.readState == halfOpen { s.readState = halfReset } if s.writeState == halfOpen { s.writeState = halfReset } s.state = streamFinished s.notifyWaiting() s.stateLock.Unlock() s.readDeadline.set(time.Time{}) s.writeDeadline.set(time.Time{}) } // called when fully closed to release any system resources. func (s *Stream) cleanup() { s.session.closeStream(s.id) s.readDeadline.set(time.Time{}) s.writeDeadline.set(time.Time{}) } // processFlags is used to update the state of the stream // based on set flags, if any. Lock must be held func (s *Stream) processFlags(flags uint16) { // Close the stream without holding the state lock var closeStream bool defer func() { if closeStream { s.cleanup() } }() if flags&flagACK == flagACK { s.stateLock.Lock() if s.state == streamSYNSent { s.state = streamEstablished } s.stateLock.Unlock() s.session.establishStream(s.id) } if flags&flagFIN == flagFIN { var notify bool s.stateLock.Lock() if s.readState == halfOpen { s.readState = halfClosed if s.writeState != halfOpen { // We're now fully closed. closeStream = true s.state = streamFinished } notify = true } s.stateLock.Unlock() if notify { s.notifyWaiting() } } if flags&flagRST == flagRST { s.stateLock.Lock() if s.readState == halfOpen { s.readState = halfReset } if s.writeState == halfOpen { s.writeState = halfReset } s.state = streamFinished s.stateLock.Unlock() closeStream = true s.notifyWaiting() } } // notifyWaiting notifies all the waiting channels func (s *Stream) notifyWaiting() { asyncNotify(s.recvNotifyCh) asyncNotify(s.sendNotifyCh) } // incrSendWindow updates the size of our send window func (s *Stream) incrSendWindow(hdr header, flags uint16) { s.processFlags(flags) // Increase window, unblock a sender atomic.AddUint32(&s.sendWindow, hdr.Length()) asyncNotify(s.sendNotifyCh) } // readData is used to handle a data frame func (s *Stream) readData(hdr header, flags uint16, conn io.Reader) error { s.processFlags(flags) // Check that our recv window is not exceeded length := hdr.Length() if length == 0 { return nil } // Copy into buffer if err := s.recvBuf.Append(conn, length); err != nil { s.session.logger.Printf("[ERR] yamux: Failed to read stream data on stream %d: %v", s.id, err) return err } // Unblock the reader asyncNotify(s.recvNotifyCh) return nil } // SetDeadline sets the read and write deadlines func (s *Stream) SetDeadline(t time.Time) error { if err := s.SetReadDeadline(t); err != nil { return err } if err := s.SetWriteDeadline(t); err != nil { return err } return nil } // SetReadDeadline sets the deadline for future Read calls. func (s *Stream) SetReadDeadline(t time.Time) error { s.stateLock.Lock() defer s.stateLock.Unlock() if s.readState == halfOpen { s.readDeadline.set(t) } return nil } // SetWriteDeadline sets the deadline for future Write calls func (s *Stream) SetWriteDeadline(t time.Time) error { s.stateLock.Lock() defer s.stateLock.Unlock() if s.writeState == halfOpen { s.writeDeadline.set(t) } return nil }