From 659df606f6d5a67f508268829a3bee5b7490474b Mon Sep 17 00:00:00 2001 From: 0xCA Date: Fri, 17 Nov 2023 17:11:39 +0500 Subject: [PATCH] Config requests, TgUseridToClientID cache, fixes --- handler/routes.go | 2 +- main.go | 13 ++++-- store/jsondb/jsondb.go | 21 ++++++++++ telegram/bot.go | 89 ++++++++++++++++++++++++++++++++++++------ util/cache.go | 4 ++ util/util.go | 62 +++++++++++++++++++++++++++++ 6 files changed, 176 insertions(+), 15 deletions(-) diff --git a/handler/routes.go b/handler/routes.go index 779e0b7..1914cc5 100644 --- a/handler/routes.go +++ b/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()}) diff --git a/main.go b/main.go index 49a8d67..6f2133e 100644 --- a/main.go +++ b/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 } diff --git a/store/jsondb/jsondb.go b/store/jsondb/jsondb.go index 757ccdc..ff8d1bb 100644 --- a/store/jsondb/jsondb.go +++ b/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) } diff --git a/telegram/bot.go b/telegram/bot.go index d0c528a..459e7b0 100644 --- a/telegram/bot.go +++ b/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 } diff --git a/util/cache.go b/util/cache.go index 340f73d..8037c30 100644 --- a/util/cache.go +++ b/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 diff --git a/util/util.go b/util/util.go index 26a7c5e..9f7cb88 100644 --- a/util/util.go +++ b/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 +}