syncmap

Author: lz1998

pkg/bot/bot.go

@@ -8,10 +8,11 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+//go:generate go run github.com/a8m/syncmap -o "gen_client_map.go" -pkg bot -name ClientMap "map[int64]*client.QQClient"
+//go:generate go run github.com/a8m/syncmap -o "gen_token_map.go" -pkg bot -name TokenMap "map[int64][]byte"
 var (
-	// TODO sync
-	Clients     = map[int64]*client.QQClient{}
-	LoginTokens = map[int64][]byte{}
+	Clients     ClientMap
+	LoginTokens TokenMap
 )
 
 func InitLog(cli *client.QQClient) {
@@ -48,7 +49,7 @@ func Login(cli *client.QQClient) (bool, error) {
 }
 
 func SetRelogin(cli *client.QQClient, retryInterval int, retryCount int) {
-	LoginTokens[cli.Uin] = cli.GenToken()
+	LoginTokens.Store(cli.Uin, cli.GenToken())
 	cli.OnDisconnected(func(bot *client.QQClient, e *client.ClientDisconnectedEvent) {
 		if bot.Online {
 			return
@@ -68,14 +69,16 @@ func SetRelogin(cli *client.QQClient, retryInterval int, retryCount int) {
 			times++
 			time.Sleep(time.Second * time.Duration(retryInterval))
 
-			// 尝试token登录
-			if err := bot.TokenLogin(LoginTokens[bot.Uin]); err != nil {
-				log.Errorf("failed to relogin with token, try to login with password, %+v", err)
-				bot.Disconnect()
-			} else {
-				LoginTokens[cli.Uin] = bot.GenToken()
-				log.Info("succeed to relogin with token")
-				return
+			if token, ok := LoginTokens.Load(bot.Uin); ok {
+				// 尝试token登录
+				if err := bot.TokenLogin(token); err != nil {
+					log.Errorf("failed to relogin with token, try to login with password, %+v", err)
+					bot.Disconnect()
+				} else {
+					LoginTokens.Store(bot.Uin, bot.GenToken())
+					log.Info("succeed to relogin with token")
+					return
+				}
 			}
 
 			time.Sleep(time.Second)
@@ -89,7 +92,7 @@ func SetRelogin(cli *client.QQClient, retryInterval int, retryCount int) {
 				continue
 			}
 			if ok {
-				LoginTokens[bot.Uin] = bot.GenToken()
+				LoginTokens.Store(bot.Uin, bot.GenToken())
 				log.Info("重连成功")
 				return
 			}
@@ -102,15 +105,17 @@ func SetRelogin(cli *client.QQClient, retryInterval int, retryCount int) {
 // ReleaseClient 断开连接并释放资源
 func ReleaseClient(cli *client.QQClient) {
 	cli.Release()
-	delete(Clients, cli.Uin) // 必须先删Clients，影响IsClientExist
-	delete(LoginTokens, cli.Uin)
-	for _, wsServer := range RemoteServers[cli.Uin] {
-		wsServer.Close()
+	Clients.Delete(cli.Uin) // 必须先删Clients，影响IsClientExist
+	LoginTokens.Delete(cli.Uin)
+	if wsServers, ok := RemoteServers.Load(cli.Uin); ok {
+		for _, wsServer := range wsServers {
+			wsServer.Close()
+		}
 	}
-	delete(RemoteServers, cli.Uin)
+	RemoteServers.Delete(cli.Uin)
 }
 
 func IsClientExist(uin int64) bool {
-	_, ok := Clients[uin]
+	_, ok := Clients.Load(uin)
 	return ok
 }

pkg/bot/captcha.go

@@ -20,12 +20,14 @@ type WaitingCaptcha struct {
 	Prom    *promise.Promise
 }
 
+//go:generate go run github.com/a8m/syncmap -o "gen_captcha_map.go" -pkg bot -name CaptchaMap "map[int64]*WaitingCaptcha"
+
 // TODO sync
-var WaitingCaptchas = map[int64]*WaitingCaptcha{}
+var WaitingCaptchas CaptchaMap
 
 func ProcessLoginRsp(cli *client.QQClient, rsp *client.LoginResponse) (bool, error) {
 	if rsp.Success {
-		delete(WaitingCaptchas, cli.Uin)
+		WaitingCaptchas.Delete(cli.Uin)
 		return true, nil
 	}
 	if rsp.Error == client.SMSOrVerifyNeededError {
@@ -40,15 +42,15 @@ func ProcessLoginRsp(cli *client.QQClient, rsp *client.LoginResponse) (bool, err
 	case client.SliderNeededError:
 		log.Infof("遇到滑块验证码，根据README提示操作 https://github.com/protobufbot/Go-Mirai-Client (顺便star)")
 		prom := promise.NewPromise()
-		WaitingCaptchas[cli.Uin] = &WaitingCaptcha{
+		WaitingCaptchas.Store(cli.Uin, &WaitingCaptcha{
 			Captcha: &dto.Bot_Captcha{
 				BotId:       cli.Uin,
 				CaptchaType: dto.Bot_Captcha_SLIDER_CAPTCHA,
 				Data:        &dto.Bot_Captcha_Url{Url: rsp.VerifyUrl},
 			},
 			Prom: prom,
-		}
-		defer delete(WaitingCaptchas, cli.Uin)
+		})
+		defer WaitingCaptchas.Delete(cli.Uin)
 		result, err := prom.Get()
 		if err != nil {
 			return false, fmt.Errorf("提交ticket错误")
@@ -63,15 +65,15 @@ func ProcessLoginRsp(cli *client.QQClient, rsp *client.LoginResponse) (bool, err
 		log.Infof("遇到图形验证码，根据README提示操作 https://github.com/protobufbot/Go-Mirai-Client (顺便star)")
 		_ = ioutil.WriteFile("captcha.jpg", rsp.CaptchaImage, 0644)
 		prom := promise.NewPromise()
-		WaitingCaptchas[cli.Uin] = &WaitingCaptcha{
+		WaitingCaptchas.Store(cli.Uin, &WaitingCaptcha{
 			Captcha: &dto.Bot_Captcha{
 				BotId:       cli.Uin,
 				CaptchaType: dto.Bot_Captcha_PIC_CAPTCHA,
 				Data:        &dto.Bot_Captcha_Image{Image: rsp.CaptchaImage},
 			},
 			Prom: prom,
-		}
-		defer delete(WaitingCaptchas, cli.Uin)
+		})
+		defer WaitingCaptchas.Delete(cli.Uin)
 		result, err := prom.Get()
 		text := result.(string)
 		rsp, err := cli.SubmitCaptcha(strings.ReplaceAll(text, "\n", ""), rsp.CaptchaSign)
@@ -86,15 +88,15 @@ func ProcessLoginRsp(cli *client.QQClient, rsp *client.LoginResponse) (bool, err
 			return false, fmt.Errorf("请求短信验证码错误，可能是太频繁")
 		}
 		prom := promise.NewPromise()
-		WaitingCaptchas[cli.Uin] = &WaitingCaptcha{
+		WaitingCaptchas.Store(cli.Uin, &WaitingCaptcha{
 			Captcha: &dto.Bot_Captcha{
 				BotId:       cli.Uin,
 				CaptchaType: dto.Bot_Captcha_SMS,
 				Data:        &dto.Bot_Captcha_Url{Url: rsp.SMSPhone},
 			},
 			Prom: prom,
-		}
-		defer delete(WaitingCaptchas, cli.Uin)
+		})
+		WaitingCaptchas.Delete(cli.Uin)
 		result, err := prom.Get()
 		if err != nil {
 			return false, fmt.Errorf("提交短信验证码错误")
@@ -108,15 +110,15 @@ func ProcessLoginRsp(cli *client.QQClient, rsp *client.LoginResponse) (bool, err
 		log.Info("设置环境变量 SMS = 1 可优先使用短信验证码")
 
 		prom := promise.NewPromise()
-		WaitingCaptchas[cli.Uin] = &WaitingCaptcha{
+		WaitingCaptchas.Store(cli.Uin, &WaitingCaptcha{
 			Captcha: &dto.Bot_Captcha{
 				BotId:       cli.Uin,
 				CaptchaType: dto.Bot_Captcha_UNSAFE_DEVICE_LOGIN_VERIFY,
 				Data:        &dto.Bot_Captcha_Url{Url: rsp.VerifyUrl},
 			},
 			Prom: prom,
-		}
-		defer delete(WaitingCaptchas, cli.Uin)
+		})
+		defer WaitingCaptchas.Delete(cli.Uin)
 		_, err := prom.Get()
 		cli.Disconnect()
 		time.Sleep(3 * time.Second)