123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749 |
- package yamux
- import (
- "bufio"
- "bytes"
- "context"
- "fmt"
- "io"
- "io/ioutil"
- "log"
- "math"
- "net"
- "strings"
- "sync"
- "sync/atomic"
- "time"
- )
- type Session struct {
-
-
- remoteGoAway int32
-
-
- localGoAway int32
-
-
- nextStreamID uint32
-
- config *Config
-
- logger Logger
-
- conn io.ReadWriteCloser
-
- bufRead *bufio.Reader
-
- pings map[uint32]chan struct{}
- pingID uint32
- pingLock sync.Mutex
-
-
-
- streams map[uint32]*Stream
- inflight map[uint32]struct{}
- streamLock sync.Mutex
-
-
-
- synCh chan struct{}
-
- acceptCh chan *Stream
-
-
- sendCh chan *sendReady
-
-
- recvDoneCh chan struct{}
- sendDoneCh chan struct{}
-
- shutdown bool
- shutdownErr error
- shutdownCh chan struct{}
- shutdownLock sync.Mutex
- shutdownErrLock sync.Mutex
- }
- type sendReady struct {
- Hdr []byte
- mu sync.Mutex
- Body []byte
- Err chan error
- }
- func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
- logger := config.Logger
- if logger == nil {
- logger = log.New(config.LogOutput, "", log.LstdFlags)
- }
- s := &Session{
- config: config,
- logger: logger,
- conn: conn,
- bufRead: bufio.NewReader(conn),
- pings: make(map[uint32]chan struct{}),
- streams: make(map[uint32]*Stream),
- inflight: make(map[uint32]struct{}),
- synCh: make(chan struct{}, config.AcceptBacklog),
- acceptCh: make(chan *Stream, config.AcceptBacklog),
- sendCh: make(chan *sendReady, 64),
- recvDoneCh: make(chan struct{}),
- sendDoneCh: make(chan struct{}),
- shutdownCh: make(chan struct{}),
- }
- if client {
- s.nextStreamID = 1
- } else {
- s.nextStreamID = 2
- }
- go s.recv()
- go s.send()
- if config.EnableKeepAlive {
- go s.keepalive()
- }
- return s
- }
- func (s *Session) IsClosed() bool {
- select {
- case <-s.shutdownCh:
- return true
- default:
- return false
- }
- }
- func (s *Session) CloseChan() <-chan struct{} {
- return s.shutdownCh
- }
- func (s *Session) NumStreams() int {
- s.streamLock.Lock()
- num := len(s.streams)
- s.streamLock.Unlock()
- return num
- }
- func (s *Session) Open() (net.Conn, error) {
- conn, err := s.OpenStream()
- if err != nil {
- return nil, err
- }
- return conn, nil
- }
- func (s *Session) OpenStream() (*Stream, error) {
- if s.IsClosed() {
- return nil, ErrSessionShutdown
- }
- if atomic.LoadInt32(&s.remoteGoAway) == 1 {
- return nil, ErrRemoteGoAway
- }
-
- select {
- case s.synCh <- struct{}{}:
- case <-s.shutdownCh:
- return nil, ErrSessionShutdown
- }
- GET_ID:
-
- id := atomic.LoadUint32(&s.nextStreamID)
- if id >= math.MaxUint32-1 {
- return nil, ErrStreamsExhausted
- }
- if !atomic.CompareAndSwapUint32(&s.nextStreamID, id, id+2) {
- goto GET_ID
- }
-
- stream := newStream(s, id, streamInit)
- s.streamLock.Lock()
- s.streams[id] = stream
- s.inflight[id] = struct{}{}
- s.streamLock.Unlock()
- if s.config.StreamOpenTimeout > 0 {
- go s.setOpenTimeout(stream)
- }
-
- if err := stream.sendWindowUpdate(); err != nil {
- select {
- case <-s.synCh:
- default:
- s.logger.Printf("[ERR] yamux: aborted stream open without inflight syn semaphore")
- }
- return nil, err
- }
- return stream, nil
- }
- func (s *Session) setOpenTimeout(stream *Stream) {
- timer := time.NewTimer(s.config.StreamOpenTimeout)
- defer timer.Stop()
- select {
- case <-stream.establishCh:
- return
- case <-s.shutdownCh:
- return
- case <-timer.C:
-
-
- s.logger.Printf("[ERR] yamux: aborted stream open (destination=%s): %v", s.RemoteAddr().String(), ErrTimeout.err)
- s.Close()
- }
- }
- func (s *Session) Accept() (net.Conn, error) {
- conn, err := s.AcceptStream()
- if err != nil {
- return nil, err
- }
- return conn, err
- }
- func (s *Session) AcceptStream() (*Stream, error) {
- select {
- case stream := <-s.acceptCh:
- if err := stream.sendWindowUpdate(); err != nil {
- return nil, err
- }
- return stream, nil
- case <-s.shutdownCh:
- return nil, s.shutdownErr
- }
- }
- func (s *Session) AcceptStreamWithContext(ctx context.Context) (*Stream, error) {
- select {
- case <-ctx.Done():
- return nil, ctx.Err()
- case stream := <-s.acceptCh:
- if err := stream.sendWindowUpdate(); err != nil {
- return nil, err
- }
- return stream, nil
- case <-s.shutdownCh:
- return nil, s.shutdownErr
- }
- }
- func (s *Session) Close() error {
- s.shutdownLock.Lock()
- defer s.shutdownLock.Unlock()
- if s.shutdown {
- return nil
- }
- s.shutdown = true
- s.shutdownErrLock.Lock()
- if s.shutdownErr == nil {
- s.shutdownErr = ErrSessionShutdown
- }
- s.shutdownErrLock.Unlock()
- close(s.shutdownCh)
- s.conn.Close()
- <-s.recvDoneCh
- s.streamLock.Lock()
- defer s.streamLock.Unlock()
- for _, stream := range s.streams {
- stream.forceClose()
- }
- <-s.sendDoneCh
- return nil
- }
- func (s *Session) exitErr(err error) {
- s.shutdownErrLock.Lock()
- if s.shutdownErr == nil {
- s.shutdownErr = err
- }
- s.shutdownErrLock.Unlock()
- s.Close()
- }
- func (s *Session) GoAway() error {
- return s.waitForSend(s.goAway(goAwayNormal), nil)
- }
- func (s *Session) goAway(reason uint32) header {
- atomic.SwapInt32(&s.localGoAway, 1)
- hdr := header(make([]byte, headerSize))
- hdr.encode(typeGoAway, 0, 0, reason)
- return hdr
- }
- func (s *Session) Ping() (time.Duration, error) {
-
- ch := make(chan struct{})
-
- s.pingLock.Lock()
- id := s.pingID
- s.pingID++
- s.pings[id] = ch
- s.pingLock.Unlock()
-
- hdr := header(make([]byte, headerSize))
- hdr.encode(typePing, flagSYN, 0, id)
- if err := s.waitForSend(hdr, nil); err != nil {
- return 0, err
- }
-
- start := time.Now()
- select {
- case <-ch:
- case <-time.After(s.config.ConnectionWriteTimeout):
- s.pingLock.Lock()
- delete(s.pings, id)
- s.pingLock.Unlock()
- return 0, ErrTimeout
- case <-s.shutdownCh:
- return 0, ErrSessionShutdown
- }
-
- return time.Now().Sub(start), nil
- }
- func (s *Session) keepalive() {
- for {
- select {
- case <-time.After(s.config.KeepAliveInterval):
- _, err := s.Ping()
- if err != nil {
- if err != ErrSessionShutdown {
- s.logger.Printf("[ERR] yamux: keepalive failed: %v", err)
- s.exitErr(ErrKeepAliveTimeout)
- }
- return
- }
- case <-s.shutdownCh:
- return
- }
- }
- }
- func (s *Session) waitForSend(hdr header, body []byte) error {
- errCh := make(chan error, 1)
- return s.waitForSendErr(hdr, body, errCh)
- }
- func (s *Session) waitForSendErr(hdr header, body []byte, errCh chan error) error {
- t := timerPool.Get()
- timer := t.(*time.Timer)
- timer.Reset(s.config.ConnectionWriteTimeout)
- defer func() {
- timer.Stop()
- select {
- case <-timer.C:
- default:
- }
- timerPool.Put(t)
- }()
- ready := &sendReady{Hdr: hdr, Body: body, Err: errCh}
- select {
- case s.sendCh <- ready:
- case <-s.shutdownCh:
- return ErrSessionShutdown
- case <-timer.C:
- return ErrConnectionWriteTimeout
- }
- bodyCopy := func() {
- if body == nil {
- return
- }
-
-
-
-
- ready.mu.Lock()
- defer ready.mu.Unlock()
- if ready.Body == nil {
- return
- }
- newBody := make([]byte, len(body))
- copy(newBody, body)
- ready.Body = newBody
- }
- select {
- case err := <-errCh:
- return err
- case <-s.shutdownCh:
- bodyCopy()
- return ErrSessionShutdown
- case <-timer.C:
- bodyCopy()
- return ErrConnectionWriteTimeout
- }
- }
- func (s *Session) sendNoWait(hdr header) error {
- t := timerPool.Get()
- timer := t.(*time.Timer)
- timer.Reset(s.config.ConnectionWriteTimeout)
- defer func() {
- timer.Stop()
- select {
- case <-timer.C:
- default:
- }
- timerPool.Put(t)
- }()
- select {
- case s.sendCh <- &sendReady{Hdr: hdr}:
- return nil
- case <-s.shutdownCh:
- return ErrSessionShutdown
- case <-timer.C:
- return ErrConnectionWriteTimeout
- }
- }
- func (s *Session) send() {
- if err := s.sendLoop(); err != nil {
- s.exitErr(err)
- }
- }
- func (s *Session) sendLoop() error {
- defer close(s.sendDoneCh)
- var bodyBuf bytes.Buffer
- for {
- bodyBuf.Reset()
- select {
- case ready := <-s.sendCh:
-
- if ready.Hdr != nil {
- _, err := s.conn.Write(ready.Hdr)
- if err != nil {
- s.logger.Printf("[ERR] yamux: Failed to write header: %v", err)
- asyncSendErr(ready.Err, err)
- return err
- }
- }
- ready.mu.Lock()
- if ready.Body != nil {
-
-
- _, err := bodyBuf.Write(ready.Body)
- if err != nil {
- ready.Body = nil
- ready.mu.Unlock()
- s.logger.Printf("[ERR] yamux: Failed to copy body into buffer: %v", err)
- asyncSendErr(ready.Err, err)
- return err
- }
- ready.Body = nil
- }
- ready.mu.Unlock()
- if bodyBuf.Len() > 0 {
-
- _, err := s.conn.Write(bodyBuf.Bytes())
- if err != nil {
- s.logger.Printf("[ERR] yamux: Failed to write body: %v", err)
- asyncSendErr(ready.Err, err)
- return err
- }
- }
-
- asyncSendErr(ready.Err, nil)
- case <-s.shutdownCh:
- return nil
- }
- }
- }
- func (s *Session) recv() {
- if err := s.recvLoop(); err != nil {
- s.exitErr(err)
- }
- }
- var (
- handlers = []func(*Session, header) error{
- typeData: (*Session).handleStreamMessage,
- typeWindowUpdate: (*Session).handleStreamMessage,
- typePing: (*Session).handlePing,
- typeGoAway: (*Session).handleGoAway,
- }
- )
- func (s *Session) recvLoop() error {
- defer close(s.recvDoneCh)
- hdr := header(make([]byte, headerSize))
- for {
-
- if _, err := io.ReadFull(s.bufRead, hdr); err != nil {
- if err != io.EOF && !strings.Contains(err.Error(), "closed") && !strings.Contains(err.Error(), "reset by peer") {
- s.logger.Printf("[ERR] yamux: Failed to read header: %v", err)
- }
- return err
- }
-
- if hdr.Version() != protoVersion {
- s.logger.Printf("[ERR] yamux: Invalid protocol version: %d", hdr.Version())
- return ErrInvalidVersion
- }
- mt := hdr.MsgType()
- if mt < typeData || mt > typeGoAway {
- return ErrInvalidMsgType
- }
- if err := handlers[mt](s, hdr); err != nil {
- return err
- }
- }
- }
- func (s *Session) handleStreamMessage(hdr header) error {
-
- id := hdr.StreamID()
- flags := hdr.Flags()
- if flags&flagSYN == flagSYN {
- if err := s.incomingStream(id); err != nil {
- return err
- }
- }
-
- s.streamLock.Lock()
- stream := s.streams[id]
- s.streamLock.Unlock()
-
- if stream == nil {
-
- if hdr.MsgType() == typeData && hdr.Length() > 0 {
- s.logger.Printf("[WARN] yamux: Discarding data for stream: %d", id)
- if _, err := io.CopyN(ioutil.Discard, s.bufRead, int64(hdr.Length())); err != nil {
- s.logger.Printf("[ERR] yamux: Failed to discard data: %v", err)
- return nil
- }
- } else {
- s.logger.Printf("[WARN] yamux: frame for missing stream: %v", hdr)
- }
- return nil
- }
-
- if hdr.MsgType() == typeWindowUpdate {
- if err := stream.incrSendWindow(hdr, flags); err != nil {
- if sendErr := s.sendNoWait(s.goAway(goAwayProtoErr)); sendErr != nil {
- s.logger.Printf("[WARN] yamux: failed to send go away: %v", sendErr)
- }
- return err
- }
- return nil
- }
-
- if err := stream.readData(hdr, flags, s.bufRead); err != nil {
- if sendErr := s.sendNoWait(s.goAway(goAwayProtoErr)); sendErr != nil {
- s.logger.Printf("[WARN] yamux: failed to send go away: %v", sendErr)
- }
- return err
- }
- return nil
- }
- func (s *Session) handlePing(hdr header) error {
- flags := hdr.Flags()
- pingID := hdr.Length()
-
-
- if flags&flagSYN == flagSYN {
- go func() {
- hdr := header(make([]byte, headerSize))
- hdr.encode(typePing, flagACK, 0, pingID)
- if err := s.sendNoWait(hdr); err != nil {
- s.logger.Printf("[WARN] yamux: failed to send ping reply: %v", err)
- }
- }()
- return nil
- }
-
- s.pingLock.Lock()
- ch := s.pings[pingID]
- if ch != nil {
- delete(s.pings, pingID)
- close(ch)
- }
- s.pingLock.Unlock()
- return nil
- }
- func (s *Session) handleGoAway(hdr header) error {
- code := hdr.Length()
- switch code {
- case goAwayNormal:
- atomic.SwapInt32(&s.remoteGoAway, 1)
- case goAwayProtoErr:
- s.logger.Printf("[ERR] yamux: received protocol error go away")
- return fmt.Errorf("yamux protocol error")
- case goAwayInternalErr:
- s.logger.Printf("[ERR] yamux: received internal error go away")
- return fmt.Errorf("remote yamux internal error")
- default:
- s.logger.Printf("[ERR] yamux: received unexpected go away")
- return fmt.Errorf("unexpected go away received")
- }
- return nil
- }
- func (s *Session) incomingStream(id uint32) error {
-
- if atomic.LoadInt32(&s.localGoAway) == 1 {
- hdr := header(make([]byte, headerSize))
- hdr.encode(typeWindowUpdate, flagRST, id, 0)
- return s.sendNoWait(hdr)
- }
-
- stream := newStream(s, id, streamSYNReceived)
- s.streamLock.Lock()
- defer s.streamLock.Unlock()
-
- if _, ok := s.streams[id]; ok {
- s.logger.Printf("[ERR] yamux: duplicate stream declared")
- if sendErr := s.sendNoWait(s.goAway(goAwayProtoErr)); sendErr != nil {
- s.logger.Printf("[WARN] yamux: failed to send go away: %v", sendErr)
- }
- return ErrDuplicateStream
- }
-
- s.streams[id] = stream
-
- select {
- case s.acceptCh <- stream:
- return nil
- default:
-
- s.logger.Printf("[WARN] yamux: backlog exceeded, forcing connection reset")
- delete(s.streams, id)
- hdr := header(make([]byte, headerSize))
- hdr.encode(typeWindowUpdate, flagRST, id, 0)
- return s.sendNoWait(hdr)
- }
- }
- func (s *Session) closeStream(id uint32) {
- s.streamLock.Lock()
- if _, ok := s.inflight[id]; ok {
- select {
- case <-s.synCh:
- default:
- s.logger.Printf("[ERR] yamux: SYN tracking out of sync")
- }
- }
- delete(s.streams, id)
- s.streamLock.Unlock()
- }
- func (s *Session) establishStream(id uint32) {
- s.streamLock.Lock()
- if _, ok := s.inflight[id]; ok {
- delete(s.inflight, id)
- } else {
- s.logger.Printf("[ERR] yamux: established stream without inflight SYN (no tracking entry)")
- }
- select {
- case <-s.synCh:
- default:
- s.logger.Printf("[ERR] yamux: established stream without inflight SYN (didn't have semaphore)")
- }
- s.streamLock.Unlock()
- }
|