Config requests, TgUseridToClientID cache, fixes

handler/routes.go

@@ -611,7 +611,7 @@ func SendTelegramClient(db store.IStore) echo.HandlerFunc {
 			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, "userid: " + err.Error()})
 		}
 
-		err = telegram.SendConfig(userid, clientData.Client.Name, configData, qrData)
+		err = telegram.SendConfig(userid, clientData.Client.Name, configData, qrData, false)
 
 		if err != nil {
 			return c.JSON(http.StatusInternalServerError, jsonHTTPResponse{false, err.Error()})

main.go

@@ -263,7 +263,14 @@ func main() {
 	// serves other static files
 	app.GET(util.BasePath+"/static/*", echo.WrapHandler(http.StripPrefix(util.BasePath+"/static/", assetHandler)))
 
-	initTelegram(db, util.BuildClientConfig)
+	initDeps := telegram.TgBotInitDependencies{
+		DB:                      db,
+		BuildClientConfig:       util.BuildClientConfig,
+		TgUseridToClientID:      util.TgUseridToClientID,
+		TgUseridToClientIDMutex: &util.TgUseridToClientIDMutex,
+	}
+
+	initTelegram(initDeps)
 
 	if strings.HasPrefix(util.BindAddress, "unix://") {
 		// Listen on unix domain socket.
@@ -314,10 +321,10 @@ func initServerConfig(db store.IStore, tmplDir fs.FS) {
 	}
 }
 
-func initTelegram(db store.IStore, buildClientConfig telegram.BuildClientConfig) {
+func initTelegram(initDeps telegram.TgBotInitDependencies) {
 	go func() {
 		for {
-			err := telegram.Start(db, buildClientConfig)
+			err := telegram.Start(initDeps)
 			if err == nil {
 				break
 			}

store/jsondb/jsondb.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"strconv"
 	"time"
 
 	"github.com/sdomino/scribble"
@@ -161,6 +162,20 @@ func (o *JsonDB) Init() error {
 		}
 	}
 
+	// init cache
+	clients, err := o.GetClients(false)
+	if err != nil {
+		return nil
+	}
+	for _, cl := range clients {
+		client := cl.Client
+		if len(client.TgUserid) > 3 {
+			if userid, err := strconv.ParseInt(client.TgUserid, 10, 64); err == nil {
+				util.UpdateTgToClientID(userid, client.ID)
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -314,6 +329,11 @@ func (o *JsonDB) GetClientByID(clientID string, qrCodeSettings model.QRCodeSetti
 func (o *JsonDB) SaveClient(client model.Client) error {
 	clientPath := path.Join(path.Join(o.dbPath, "clients"), client.ID+".json")
 	output := o.conn.Write("clients", client.ID, client)
+	if output == nil && len(client.TgUserid) > 3 {
+		if userid, err := strconv.ParseInt(client.TgUserid, 10, 64); err == nil {
+			util.UpdateTgToClientID(userid, client.ID)
+		}
+	}
 	err := util.ManagePerms(clientPath)
 	if err != nil {
 		return err
@@ -322,6 +342,7 @@ func (o *JsonDB) SaveClient(client model.Client) error {
 }
 
 func (o *JsonDB) DeleteClient(clientID string) error {
+	util.RemoveTgToClientID(clientID)
 	return o.conn.Delete("clients", clientID)
 }
 

telegram/bot.go

@@ -2,6 +2,7 @@ package telegram
 
 import (
 	"fmt"
+	"strconv"
 	"sync"
 	"time"
 
@@ -9,10 +10,18 @@ import (
 	"github.com/labstack/gommon/log"
 	"github.com/ngoduykhanh/wireguard-ui/model"
 	"github.com/ngoduykhanh/wireguard-ui/store"
+	"github.com/skip2/go-qrcode"
 )
 
 type BuildClientConfig func(client model.Client, server model.Server, setting model.GlobalSetting) string
 
+type TgBotInitDependencies struct {
+	DB                      store.IStore
+	BuildClientConfig       BuildClientConfig
+	TgUseridToClientID      map[int64]([]string)
+	TgUseridToClientIDMutex *sync.RWMutex
+}
+
 var (
 	TelegramToken            string
 	TelegramAllowConfRequest bool
@@ -23,13 +32,21 @@ var (
 	TgBotMutex sync.RWMutex
 
 	floodWait = make(map[int64]int64, 0)
+
+	qrCodeSettings = model.QRCodeSettings{
+		Enabled:    true,
+		IncludeDNS: true,
+		IncludeMTU: true,
+	}
 )
 
-func Start(db store.IStore, buildClientConfig BuildClientConfig) (err error) {
+func Start(initDeps TgBotInitDependencies) (err error) {
+	ticker := time.NewTicker(time.Minute)
 	defer func() {
 		TgBotMutex.Lock()
 		TgBot = nil
 		TgBotMutex.Unlock()
+		ticker.Stop()
 		if r := recover(); r != nil {
 			err = fmt.Errorf("[PANIC] recovered from panic: %v", r)
 		}
@@ -56,27 +73,75 @@ func Start(db store.IStore, buildClientConfig BuildClientConfig) (err error) {
 		fmt.Printf("[Telegram] Authorized as %s\n", res.Result.Username)
 	}
 
-	if !TelegramAllowConfRequest {
-		return
-	}
-
-	ticker := time.NewTicker(time.Minute)
 	go func() {
 		for range ticker.C {
 			updateFloodWait()
 		}
 	}()
 
+	if !TelegramAllowConfRequest {
+		return
+	}
+
 	updatesChan := echotron.PollingUpdatesOptions(token, false, echotron.UpdateOptions{AllowedUpdates: []echotron.UpdateType{echotron.MessageUpdate}})
 	for update := range updatesChan {
 		if update.Message != nil {
-			floodWait[update.Message.Chat.ID] = time.Now().Unix()
+			userid := update.Message.Chat.ID
+			if _, wait := floodWait[userid]; wait {
+				bot.SendMessage(
+					fmt.Sprintf("You can only request your configs once per %d minutes", TelegramFloodWait),
+					userid,
+					&echotron.MessageOptions{
+						ReplyToMessageID: update.Message.ID,
+					})
+				continue
+			}
+			floodWait[userid] = time.Now().Unix()
+
+			initDeps.TgUseridToClientIDMutex.RLock()
+			if clids, found := initDeps.TgUseridToClientID[userid]; found && len(clids) > 0 {
+				initDeps.TgUseridToClientIDMutex.RUnlock()
+
+				for _, clid := range clids {
+					func(clid string) {
+						clientData, err := initDeps.DB.GetClientByID(clid, qrCodeSettings)
+						if err != nil {
+							return
+						}
+
+						// build config
+						server, _ := initDeps.DB.GetServer()
+						globalSettings, _ := initDeps.DB.GetGlobalSettings()
+						config := initDeps.BuildClientConfig(*clientData.Client, server, globalSettings)
+						configData := []byte(config)
+						var qrData []byte
+
+						if clientData.Client.PrivateKey != "" {
+							qrData, err = qrcode.Encode(config, qrcode.Medium, 512)
+							if err != nil {
+								return
+							}
+						}
+
+						userid, err := strconv.ParseInt(clientData.Client.TgUserid, 10, 64)
+						if err != nil {
+							return
+						}
+
+						SendConfig(userid, clientData.Client.Name, configData, qrData, true)
+					}(clid)
+					time.Sleep(2 * time.Second)
+				}
+			} else {
+				initDeps.TgUseridToClientIDMutex.RUnlock()
+			}
+
 		}
 	}
 	return err
 }
 
-func SendConfig(userid int64, clientName string, confData, qrData []byte) error {
+func SendConfig(userid int64, clientName string, confData, qrData []byte, ignoreFloodWait bool) error {
 	TgBotMutex.RLock()
 	defer TgBotMutex.RUnlock()
 
@@ -84,10 +149,14 @@ func SendConfig(userid int64, clientName string, confData, qrData []byte) error
 		return fmt.Errorf("telegram bot is not configured or not available")
 	}
 
-	if _, wait := floodWait[userid]; wait {
+	if _, wait := floodWait[userid]; wait && !ignoreFloodWait {
 		return fmt.Errorf("this client already got their config less than %d minutes ago", TelegramFloodWait)
 	}
 
+	if !ignoreFloodWait {
+		floodWait[userid] = time.Now().Unix()
+	}
+
 	qrAttachment := echotron.NewInputFileBytes("qr.png", qrData)
 	_, err := TgBot.SendPhoto(qrAttachment, userid, &echotron.PhotoOptions{Caption: clientName})
 	if err != nil {
@@ -101,8 +170,6 @@ func SendConfig(userid int64, clientName string, confData, qrData []byte) error
 		log.Error(err)
 		return fmt.Errorf("unable to send conf file")
 	}
-
-	floodWait[userid] = time.Now().Unix()
 	return nil
 }
 

util/cache.go

@@ -1,3 +1,7 @@
 package util
 
+import "sync"
+
 var IPToSubnetRange = map[string]uint16{}
+var TgUseridToClientID = map[int64]([]string){}
+var TgUseridToClientIDMutex sync.RWMutex

util/util.go

@@ -708,3 +708,65 @@ func ManagePerms(path string) error {
 	err := os.Chmod(path, 0600)
 	return err
 }
+
+func AddTgToClientID(userid int64, clientID string) {
+	TgUseridToClientIDMutex.Lock()
+	defer TgUseridToClientIDMutex.Unlock()
+
+	if _, ok := TgUseridToClientID[userid]; ok && TgUseridToClientID[userid] != nil {
+		TgUseridToClientID[userid] = append(TgUseridToClientID[userid], clientID)
+	} else {
+		TgUseridToClientID[userid] = []string{clientID}
+	}
+}
+
+func UpdateTgToClientID(userid int64, clientID string) {
+	TgUseridToClientIDMutex.Lock()
+	defer TgUseridToClientIDMutex.Unlock()
+
+	// Detach clientID from any existing userid
+	for uid, cls := range TgUseridToClientID {
+		if cls != nil {
+			filtered := filterStringSlice(cls, clientID)
+			if len(filtered) > 0 {
+				TgUseridToClientID[uid] = filtered
+			} else {
+				delete(TgUseridToClientID, uid)
+			}
+		}
+	}
+
+	// Attach it to the new one
+	if _, ok := TgUseridToClientID[userid]; ok && TgUseridToClientID[userid] != nil {
+		TgUseridToClientID[userid] = append(TgUseridToClientID[userid], clientID)
+	} else {
+		TgUseridToClientID[userid] = []string{clientID}
+	}
+}
+
+func RemoveTgToClientID(clientID string) {
+	TgUseridToClientIDMutex.Lock()
+	defer TgUseridToClientIDMutex.Unlock()
+
+	// Detach clientID from any existing userid
+	for uid, cls := range TgUseridToClientID {
+		if cls != nil {
+			filtered := filterStringSlice(cls, clientID)
+			if len(filtered) > 0 {
+				TgUseridToClientID[uid] = filtered
+			} else {
+				delete(TgUseridToClientID, uid)
+			}
+		}
+	}
+}
+
+func filterStringSlice(s []string, excludedStr string) []string {
+	filtered := s[:0]
+	for _, v := range s {
+		if v != excludedStr {
+			filtered = append(filtered, v)
+		}
+	}
+	return filtered
+}