17 // Session is used to wrap a reliable ordered connection and to
18 // multiplex it into multiple streams.
20 // remoteGoAway indicates the remote side does
21 // not want futher connections. Must be first for alignment.
24 // localGoAway indicates that we should stop
25 // accepting futher connections. Must be first for alignment.
28 // nextStreamID is the next stream we should
29 // send. This depends if we are a client/server.
32 // config holds our configuration
35 // logger is used for our logs
38 // conn is the underlying connection
39 conn io.ReadWriteCloser
41 // bufRead is a buffered reader
44 // pings is used to track inflight pings
45 pings map[uint32]chan struct{}
49 // streams maps a stream id to a stream, and inflight has an entry
50 // for any outgoing stream that has not yet been established. Both are
51 // protected by streamLock.
52 streams map[uint32]*Stream
53 inflight map[uint32]struct{}
56 // synCh acts like a semaphore. It is sized to the AcceptBacklog which
57 // is assumed to be symmetric between the client and server. This allows
58 // the client to avoid exceeding the backlog and instead blocks the open.
61 // acceptCh is used to pass ready streams to the client
64 // sendCh is used to mark a stream as ready to send,
65 // or to send a header out directly.
68 // recvDoneCh is closed when recv() exits to avoid a race
69 // between stream registration and stream shutdown
70 recvDoneCh chan struct{}
72 // shutdown is used to safely close a session
75 shutdownCh chan struct{}
76 shutdownLock sync.Mutex
79 // sendReady is used to either mark a stream as ready
80 // or to directly send a header
81 type sendReady struct {
87 // newSession is used to construct a new session
88 func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
91 logger: log.New(config.LogOutput, "", log.LstdFlags),
93 bufRead: bufio.NewReader(conn),
94 pings: make(map[uint32]chan struct{}),
95 streams: make(map[uint32]*Stream),
96 inflight: make(map[uint32]struct{}),
97 synCh: make(chan struct{}, config.AcceptBacklog),
98 acceptCh: make(chan *Stream, config.AcceptBacklog),
99 sendCh: make(chan sendReady, 64),
100 recvDoneCh: make(chan struct{}),
101 shutdownCh: make(chan struct{}),
110 if config.EnableKeepAlive {
116 // IsClosed does a safe check to see if we have shutdown
117 func (s *Session) IsClosed() bool {
126 // CloseChan returns a read-only channel which is closed as
127 // soon as the session is closed.
128 func (s *Session) CloseChan() <-chan struct{} {
132 // NumStreams returns the number of currently open streams
133 func (s *Session) NumStreams() int {
135 num := len(s.streams)
136 s.streamLock.Unlock()
140 // Open is used to create a new stream as a net.Conn
141 func (s *Session) Open() (net.Conn, error) {
142 conn, err := s.OpenStream()
149 // OpenStream is used to create a new stream
150 func (s *Session) OpenStream() (*Stream, error) {
152 return nil, ErrSessionShutdown
154 if atomic.LoadInt32(&s.remoteGoAway) == 1 {
155 return nil, ErrRemoteGoAway
158 // Block if we have too many inflight SYNs
160 case s.synCh <- struct{}{}:
162 return nil, ErrSessionShutdown
166 // Get an ID, and check for stream exhaustion
167 id := atomic.LoadUint32(&s.nextStreamID)
168 if id >= math.MaxUint32-1 {
169 return nil, ErrStreamsExhausted
171 if !atomic.CompareAndSwapUint32(&s.nextStreamID, id, id+2) {
175 // Register the stream
176 stream := newStream(s, id, streamInit)
178 s.streams[id] = stream
179 s.inflight[id] = struct{}{}
180 s.streamLock.Unlock()
182 // Send the window update to create
183 if err := stream.sendWindowUpdate(); err != nil {
187 s.logger.Printf("[ERR] yamux: aborted stream open without inflight syn semaphore")
194 // Accept is used to block until the next available stream
195 // is ready to be accepted.
196 func (s *Session) Accept() (net.Conn, error) {
197 conn, err := s.AcceptStream()
204 // AcceptStream is used to block until the next available stream
205 // is ready to be accepted.
206 func (s *Session) AcceptStream() (*Stream, error) {
208 case stream := <-s.acceptCh:
209 if err := stream.sendWindowUpdate(); err != nil {
214 return nil, s.shutdownErr
218 // Close is used to close the session and all streams.
219 // Attempts to send a GoAway before closing the connection.
220 func (s *Session) Close() error {
221 s.shutdownLock.Lock()
222 defer s.shutdownLock.Unlock()
228 if s.shutdownErr == nil {
229 s.shutdownErr = ErrSessionShutdown
236 defer s.streamLock.Unlock()
237 for _, stream := range s.streams {
243 // exitErr is used to handle an error that is causing the
244 // session to terminate.
245 func (s *Session) exitErr(err error) {
246 s.shutdownLock.Lock()
247 if s.shutdownErr == nil {
250 s.shutdownLock.Unlock()
254 // GoAway can be used to prevent accepting further
255 // connections. It does not close the underlying conn.
256 func (s *Session) GoAway() error {
257 return s.waitForSend(s.goAway(goAwayNormal), nil)
260 // goAway is used to send a goAway message
261 func (s *Session) goAway(reason uint32) header {
262 atomic.SwapInt32(&s.localGoAway, 1)
263 hdr := header(make([]byte, headerSize))
264 hdr.encode(typeGoAway, 0, 0, reason)
268 // Ping is used to measure the RTT response time
269 func (s *Session) Ping() (time.Duration, error) {
270 // Get a channel for the ping
271 ch := make(chan struct{})
273 // Get a new ping id, mark as pending
280 // Send the ping request
281 hdr := header(make([]byte, headerSize))
282 hdr.encode(typePing, flagSYN, 0, id)
283 if err := s.waitForSend(hdr, nil); err != nil {
287 // Wait for a response
291 case <-time.After(s.config.ConnectionWriteTimeout):
293 delete(s.pings, id) // Ignore it if a response comes later.
297 return 0, ErrSessionShutdown
301 return time.Now().Sub(start), nil
304 // keepalive is a long running goroutine that periodically does
305 // a ping to keep the connection alive.
306 func (s *Session) keepalive() {
309 case <-time.After(s.config.KeepAliveInterval):
312 if err != ErrSessionShutdown {
313 s.logger.Printf("[ERR] yamux: keepalive failed: %v", err)
314 s.exitErr(ErrKeepAliveTimeout)
324 // waitForSendErr waits to send a header, checking for a potential shutdown
325 func (s *Session) waitForSend(hdr header, body io.Reader) error {
326 errCh := make(chan error, 1)
327 return s.waitForSendErr(hdr, body, errCh)
330 // waitForSendErr waits to send a header with optional data, checking for a
331 // potential shutdown. Since there's the expectation that sends can happen
332 // in a timely manner, we enforce the connection write timeout here.
333 func (s *Session) waitForSendErr(hdr header, body io.Reader, errCh chan error) error {
335 timer := t.(*time.Timer)
336 timer.Reset(s.config.ConnectionWriteTimeout)
346 ready := sendReady{Hdr: hdr, Body: body, Err: errCh}
348 case s.sendCh <- ready:
350 return ErrSessionShutdown
352 return ErrConnectionWriteTimeout
359 return ErrSessionShutdown
361 return ErrConnectionWriteTimeout
365 // sendNoWait does a send without waiting. Since there's the expectation that
366 // the send happens right here, we enforce the connection write timeout if we
367 // can't queue the header to be sent.
368 func (s *Session) sendNoWait(hdr header) error {
370 timer := t.(*time.Timer)
371 timer.Reset(s.config.ConnectionWriteTimeout)
382 case s.sendCh <- sendReady{Hdr: hdr}:
385 return ErrSessionShutdown
387 return ErrConnectionWriteTimeout
391 // send is a long running goroutine that sends data
392 func (s *Session) send() {
395 case ready := <-s.sendCh:
396 // Send a header if ready
397 if ready.Hdr != nil {
399 for sent < len(ready.Hdr) {
400 n, err := s.conn.Write(ready.Hdr[sent:])
402 s.logger.Printf("[ERR] yamux: Failed to write header: %v", err)
403 asyncSendErr(ready.Err, err)
411 // Send data from a body if given
412 if ready.Body != nil {
413 _, err := io.Copy(s.conn, ready.Body)
415 s.logger.Printf("[ERR] yamux: Failed to write body: %v", err)
416 asyncSendErr(ready.Err, err)
422 // No error, successful send
423 asyncSendErr(ready.Err, nil)
430 // recv is a long running goroutine that accepts new data
431 func (s *Session) recv() {
432 if err := s.recvLoop(); err != nil {
437 // Ensure that the index of the handler (typeData/typeWindowUpdate/etc) matches the message type
439 handlers = []func(*Session, header) error{
440 typeData: (*Session).handleStreamMessage,
441 typeWindowUpdate: (*Session).handleStreamMessage,
442 typePing: (*Session).handlePing,
443 typeGoAway: (*Session).handleGoAway,
447 // recvLoop continues to receive data until a fatal error is encountered
448 func (s *Session) recvLoop() error {
449 defer close(s.recvDoneCh)
450 hdr := header(make([]byte, headerSize))
453 if _, err := io.ReadFull(s.bufRead, hdr); err != nil {
454 if err != io.EOF && !strings.Contains(err.Error(), "closed") && !strings.Contains(err.Error(), "reset by peer") {
455 s.logger.Printf("[ERR] yamux: Failed to read header: %v", err)
460 // Verify the version
461 if hdr.Version() != protoVersion {
462 s.logger.Printf("[ERR] yamux: Invalid protocol version: %d", hdr.Version())
463 return ErrInvalidVersion
467 if mt < typeData || mt > typeGoAway {
468 return ErrInvalidMsgType
471 if err := handlers[mt](s, hdr); err != nil {
477 // handleStreamMessage handles either a data or window update frame
478 func (s *Session) handleStreamMessage(hdr header) error {
479 // Check for a new stream creation
482 if flags&flagSYN == flagSYN {
483 if err := s.incomingStream(id); err != nil {
490 stream := s.streams[id]
491 s.streamLock.Unlock()
493 // If we do not have a stream, likely we sent a RST
495 // Drain any data on the wire
496 if hdr.MsgType() == typeData && hdr.Length() > 0 {
497 s.logger.Printf("[WARN] yamux: Discarding data for stream: %d", id)
498 if _, err := io.CopyN(ioutil.Discard, s.bufRead, int64(hdr.Length())); err != nil {
499 s.logger.Printf("[ERR] yamux: Failed to discard data: %v", err)
503 s.logger.Printf("[WARN] yamux: frame for missing stream: %v", hdr)
508 // Check if this is a window update
509 if hdr.MsgType() == typeWindowUpdate {
510 if err := stream.incrSendWindow(hdr, flags); err != nil {
511 if sendErr := s.sendNoWait(s.goAway(goAwayProtoErr)); sendErr != nil {
512 s.logger.Printf("[WARN] yamux: failed to send go away: %v", sendErr)
520 if err := stream.readData(hdr, flags, s.bufRead); err != nil {
521 if sendErr := s.sendNoWait(s.goAway(goAwayProtoErr)); sendErr != nil {
522 s.logger.Printf("[WARN] yamux: failed to send go away: %v", sendErr)
529 // handlePing is invokde for a typePing frame
530 func (s *Session) handlePing(hdr header) error {
532 pingID := hdr.Length()
534 // Check if this is a query, respond back in a separate context so we
535 // don't interfere with the receiving thread blocking for the write.
536 if flags&flagSYN == flagSYN {
538 hdr := header(make([]byte, headerSize))
539 hdr.encode(typePing, flagACK, 0, pingID)
540 if err := s.sendNoWait(hdr); err != nil {
541 s.logger.Printf("[WARN] yamux: failed to send ping reply: %v", err)
549 ch := s.pings[pingID]
551 delete(s.pings, pingID)
558 // handleGoAway is invokde for a typeGoAway frame
559 func (s *Session) handleGoAway(hdr header) error {
563 atomic.SwapInt32(&s.remoteGoAway, 1)
565 s.logger.Printf("[ERR] yamux: received protocol error go away")
566 return fmt.Errorf("yamux protocol error")
567 case goAwayInternalErr:
568 s.logger.Printf("[ERR] yamux: received internal error go away")
569 return fmt.Errorf("remote yamux internal error")
571 s.logger.Printf("[ERR] yamux: received unexpected go away")
572 return fmt.Errorf("unexpected go away received")
577 // incomingStream is used to create a new incoming stream
578 func (s *Session) incomingStream(id uint32) error {
579 // Reject immediately if we are doing a go away
580 if atomic.LoadInt32(&s.localGoAway) == 1 {
581 hdr := header(make([]byte, headerSize))
582 hdr.encode(typeWindowUpdate, flagRST, id, 0)
583 return s.sendNoWait(hdr)
586 // Allocate a new stream
587 stream := newStream(s, id, streamSYNReceived)
590 defer s.streamLock.Unlock()
592 // Check if stream already exists
593 if _, ok := s.streams[id]; ok {
594 s.logger.Printf("[ERR] yamux: duplicate stream declared")
595 if sendErr := s.sendNoWait(s.goAway(goAwayProtoErr)); sendErr != nil {
596 s.logger.Printf("[WARN] yamux: failed to send go away: %v", sendErr)
598 return ErrDuplicateStream
601 // Register the stream
602 s.streams[id] = stream
604 // Check if we've exceeded the backlog
606 case s.acceptCh <- stream:
609 // Backlog exceeded! RST the stream
610 s.logger.Printf("[WARN] yamux: backlog exceeded, forcing connection reset")
611 delete(s.streams, id)
612 stream.sendHdr.encode(typeWindowUpdate, flagRST, id, 0)
613 return s.sendNoWait(stream.sendHdr)
617 // closeStream is used to close a stream once both sides have
618 // issued a close. If there was an in-flight SYN and the stream
619 // was not yet established, then this will give the credit back.
620 func (s *Session) closeStream(id uint32) {
622 if _, ok := s.inflight[id]; ok {
626 s.logger.Printf("[ERR] yamux: SYN tracking out of sync")
629 delete(s.streams, id)
630 s.streamLock.Unlock()
633 // establishStream is used to mark a stream that was in the
634 // SYN Sent state as established.
635 func (s *Session) establishStream(id uint32) {
637 if _, ok := s.inflight[id]; ok {
638 delete(s.inflight, id)
640 s.logger.Printf("[ERR] yamux: established stream without inflight SYN (no tracking entry)")
645 s.logger.Printf("[ERR] yamux: established stream without inflight SYN (didn't have semaphore)")
647 s.streamLock.Unlock()