From 2d17133a7b38457ff5ea1d9246dccc411b864b03 Mon Sep 17 00:00:00 2001
From: Mrs4s <mrs4sxiaoshi@gmail.com>
Date: Tue, 7 Dec 2021 13:29:01 +0800
Subject: [PATCH] fix: reconnect on decrypt flag error

---
 client/network.go   | 3 +++
 utils/connection.go | 8 ++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/client/network.go b/client/network.go
index e414852c..7ae16f86 100644
--- a/client/network.go
+++ b/client/network.go
@@ -311,6 +311,9 @@ func (c *QQClient) netLoop() {
 			pkt.Payload, err = pkt.DecryptPayload(c.ecdh.InitialShareKey, c.RandomKey, c.sigInfo.wtSessionTicketKey)
 			if err != nil {
 				c.Error("decrypt payload error: %v", err)
+				if errors.Is(err, packets.ErrUnknownFlag) {
+					go c.quickReconnect()
+				}
 				continue
 			}
 		}
diff --git a/utils/connection.go b/utils/connection.go
index 6172c87a..77f66da8 100644
--- a/utils/connection.go
+++ b/utils/connection.go
@@ -12,6 +12,7 @@ import (
 type TCPListener struct {
 	lock                 sync.RWMutex
 	conn                 net.Conn
+	connected            bool
 	plannedDisconnect    func(*TCPListener)
 	unexpectedDisconnect func(*TCPListener, error)
 }
@@ -42,6 +43,7 @@ func (t *TCPListener) Connect(addr *net.TCPAddr) error {
 	t.lock.Lock()
 	defer t.lock.Unlock()
 	t.conn = conn
+	t.connected = true
 	return nil
 }
 
@@ -103,16 +105,18 @@ func (t *TCPListener) close() {
 func (t *TCPListener) invokePlannedDisconnect() {
 	t.lock.RLock()
 	defer t.lock.RUnlock()
-	if t.plannedDisconnect != nil {
+	if t.plannedDisconnect != nil && t.connected {
 		go t.plannedDisconnect(t)
+		t.connected = false
 	}
 }
 
 func (t *TCPListener) invokeUnexpectedDisconnect(err error) {
 	t.lock.RLock()
 	defer t.lock.RUnlock()
-	if t.unexpectedDisconnect != nil {
+	if t.unexpectedDisconnect != nil && t.connected {
 		go t.unexpectedDisconnect(t, err)
+		t.connected = false
 	}
 }