Files
mtproxy_checker/cmd/tdlib_ping/tdlib.go
T

293 lines
6.9 KiB
Go

//go:build tdlib && cgo
package main
/*
#cgo LDFLAGS: -ltdjson
#cgo CFLAGS: -I/usr/include -I/usr/local/include
#include <stdlib.h>
#include <td/telegram/td_json_client.h>
*/
import "C"
import (
"encoding/hex"
"encoding/json"
"fmt"
"os"
"strconv"
"strings"
"time"
"unsafe"
"mtproxy_checker/internal/parseurl"
"mtproxy_checker/internal/secret"
)
func main() {
if len(os.Args) < 2 {
outJSON(false, "usage: tdlib_ping <tg://proxy?...>", 2)
}
run(os.Args[1])
}
func outJSON(ok bool, msg string, code int) {
enc := json.NewEncoder(os.Stdout)
enc.SetEscapeHTML(false)
_ = enc.Encode(map[string]interface{}{
"ok": ok, "error": msg, "exit_code": code,
})
if ok {
os.Exit(0)
}
switch code {
case 1:
os.Exit(1)
case 4:
os.Exit(4)
default:
os.Exit(2)
}
}
func run(proxyURL string) {
deadlineMs, _ := strconv.Atoi(os.Getenv("MTPROXY_TDLIB_PING_MS"))
if deadlineMs <= 0 {
deadlineMs = 45000
}
deadline := time.Now().Add(time.Duration(deadlineMs) * time.Millisecond)
t, err := parseurl.ParseTGProxy(proxyURL)
if err != nil {
outJSON(false, err.Error(), 1)
}
parsed, err := secret.Parse(t.Secret)
if err != nil {
outJSON(false, err.Error(), 1)
}
secHex := mtprotoSecretHex(parsed)
apiID := int64(12345)
if v := strings.TrimSpace(os.Getenv("MTPROXY_TD_API_ID")); v != "" {
if n, e := strconv.ParseInt(v, 10, 32); e == nil {
apiID = n
}
}
apiHash := strings.TrimSpace(os.Getenv("MTPROXY_TD_API_HASH"))
if apiHash == "" {
apiHash = "0123456789abcdef0123456789abcdef"
}
useTest := strings.TrimSpace(os.Getenv("MTPROXY_TD_USE_TEST_DC")) == "1"
dbBase := strings.TrimSpace(os.Getenv("MTPROXY_TDLIB_DATABASE_DIR"))
var cleanupTemp string
if dbBase == "" {
cleanupTemp, err = os.MkdirTemp("", "mtproxy_tdlib_")
if err != nil {
outJSON(false, err.Error(), 2)
}
dbBase = cleanupTemp
defer func() { _ = os.RemoveAll(cleanupTemp) }()
}
verb := `{"@type":"setLogVerbosityLevel","new_verbosity_level":1}`
if cs := C.CString(verb); cs != nil {
_ = C.td_execute(cs)
C.free(unsafe.Pointer(cs))
}
clientID := int(C.td_create_client_id())
defer closeClient(clientID)
if !waitAuthState(clientID, "authorizationStateWaitTdlibParameters", deadline) {
outJSON(false, "timeout waiting for authorizationStateWaitTdlibParameters", 4)
}
params := map[string]interface{}{
"@type": "setTdlibParameters",
"use_test_dc": useTest,
"database_directory": dbBase + "/db",
"files_directory": dbBase + "/files",
"use_file_database": true,
"use_chat_info_database": true,
"use_message_database": true,
"use_secret_chats": true,
"api_id": apiID,
"api_hash": apiHash,
"system_language_code": "en",
"device_model": "mtproxy-tdlib-ping",
"application_version": "1.0",
}
if err := os.MkdirAll(dbBase+"/db", 0o700); err != nil {
outJSON(false, err.Error(), 2)
}
if err := os.MkdirAll(dbBase+"/files", 0o700); err != nil {
outJSON(false, err.Error(), 2)
}
tdSend(clientID, params)
for time.Now().Before(deadline) {
ev := tdReceive(1.0)
if ev == nil {
continue
}
if typ, _ := ev["@type"].(string); typ == "error" {
outJSON(false, fmt.Sprintf("tdlib: %v", ev["message"]), 2)
}
switch authStateType(ev) {
case "authorizationStateWaitEncryptionKey":
tdSend(clientID, map[string]interface{}{
"@type": "checkDatabaseEncryptionKey", "encryption_key": "",
})
case "authorizationStateReady":
goto doProxy
case "authorizationStateWaitPhoneNumber":
outJSON(false, "TDLib authorizationStateWaitPhoneNumber: use a persistent MTPROXY_TDLIB_DATABASE_DIR after TDLib login, or set MTPROXY_TD_USE_TEST_DC=1", 2)
case "authorizationStateClosed":
outJSON(false, "unexpected authorizationStateClosed before ping", 2)
}
}
outJSON(false, "timeout in TDLib authorization", 4)
doProxy:
addBody := map[string]interface{}{
"@type": "addProxy",
"server": t.Host, "port": t.Port, "enable": true,
"type": map[string]interface{}{
"@type": "proxyTypeMtproto", "secret": secHex,
},
"@extra": "addproxy",
}
tdSend(clientID, addBody)
var proxyID int64
for time.Now().Before(deadline) {
ev := tdReceive(1.0)
if ev == nil {
continue
}
if matchExtra(ev, "addproxy") {
if typ, _ := ev["@type"].(string); typ == "error" {
outJSON(false, fmt.Sprintf("addProxy: %v", ev["message"]), 2)
}
if typ, _ := ev["@type"].(string); typ == "proxy" {
if idf, ok := ev["id"].(float64); ok {
proxyID = int64(idf)
break
}
}
outJSON(false, fmt.Sprintf("addProxy unexpected: %#v", ev), 2)
}
}
if proxyID == 0 {
outJSON(false, "timeout waiting for addProxy response", 4)
}
tdSend(clientID, map[string]interface{}{
"@type": "pingProxy", "proxy_id": proxyID, "@extra": "pingpx",
})
for time.Now().Before(deadline) {
ev := tdReceive(1.0)
if ev == nil {
continue
}
if matchExtra(ev, "pingpx") {
if typ, _ := ev["@type"].(string); typ == "ok" {
outJSON(true, "", 0)
}
if typ, _ := ev["@type"].(string); typ == "error" {
outJSON(false, fmt.Sprintf("pingProxy: %v", ev["message"]), 2)
}
outJSON(false, fmt.Sprintf("pingProxy unexpected: %#v", ev), 2)
}
}
outJSON(false, "timeout waiting for pingProxy", 4)
}
func mtprotoSecretHex(p *secret.Parsed) string {
var full []byte
switch p.Kind {
case secret.KindDD:
full = append([]byte{0xdd}, p.Key...)
case secret.KindEE:
full = append(append([]byte{0xee}, p.Key...), p.Domain...)
default:
return ""
}
return hex.EncodeToString(full)
}
func authStateType(m map[string]interface{}) string {
if m == nil {
return ""
}
if m["@type"] != "updateAuthorizationState" {
return ""
}
v, _ := m["authorization_state"].(map[string]interface{})
if v == nil {
return ""
}
t, _ := v["@type"].(string)
return t
}
func matchExtra(m map[string]interface{}, want string) bool {
if m == nil {
return false
}
ext, ok := m["@extra"].(string)
return ok && ext == want
}
func waitAuthState(clientID int, want string, deadline time.Time) bool {
for time.Now().Before(deadline) {
ev := tdReceive(1.0)
if ev == nil {
continue
}
if authStateType(ev) == want {
return true
}
}
return false
}
func tdSend(clientID int, v interface{}) {
b, err := json.Marshal(v)
if err != nil {
outJSON(false, err.Error(), 2)
}
cs := C.CString(string(b))
defer C.free(unsafe.Pointer(cs))
C.td_send(C.int(clientID), cs)
}
func tdReceive(maxSec float64) map[string]interface{} {
cs := C.td_receive(C.double(maxSec))
if cs == nil {
return nil
}
s := C.GoString(cs)
var m map[string]interface{}
if err := json.Unmarshal([]byte(s), &m); err != nil {
return nil
}
return m
}
func closeClient(clientID int) {
tdSend(clientID, map[string]interface{}{"@type": "close"})
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
ev := tdReceive(0.5)
if ev == nil {
continue
}
if authStateType(ev) == "authorizationStateClosed" {
break
}
}
}