OpenFlare/internal/apps/relay/frps/manager_test.go

// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0

package frps

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"testing"
	"time"

	service "github.com/Rain-kl/Wavelet/pkg/protocol"
)

// Helper to write control file for the dummy script
func writeControl(t *testing.T, dir string, exitCode int, delaySeconds int) {
	t.Helper()
	controlPath := filepath.Join(dir, "control.txt")
	content := fmt.Sprintf("%d %d\n", exitCode, delaySeconds)
	err := os.WriteFile(controlPath, []byte(content), 0644)
	if err != nil {
		t.Fatalf("failed to write control file: %v", err)
	}
}

// Setup a dummy executable script that reads control.txt to decide exit code and sleep duration
func setupDummyScript(t *testing.T) (string, string) {
	t.Helper()
	dir := t.TempDir()
	scriptPath := filepath.Join(dir, "dummy_frps")

	// On macOS/Linux, we write a shell script
	scriptContent := fmt.Sprintf(`#!/bin/sh
control_file="%s/control.txt"
EXIT_CODE=0
DELAY=0
if [ -f "$control_file" ]; then
    read -r EXIT_CODE DELAY < "$control_file"
fi
if [ -n "$DELAY" ] && [ "$DELAY" -gt 0 ] 2>/dev/null; then
    sleep "$DELAY"
fi
exit "${EXIT_CODE:-0}"
`, dir)

	err := os.WriteFile(scriptPath, []byte(scriptContent), 0755)
	if err != nil {
		t.Fatalf("failed to write dummy script: %v", err)
	}

	return scriptPath, dir
}

// Helper to poll for status to eliminate timing flakiness in tests
func assertStatusEventually(t *testing.T, m *Manager, expectedStatus string, timeout time.Duration) {
	t.Helper()
	if timeout < 6*time.Second {
		timeout = 6 * time.Second
	}
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		rt := m.GetRuntimeStatus()
		if rt.Status == expectedStatus {
			return
		}
		time.Sleep(50 * time.Millisecond)
	}
	rt := m.GetRuntimeStatus()
	t.Fatalf("expected status eventually %s, got %s (err: %s)", expectedStatus, rt.Status, rt.LastError)
}

func assertCommandExitedEventually(t *testing.T, cmd *exec.Cmd, timeout time.Duration) {
	t.Helper()
	if timeout < 6*time.Second {
		timeout = 6 * time.Second
	}

	done := make(chan error, 1)
	go func() {
		done <- cmd.Wait()
	}()

	select {
	case <-time.After(timeout):
		t.Fatalf("expected process pid=%d to exit within %s", cmd.Process.Pid, timeout)
	case <-done:
	}
}

func TestStartProcessSuccess(t *testing.T) {
	scriptPath, dir := setupDummyScript(t)
	writeControl(t, dir, 0, 5) // exit code 0, sleep 5s

	m := NewManager(scriptPath, dir, "agent-token")
	defer m.Stop()

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	assertStatusEventually(t, m, "healthy", 2*time.Second)

	rt := m.GetRuntimeStatus()
	if !rt.ProcessAlive {
		t.Error("expected process to be alive")
	}
}

func TestStartProcessFailureAndBackoff(t *testing.T) {
	dir := t.TempDir()
	invalidScriptPath := filepath.Join(dir, "non_existent_frps")

	m := NewManager(invalidScriptPath, dir, "agent-token")
	defer m.Stop()

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	assertStatusEventually(t, m, "unhealthy", 2*time.Second)

	rt := m.GetRuntimeStatus()
	if !strings.Contains(rt.LastError, "failed to start") {
		t.Errorf("expected error message containing 'failed to start', got %s", rt.LastError)
	}

	// Correct the path to dummy script
	scriptPath, _ := setupDummyScript(t)
	writeControl(t, filepath.Dir(scriptPath), 0, 5)

	m.mu.Lock()
	m.frpsPath = scriptPath
	m.mu.Unlock()

	// Wait for backoff retry (1s backoff)
	assertStatusEventually(t, m, "healthy", 3*time.Second)

	rt = m.GetRuntimeStatus()
	if !rt.ProcessAlive {
		t.Error("expected process to be alive now")
	}
}

func TestUnexpectedExitAndAutorestart(t *testing.T) {
	scriptPath, dir := setupDummyScript(t)
	// Start with immediate exit code 1
	writeControl(t, dir, 1, 0)

	m := NewManager(scriptPath, dir, "agent-token")
	defer m.Stop()

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	assertStatusEventually(t, m, "unhealthy", 2*time.Second)

	rt := m.GetRuntimeStatus()
	if !strings.Contains(rt.LastError, "exited with error") {
		t.Errorf("expected exit error, got %s", rt.LastError)
	}

	// Change control to be healthy (runs for 5s, exit 0)
	writeControl(t, dir, 0, 5)

	// Wait for the retry to fire (backoff was 1s)
	assertStatusEventually(t, m, "healthy", 3*time.Second)
}

func TestBackoffReset(t *testing.T) {
	scriptPath, dir := setupDummyScript(t)
	// Rapid exit to increase backoff
	writeControl(t, dir, 1, 0)

	m := NewManager(scriptPath, dir, "agent-token")
	defer m.Stop()

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	// Crashed once, backoff is 2s
	assertStatusEventually(t, m, "unhealthy", 2*time.Second)

	// Now make it run successfully for 11 seconds (exit code 0, sleep 11s)
	writeControl(t, dir, 0, 11)

	// Wait for next retry to start running
	assertStatusEventually(t, m, "healthy", 4*time.Second)

	// Wait for process to run for 10.5 seconds to trigger backoff reset
	time.Sleep(10500 * time.Millisecond)

	// Now make it crash again (exit code 1, sleep 0s)
	writeControl(t, dir, 1, 0)

	// Wait for it to finish and crash
	assertStatusEventually(t, m, "unhealthy", 3*time.Second)

	// It crashed. Since it ran for > 10s, backoff should have been reset to 1s.
	// We make it healthy again (exit code 0, sleep 5)
	writeControl(t, dir, 0, 5)

	// Wait 1.5 seconds. If backoff was reset to 1s, it should be healthy now.
	assertStatusEventually(t, m, "healthy", 2*time.Second)
}

func TestImmediateRestartOnSameConfigDeadProcess(t *testing.T) {
	scriptPath, dir := setupDummyScript(t)
	// Crashes immediately
	writeControl(t, dir, 1, 0)

	m := NewManager(scriptPath, dir, "agent-token")
	defer m.Stop()

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	// Let it crash
	assertStatusEventually(t, m, "unhealthy", 2*time.Second)

	// Make it start successfully
	writeControl(t, dir, 0, 5)

	// Send same config block to trigger immediate restart bypass of backoff sleep
	m.UpdateConfig(context.Background(), cfg)

	// Check if it started immediately
	assertStatusEventually(t, m, "healthy", 2*time.Second)
}

func TestSupervisorGenerationInterrupt(t *testing.T) {
	scriptPath, dir := setupDummyScript(t)
	writeControl(t, dir, 0, 10)

	m := NewManager(scriptPath, dir, "agent-token")
	defer m.Stop()

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	assertStatusEventually(t, m, "healthy", 2*time.Second)

	m.mu.Lock()
	gen1 := m.generation
	cmd1 := m.cmd
	m.mu.Unlock()

	if cmd1 == nil {
		t.Fatal("expected active process")
	}

	// Update configuration with new bind port to trigger new generation
	cfg2 := &service.RelayConfig{
		BindPort:         7001,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}
	m.UpdateConfig(context.Background(), cfg2)

	assertStatusEventually(t, m, "healthy", 2*time.Second)

	m.mu.Lock()
	gen2 := m.generation
	cmd2 := m.cmd
	m.mu.Unlock()

	if gen2 <= gen1 {
		t.Errorf("expected generation incremented, got gen1=%d gen2=%d", gen1, gen2)
	}
	if cmd2 == cmd1 {
		t.Error("expected old process killed and new command started")
	}

	// Verify old process is actually killed：不要对受管 Cmd 调用 Wait（旧 supervise
	// goroutine 拥有 Wait 权，并发 Wait 会与 os/exec 内部状态竞争），改为探测
	// 进程是否已被收割（Signal(0) 在 Wait 后即报错）。
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		if err := cmd1.Process.Signal(syscall.Signal(0)); err != nil {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Error("expected first process to be killed")
}

func TestUpdateConfigKillsOrphanProcessBeforeRestart(t *testing.T) {
	scriptPath, dir := setupDummyScript(t)
	writeControl(t, dir, 0, 5)

	m := NewManager(scriptPath, dir, "agent-token")
	defer m.Stop()

	orphan := exec.Command("sh", "-c", "sleep 30")
	if err := orphan.Start(); err != nil {
		t.Fatalf("failed to start orphan process: %v", err)
	}
	t.Cleanup(func() {
		if orphan.Process != nil {
			_ = orphan.Process.Kill()
		}
	})

	if err := os.WriteFile(m.pidPath, []byte(fmt.Sprintf("%d", orphan.Process.Pid)), 0o644); err != nil {
		t.Fatalf("failed to seed orphan pid file: %v", err)
	}

	cfg := &service.RelayConfig{
		BindPort:         7000,
		VhostHTTPPort:    8080,
		AuthToken:        "test-auth",
		WebServerEnabled: false,
	}

	m.UpdateConfig(context.Background(), cfg)

	assertCommandExitedEventually(t, orphan, 2*time.Second)
	assertStatusEventually(t, m, "healthy", 2*time.Second)
}