fix: anytls client close (#1871)

anytls · web-flow · commit a7e56f1c431b · 2025-03-02T10:47:10.000+08:00
Co-authored-by: anytls &lt;anytls&gt;
diff --git a/transport/anytls/session/client.go b/transport/anytls/session/client.go
@@ -21,10 +21,14 @@ type Client struct {
 
 	dialOut util.DialOutFunc
 
-	sessionCounter  atomic.Uint64
+	sessionCounter atomic.Uint64
+
 	idleSession     *skiplist.SkipList[uint64, *Session]
 	idleSessionLock sync.Mutex
 
+	sessions     map[uint64]*Session
+	sessionsLock sync.Mutex
+
 	padding *atomic.TypedValue[*padding.PaddingFactory]
 
 	idleSessionTimeout time.Duration
@@ -33,6 +37,7 @@ type Client struct {
 
 func NewClient(ctx context.Context, dialOut util.DialOutFunc, _padding *atomic.TypedValue[*padding.PaddingFactory], idleSessionCheckInterval, idleSessionTimeout time.Duration, minIdleSession int) *Client {
 	c := &Client{
+		sessions:           make(map[uint64]*Session),
 		dialOut:            dialOut,
 		padding:            _padding,
 		idleSessionTimeout: idleSessionTimeout,
@@ -130,15 +135,35 @@ func (c *Client) createSession(ctx context.Context) (*Session, error) {
 		c.idleSessionLock.Lock()
 		c.idleSession.Remove(math.MaxUint64 - session.seq)
 		c.idleSessionLock.Unlock()
+
+		c.sessionsLock.Lock()
+		delete(c.sessions, session.seq)
+		c.sessionsLock.Unlock()
 	}
+
+	c.sessionsLock.Lock()
+	c.sessions[session.seq] = session
+	c.sessionsLock.Unlock()
+
 	session.Run()
 	return session, nil
 }
 
 func (c *Client) Close() error {
 	c.dieCancel()
-	c.minIdleSession = 0
-	go c.idleCleanupExpTime(time.Time{})
+
+	c.sessionsLock.Lock()
+	sessionToClose := make([]*Session, 0, len(c.sessions))
+	for seq, session := range c.sessions {
+		sessionToClose = append(sessionToClose, session)
+		delete(c.sessions, seq)
+	}
+	c.sessionsLock.Unlock()
+
+	for _, session := range sessionToClose {
+		session.Close()
+	}
+
 	return nil
 }
 
