github.com/EagleQL/Xray-core@v1.4.3/common/mux/session.go (about)

     1  package mux
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/xtls/xray-core/common"
     7  	"github.com/xtls/xray-core/common/buf"
     8  	"github.com/xtls/xray-core/common/net"
     9  	"github.com/xtls/xray-core/common/protocol"
    10  )
    11  
    12  type SessionManager struct {
    13  	sync.RWMutex
    14  	sessions map[uint16]*Session
    15  	count    uint16
    16  	closed   bool
    17  }
    18  
    19  func NewSessionManager() *SessionManager {
    20  	return &SessionManager{
    21  		count:    0,
    22  		sessions: make(map[uint16]*Session, 16),
    23  	}
    24  }
    25  
    26  func (m *SessionManager) Closed() bool {
    27  	m.RLock()
    28  	defer m.RUnlock()
    29  
    30  	return m.closed
    31  }
    32  
    33  func (m *SessionManager) Size() int {
    34  	m.RLock()
    35  	defer m.RUnlock()
    36  
    37  	return len(m.sessions)
    38  }
    39  
    40  func (m *SessionManager) Count() int {
    41  	m.RLock()
    42  	defer m.RUnlock()
    43  
    44  	return int(m.count)
    45  }
    46  
    47  func (m *SessionManager) Allocate() *Session {
    48  	m.Lock()
    49  	defer m.Unlock()
    50  
    51  	if m.closed {
    52  		return nil
    53  	}
    54  
    55  	m.count++
    56  	s := &Session{
    57  		ID:     m.count,
    58  		parent: m,
    59  	}
    60  	m.sessions[s.ID] = s
    61  	return s
    62  }
    63  
    64  func (m *SessionManager) Add(s *Session) {
    65  	m.Lock()
    66  	defer m.Unlock()
    67  
    68  	if m.closed {
    69  		return
    70  	}
    71  
    72  	m.count++
    73  	m.sessions[s.ID] = s
    74  }
    75  
    76  func (m *SessionManager) Remove(id uint16) {
    77  	m.Lock()
    78  	defer m.Unlock()
    79  
    80  	if m.closed {
    81  		return
    82  	}
    83  
    84  	delete(m.sessions, id)
    85  
    86  	if len(m.sessions) == 0 {
    87  		m.sessions = make(map[uint16]*Session, 16)
    88  	}
    89  }
    90  
    91  func (m *SessionManager) Get(id uint16) (*Session, bool) {
    92  	m.RLock()
    93  	defer m.RUnlock()
    94  
    95  	if m.closed {
    96  		return nil, false
    97  	}
    98  
    99  	s, found := m.sessions[id]
   100  	return s, found
   101  }
   102  
   103  func (m *SessionManager) CloseIfNoSession() bool {
   104  	m.Lock()
   105  	defer m.Unlock()
   106  
   107  	if m.closed {
   108  		return true
   109  	}
   110  
   111  	if len(m.sessions) != 0 {
   112  		return false
   113  	}
   114  
   115  	m.closed = true
   116  	return true
   117  }
   118  
   119  func (m *SessionManager) Close() error {
   120  	m.Lock()
   121  	defer m.Unlock()
   122  
   123  	if m.closed {
   124  		return nil
   125  	}
   126  
   127  	m.closed = true
   128  
   129  	for _, s := range m.sessions {
   130  		common.Close(s.input)
   131  		common.Close(s.output)
   132  	}
   133  
   134  	m.sessions = nil
   135  	return nil
   136  }
   137  
   138  // Session represents a client connection in a Mux connection.
   139  type Session struct {
   140  	input        buf.Reader
   141  	output       buf.Writer
   142  	parent       *SessionManager
   143  	ID           uint16
   144  	transferType protocol.TransferType
   145  }
   146  
   147  // Close closes all resources associated with this session.
   148  func (s *Session) Close() error {
   149  	common.Close(s.output)
   150  	common.Close(s.input)
   151  	s.parent.Remove(s.ID)
   152  	return nil
   153  }
   154  
   155  // NewReader creates a buf.Reader based on the transfer type of this Session.
   156  func (s *Session) NewReader(reader *buf.BufferedReader, dest *net.Destination) buf.Reader {
   157  	if s.transferType == protocol.TransferTypeStream {
   158  		return NewStreamReader(reader)
   159  	}
   160  	return NewPacketReader(reader, dest)
   161  }