diff --git a/main.go b/main.go index d6bce7b..37106f1 100644 --- a/main.go +++ b/main.go @@ -12,6 +12,7 @@ import ( "net/url" "os" "os/signal" + "regexp" "strconv" "strings" "syscall" @@ -385,36 +386,67 @@ func hasAnyLinks(links userLinks) bool { return len(links.Classic) > 0 || len(links.Secure) > 0 || len(links.TLS) > 0 } +// serverParamRe matches server=... in tg://proxy URLs +var serverParamRe = regexp.MustCompile(`server=[^&]+`) + +func (b *bot) rewriteLinksWithHost(links userLinks) userLinks { + host := strings.TrimSpace(b.cfg.TelemtLinkHost) + if host == "" { + return links + } + replacement := "server=" + host + rewrite := func(items []string) []string { + if len(items) == 0 { + return items + } + out := make([]string, len(items)) + for i, s := range items { + out[i] = serverParamRe.ReplaceAllString(s, replacement) + } + return out + } + return userLinks{ + Classic: rewrite(links.Classic), + Secure: rewrite(links.Secure), + TLS: rewrite(links.TLS), + } +} + func (b *bot) resolveUserLinks(ctx context.Context, username, secret string, initial userLinks) userLinks { + var result userLinks if hasAnyLinks(initial) { - return initial - } - - var fetched userInfo - path := "/v1/users/" + url.PathEscape(username) - if err := b.callTelemt(ctx, path, &fetched); err == nil && hasAnyLinks(fetched.Links) { - return fetched.Links - } - - var users []userInfo - if err := b.callTelemt(ctx, "/v1/users", &users); err == nil { - for i := range users { - if users[i].Username == username && hasAnyLinks(users[i].Links) { - return users[i].Links + result = initial + } else { + var fetched userInfo + path := "/v1/users/" + url.PathEscape(username) + if err := b.callTelemt(ctx, path, &fetched); err == nil && hasAnyLinks(fetched.Links) { + result = fetched.Links + } else { + var users []userInfo + if err := b.callTelemt(ctx, "/v1/users", &users); err == nil { + for i := range users { + if users[i].Username == username && hasAnyLinks(users[i].Links) { + result = users[i].Links + break + } + } + } + if !hasAnyLinks(result) { + if err := b.callTelemt(ctx, "/v1/stats/users", &users); err == nil { + for i := range users { + if users[i].Username == username && hasAnyLinks(users[i].Links) { + result = users[i].Links + break + } + } + } + } + if !hasAnyLinks(result) { + result = b.generateLinksFromSecret(secret) } } } - - users = nil - if err := b.callTelemt(ctx, "/v1/stats/users", &users); err == nil { - for i := range users { - if users[i].Username == username && hasAnyLinks(users[i].Links) { - return users[i].Links - } - } - } - - return b.generateLinksFromSecret(secret) + return b.rewriteLinksWithHost(result) } func (b *bot) generateLinksFromSecret(secret string) userLinks {