2017-04-12 20:31:11 +00:00
|
|
|
package mux
|
|
|
|
|
|
|
|
import (
|
|
|
|
"sync"
|
|
|
|
|
2018-04-16 22:31:10 +00:00
|
|
|
"v2ray.com/core/common"
|
2017-05-02 20:23:07 +00:00
|
|
|
"v2ray.com/core/common/buf"
|
|
|
|
"v2ray.com/core/common/protocol"
|
2017-04-12 20:31:11 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
type SessionManager struct {
|
|
|
|
sync.RWMutex
|
|
|
|
sessions map[uint16]*Session
|
2017-04-19 19:27:21 +00:00
|
|
|
count uint16
|
2017-04-13 18:14:07 +00:00
|
|
|
closed bool
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func NewSessionManager() *SessionManager {
|
|
|
|
return &SessionManager{
|
|
|
|
count: 0,
|
2017-10-27 09:40:18 +00:00
|
|
|
sessions: make(map[uint16]*Session, 16),
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2018-10-29 11:51:56 +00:00
|
|
|
func (m *SessionManager) Closed() bool {
|
|
|
|
m.RLock()
|
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
|
|
return m.closed
|
|
|
|
}
|
|
|
|
|
2017-04-12 20:31:11 +00:00
|
|
|
func (m *SessionManager) Size() int {
|
|
|
|
m.RLock()
|
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
|
|
return len(m.sessions)
|
|
|
|
}
|
|
|
|
|
2017-04-19 19:27:21 +00:00
|
|
|
func (m *SessionManager) Count() int {
|
|
|
|
m.RLock()
|
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
|
|
return int(m.count)
|
|
|
|
}
|
|
|
|
|
2017-04-13 18:14:07 +00:00
|
|
|
func (m *SessionManager) Allocate() *Session {
|
2017-04-12 20:31:11 +00:00
|
|
|
m.Lock()
|
|
|
|
defer m.Unlock()
|
|
|
|
|
2017-04-13 18:14:07 +00:00
|
|
|
if m.closed {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2017-04-12 20:31:11 +00:00
|
|
|
m.count++
|
2017-04-13 18:14:07 +00:00
|
|
|
s := &Session{
|
|
|
|
ID: m.count,
|
|
|
|
parent: m,
|
|
|
|
}
|
2017-04-12 20:31:11 +00:00
|
|
|
m.sessions[s.ID] = s
|
2017-04-13 18:14:07 +00:00
|
|
|
return s
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (m *SessionManager) Add(s *Session) {
|
|
|
|
m.Lock()
|
|
|
|
defer m.Unlock()
|
|
|
|
|
2017-10-27 09:40:18 +00:00
|
|
|
if m.closed {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2018-10-28 08:08:43 +00:00
|
|
|
m.count++
|
2017-04-12 20:31:11 +00:00
|
|
|
m.sessions[s.ID] = s
|
|
|
|
}
|
|
|
|
|
|
|
|
func (m *SessionManager) Remove(id uint16) {
|
|
|
|
m.Lock()
|
|
|
|
defer m.Unlock()
|
|
|
|
|
2017-10-27 09:40:18 +00:00
|
|
|
if m.closed {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
2017-04-12 20:31:11 +00:00
|
|
|
delete(m.sessions, id)
|
2018-04-07 21:07:30 +00:00
|
|
|
|
|
|
|
if len(m.sessions) == 0 {
|
|
|
|
m.sessions = make(map[uint16]*Session, 16)
|
|
|
|
}
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (m *SessionManager) Get(id uint16) (*Session, bool) {
|
|
|
|
m.RLock()
|
|
|
|
defer m.RUnlock()
|
|
|
|
|
2017-04-13 18:14:07 +00:00
|
|
|
if m.closed {
|
|
|
|
return nil, false
|
|
|
|
}
|
|
|
|
|
2017-04-12 20:31:11 +00:00
|
|
|
s, found := m.sessions[id]
|
|
|
|
return s, found
|
|
|
|
}
|
|
|
|
|
2017-04-13 18:14:07 +00:00
|
|
|
func (m *SessionManager) CloseIfNoSession() bool {
|
2017-04-19 19:27:21 +00:00
|
|
|
m.Lock()
|
|
|
|
defer m.Unlock()
|
2017-04-12 20:31:11 +00:00
|
|
|
|
2017-04-13 18:56:32 +00:00
|
|
|
if m.closed {
|
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2017-04-19 10:36:04 +00:00
|
|
|
if len(m.sessions) != 0 {
|
2017-04-13 18:14:07 +00:00
|
|
|
return false
|
|
|
|
}
|
|
|
|
|
2017-04-13 18:56:32 +00:00
|
|
|
m.closed = true
|
|
|
|
return true
|
|
|
|
}
|
|
|
|
|
2018-02-08 14:39:46 +00:00
|
|
|
func (m *SessionManager) Close() error {
|
2017-04-19 19:27:21 +00:00
|
|
|
m.Lock()
|
|
|
|
defer m.Unlock()
|
2017-04-13 18:56:32 +00:00
|
|
|
|
|
|
|
if m.closed {
|
2018-02-08 14:39:46 +00:00
|
|
|
return nil
|
2017-04-13 18:56:32 +00:00
|
|
|
}
|
|
|
|
|
2017-04-13 18:14:07 +00:00
|
|
|
m.closed = true
|
|
|
|
|
2017-04-12 20:31:11 +00:00
|
|
|
for _, s := range m.sessions {
|
2018-05-26 14:32:55 +00:00
|
|
|
common.Close(s.input) // nolint: errcheck
|
|
|
|
common.Close(s.output) // nolint: errcheck
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
2017-04-13 18:14:07 +00:00
|
|
|
|
2017-10-27 09:40:18 +00:00
|
|
|
m.sessions = nil
|
2018-02-08 14:39:46 +00:00
|
|
|
return nil
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
|
|
|
|
2017-10-30 11:36:31 +00:00
|
|
|
// Session represents a client connection in a Mux connection.
|
2017-04-12 20:31:11 +00:00
|
|
|
type Session struct {
|
2018-04-16 22:31:10 +00:00
|
|
|
input buf.Reader
|
|
|
|
output buf.Writer
|
2017-05-02 21:53:39 +00:00
|
|
|
parent *SessionManager
|
|
|
|
ID uint16
|
|
|
|
transferType protocol.TransferType
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
|
|
|
|
2017-10-21 18:40:31 +00:00
|
|
|
// Close closes all resources associated with this session.
|
2018-02-08 14:39:46 +00:00
|
|
|
func (s *Session) Close() error {
|
2018-05-26 14:32:55 +00:00
|
|
|
common.Close(s.output) // nolint: errcheck
|
|
|
|
common.Close(s.input) // nolint: errcheck
|
2017-05-02 21:53:39 +00:00
|
|
|
s.parent.Remove(s.ID)
|
2018-02-08 14:39:46 +00:00
|
|
|
return nil
|
2017-04-12 20:31:11 +00:00
|
|
|
}
|
2017-05-02 20:23:07 +00:00
|
|
|
|
2017-10-30 11:36:31 +00:00
|
|
|
// NewReader creates a buf.Reader based on the transfer type of this Session.
|
2017-11-25 00:05:30 +00:00
|
|
|
func (s *Session) NewReader(reader *buf.BufferedReader) buf.Reader {
|
2017-05-02 20:23:07 +00:00
|
|
|
if s.transferType == protocol.TransferTypeStream {
|
|
|
|
return NewStreamReader(reader)
|
|
|
|
}
|
|
|
|
return NewPacketReader(reader)
|
|
|
|
}
|