package mtproto import ( "encoding/json" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" "time" ) // TestMain lets the test binary re-exec itself as a stand-in for the mtg // child process: with MTG_FAKE_CHILD=1 it records its pid and blocks, so the // manager can start and stop it without a real mtg-multi binary. func TestMain(m *testing.M) { if os.Getenv("MTG_FAKE_CHILD") == "1" { if f, err := os.OpenFile(os.Getenv("MTG_FAKE_PIDFILE"), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644); err == nil { fmt.Fprintf(f, "%d\n", os.Getpid()) f.Close() } select {} } os.Exit(m.Run()) } func installFakeMtg(t *testing.T) string { t.Helper() binDir := t.TempDir() self, err := os.Executable() if err != nil { t.Fatalf("locate test binary: %v", err) } payload, err := os.ReadFile(self) if err != nil { t.Fatalf("read test binary: %v", err) } if err := os.WriteFile(filepath.Join(binDir, GetBinaryName()), payload, 0o755); err != nil { t.Fatalf("install fake mtg: %v", err) } pidFile := filepath.Join(binDir, "mtg-pids.txt") t.Setenv("XUI_BIN_FOLDER", binDir) t.Setenv("MTG_FAKE_CHILD", "1") t.Setenv("MTG_FAKE_PIDFILE", pidFile) return pidFile } func spawnCount(t *testing.T, pidFile string) int { t.Helper() data, err := os.ReadFile(pidFile) if os.IsNotExist(err) { return 0 } if err != nil { t.Fatalf("read pid file: %v", err) } return len(strings.Fields(string(data))) } func waitSpawnCount(t *testing.T, pidFile string, want int) { t.Helper() deadline := time.Now().Add(5 * time.Second) for { got := spawnCount(t, pidFile) if got == want { return } if got > want { t.Fatalf("expected %d spawn(s), got %d", want, got) } if time.Now().After(deadline) { t.Fatalf("expected %d spawn(s), still %d after timeout", want, got) } time.Sleep(20 * time.Millisecond) } } func mtgInst(id int, secrets ...SecretEntry) Instance { return Instance{Id: id, Tag: fmt.Sprintf("inbound-%d", id), Listen: "127.0.0.1", Port: 24000 + id, Secrets: secrets} } func TestEnsureActionFor(t *testing.T) { cases := []struct { name string running bool curStruct, curSecrets, newStruct, newSecrets string want ensureAction }{ {"dead process restarts", false, "s", "a", "s", "a", ensureRestart}, {"structural change restarts", true, "s1", "a", "s2", "a", ensureRestart}, {"secrets change reloads", true, "s", "a", "s", "b", ensureReload}, {"identical is a noop", true, "s", "a", "s", "a", ensureNoop}, {"dead beats a secrets-only change", false, "s", "a", "s", "b", ensureRestart}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { if got := ensureActionFor(tc.running, tc.curStruct, tc.curSecrets, tc.newStruct, tc.newSecrets); got != tc.want { t.Fatalf("ensureActionFor = %d, want %d", got, tc.want) } }) } } func TestApplySecrets(t *testing.T) { cases := []struct { name string status int want bool }{ {"ok", http.StatusOK, true}, {"not found on old binary", http.StatusNotFound, false}, {"bad request", http.StatusBadRequest, false}, {"unavailable", http.StatusServiceUnavailable, false}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { var gotMethod, gotPath string var gotBody secretsPutBody srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { gotMethod, gotPath = r.Method, r.URL.Path _ = json.NewDecoder(r.Body).Decode(&gotBody) w.WriteHeader(tc.status) })) defer srv.Close() inst := mtgInst(1, SecretEntry{Name: "alice", Secret: "ee01"}) inst.AdTag = "0123456789abcdef0123456789abcdef" if got := applySecrets(serverPort(t, srv), inst); got != tc.want { t.Fatalf("applySecrets = %v, want %v", got, tc.want) } if gotMethod != http.MethodPut || gotPath != "/secrets" { t.Fatalf("expected PUT /secrets, got %s %s", gotMethod, gotPath) } if gotBody.Secrets["alice"].Secret != "ee01" || gotBody.AdTag != "0123456789abcdef0123456789abcdef" { t.Fatalf("payload must carry the secret and ad-tag: %+v", gotBody) } }) } t.Run("refused connection", func(t *testing.T) { srv := httptest.NewServer(http.NotFoundHandler()) port := serverPort(t, srv) srv.Close() if applySecrets(port, mtgInst(1, SecretEntry{Name: "a", Secret: "ee"})) { t.Fatal("a refused connection must yield false") } }) } func TestEnsureHotReloadKeepsProcess(t *testing.T) { pidFile := installFakeMtg(t) mgr := &Manager{procs: map[int]*managed{}, swept: true} inst := mtgInst(1, SecretEntry{Name: "alice", Secret: "ee01"}) if err := mgr.Ensure(inst); err != nil { t.Fatalf("initial ensure: %v", err) } waitSpawnCount(t, pidFile, 1) orig := mgr.procs[1].proc reloaded := make(chan struct{}, 1) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodPut && r.URL.Path == "/secrets" { reloaded <- struct{}{} w.WriteHeader(http.StatusOK) return } http.NotFound(w, r) })) defer srv.Close() mgr.procs[1].apiPort = serverPort(t, srv) rekeyed := mtgInst(1, SecretEntry{Name: "alice", Secret: "ee01"}, SecretEntry{Name: "bob", Secret: "ee02"}) if err := mgr.Ensure(rekeyed); err != nil { t.Fatalf("reload ensure: %v", err) } select { case <-reloaded: case <-time.After(3 * time.Second): t.Fatal("expected a PUT /secrets request") } if got := spawnCount(t, pidFile); got != 1 { t.Fatalf("hot reload must not spawn a new process, got %d", got) } if mgr.procs[1].proc != orig { t.Fatal("hot reload must keep the same process") } if mgr.procs[1].secretsFP != rekeyed.secretsFingerprint() { t.Fatal("stored secrets fingerprint must advance after a reload") } cfg, err := os.ReadFile(configPathForID(1)) if err != nil { t.Fatalf("read config: %v", err) } if !strings.Contains(string(cfg), `"bob" = "ee02"`) { t.Fatalf("reloaded config must carry the new secret:\n%s", cfg) } if !strings.Contains(string(cfg), fmt.Sprintf("api-bind-to = \"127.0.0.1:%d\"", serverPort(t, srv))) { t.Fatalf("reload must reuse the same api port:\n%s", cfg) } mgr.StopAll() } func TestEnsureReloadFallbackRestarts(t *testing.T) { pidFile := installFakeMtg(t) mgr := &Manager{procs: map[int]*managed{}, swept: true} if err := mgr.Ensure(mtgInst(2, SecretEntry{Name: "alice", Secret: "ee01"})); err != nil { t.Fatalf("initial ensure: %v", err) } waitSpawnCount(t, pidFile, 1) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNotFound) })) defer srv.Close() mgr.procs[2].apiPort = serverPort(t, srv) if err := mgr.Ensure(mtgInst(2, SecretEntry{Name: "carol", Secret: "ee03"})); err != nil { t.Fatalf("fallback ensure: %v", err) } waitSpawnCount(t, pidFile, 2) mgr.StopAll() } func TestEnsureNoopKeepsProcess(t *testing.T) { pidFile := installFakeMtg(t) mgr := &Manager{procs: map[int]*managed{}, swept: true} inst := mtgInst(3, SecretEntry{Name: "alice", Secret: "ee01"}, SecretEntry{Name: "bob", Secret: "ee02"}) if err := mgr.Ensure(inst); err != nil { t.Fatalf("initial ensure: %v", err) } waitSpawnCount(t, pidFile, 1) if err := mgr.Ensure(inst); err != nil { t.Fatalf("repeat ensure: %v", err) } reordered := mtgInst(3, SecretEntry{Name: "bob", Secret: "ee02"}, SecretEntry{Name: "alice", Secret: "ee01"}) if err := mgr.Ensure(reordered); err != nil { t.Fatalf("reordered ensure: %v", err) } time.Sleep(300 * time.Millisecond) if got := spawnCount(t, pidFile); got != 1 { t.Fatalf("an unchanged instance must keep the one process, got %d spawns", got) } mgr.StopAll() }