github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/mux/session.go (about)

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