flvx/go-gost/x/socket/websocket_reporter_test.go

package socket

import (
	"bytes"
	"errors"
	"io"
	"net/http"
	"os"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func captureStdout(t *testing.T, fn func()) string {
	t.Helper()

	orig := os.Stdout
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatalf("create stdout pipe: %v", err)
	}
	os.Stdout = w
	defer func() {
		os.Stdout = orig
		_ = w.Close()
		_ = r.Close()
	}()

	fn()

	_ = w.Close()

	var buf bytes.Buffer
	if _, err := io.Copy(&buf, r); err != nil {
		t.Fatalf("read stdout: %v", err)
	}
	return buf.String()
}

func TestBuildWebSocketCandidatesSecureFirst(t *testing.T) {
	candidates := buildWebSocketCandidates("panel.example.com:443", "abc", "2.0.2", 1, 0, 1, "")

	if len(candidates) != 2 {
		t.Fatalf("expected 2 candidates, got %d", len(candidates))
	}
	if !strings.HasPrefix(candidates[0], "wss://") {
		t.Fatalf("expected first candidate to start with wss://, got %s", candidates[0])
	}
	if !strings.HasPrefix(candidates[1], "ws://") {
		t.Fatalf("expected second candidate to start with ws://, got %s", candidates[1])
	}
}

func TestBuildWebSocketCandidatesUsesPreferredScheme(t *testing.T) {
	candidates := buildWebSocketCandidates("panel.example.com:443", "abc", "2.0.2", 1, 0, 1, "ws")

	if len(candidates) != 2 {
		t.Fatalf("expected 2 candidates, got %d", len(candidates))
	}
	if !strings.HasPrefix(candidates[0], "ws://") {
		t.Fatalf("expected preferred ws:// candidate first, got %s", candidates[0])
	}
	if !strings.HasPrefix(candidates[1], "wss://") {
		t.Fatalf("expected fallback wss:// candidate second, got %s", candidates[1])
	}
}

func TestBuildWebSocketCandidatesNormalizesSchemePrefixedAddr(t *testing.T) {
	candidates := buildWebSocketCandidates("https://panel.example.com:443/path?q=1", "abc", "2.0.2", 0, 0, 0, "")

	if len(candidates) != 2 {
		t.Fatalf("expected 2 candidates, got %d", len(candidates))
	}
	if !strings.HasPrefix(candidates[0], "wss://panel.example.com:443/") {
		t.Fatalf("expected normalized wss candidate, got %s", candidates[0])
	}
	if !strings.HasPrefix(candidates[1], "ws://panel.example.com:443/") {
		t.Fatalf("expected normalized ws fallback candidate, got %s", candidates[1])
	}
}

func TestDialWebSocketWithFallbackTriesWSAfterWSSFailure(t *testing.T) {
	orig := wsDial
	defer func() { wsDial = orig }()

	var attempts []string
	wsDial = func(_ *websocket.Dialer, rawURL string) (*websocket.Conn, *http.Response, error) {
		attempts = append(attempts, rawURL)
		if strings.HasPrefix(rawURL, "wss://") {
			return nil, nil, errors.New("tls failed")
		}
		return &websocket.Conn{}, nil, nil
	}

	_, usedURL, err := dialWebSocketWithFallback(
		&websocket.Dialer{},
		[]string{
			"wss://panel.example.com/system-info?type=1&secret=abc",
			"ws://panel.example.com/system-info?type=1&secret=abc",
		},
	)
	if err != nil {
		t.Fatalf("expected fallback success, got err=%v", err)
	}
	if !strings.HasPrefix(usedURL, "ws://") {
		t.Fatalf("expected fallback ws:// url, got %s", usedURL)
	}
	if len(attempts) != 2 {
		t.Fatalf("expected 2 attempts, got %d", len(attempts))
	}
	if !strings.HasPrefix(attempts[0], "wss://") || !strings.HasPrefix(attempts[1], "ws://") {
		t.Fatalf("unexpected attempt order: %#v", attempts)
	}
}

func TestDetectWebSocketScheme(t *testing.T) {
	if detectWebSocketScheme("wss://panel.example.com/system-info") != "wss" {
		t.Fatalf("expected wss detection")
	}
	if detectWebSocketScheme("ws://panel.example.com/system-info") != "ws" {
		t.Fatalf("expected ws detection")
	}
	if detectWebSocketScheme("http://panel.example.com/system-info") != "" {
		t.Fatalf("expected empty detection for non-websocket scheme")
	}
}

func TestSanitizeWebSocketURL(t *testing.T) {
	raw := "wss://panel.example.com/system-info?type=1&secret=abc&version=2.0.2"
	sanitized := sanitizeWebSocketURL(raw)

	if strings.Contains(sanitized, "secret=abc") {
		t.Fatalf("expected secret to be masked, got %s", sanitized)
	}
	if !strings.Contains(sanitized, "secret=%2A%2A%2A") {
		t.Fatalf("expected masked secret in url, got %s", sanitized)
	}
}

func TestNewWebSocketReporterUsesReducedMetricInterval(t *testing.T) {
	reporter := NewWebSocketReporter("panel.example.com:443", "abc")

	if reporter.pingInterval != defaultMetricReportInterval {
		t.Fatalf("expected metric interval %s, got %s", defaultMetricReportInterval, reporter.pingInterval)
	}
}

func TestFormatWebSocketDialErrorIncludesHTTPStatus(t *testing.T) {
	err := errors.New("websocket: bad handshake")
	resp := &http.Response{
		Status: "403 Forbidden",
		Body:   io.NopCloser(strings.NewReader("forbidden")),
	}

	msg := formatWebSocketDialError(err, resp)
	if !strings.Contains(msg, "HTTP 403 Forbidden") {
		t.Fatalf("expected status in message, got %s", msg)
	}
	if !strings.Contains(msg, "forbidden") {
		t.Fatalf("expected response body in message, got %s", msg)
	}
}

func TestAgentUpgradeRestartScriptStopsLegacyGostService(t *testing.T) {
	script := buildAgentRestartScript("/tmp/flux_agent.new", "/etc/flux_agent/flux_agent")

	if !strings.Contains(script, "systemctl stop flux_agent") {
		t.Fatalf("expected script to stop flux_agent, got %s", script)
	}
	if !strings.Contains(script, "mv /tmp/flux_agent.new /etc/flux_agent/flux_agent") {
		t.Fatalf("expected script to replace the flux_agent binary, got %s", script)
	}
	if !strings.Contains(script, "systemctl stop gost") {
		t.Fatalf("expected script to stop the legacy gost service, got %s", script)
	}
	if !strings.Contains(script, "systemctl disable gost") {
		t.Fatalf("expected script to disable the legacy gost service, got %s", script)
	}
	if !strings.Contains(script, "rm -f /usr/local/bin/gost") {
		t.Fatalf("expected script to remove the legacy gost binary, got %s", script)
	}
	if !strings.Contains(script, "WorkingDirectory=/etc/gost") {
		t.Fatalf("expected script to scope cleanup to the legacy FLVX gost service definition, got %s", script)
	}
	if !strings.Contains(script, "systemctl start flux_agent") {
		t.Fatalf("expected script to restart flux_agent, got %s", script)
	}
	if strings.Contains(script, "systemctl stop flux_agent && systemctl stop gost 2>/dev/null || true") {
		t.Fatalf("expected legacy gost cleanup fallback to be scoped, got %s", script)
	}
	if runtime.GOARCH == "" {
		t.Fatalf("unexpected empty runtime arch")
	}
}

func TestStartWebSocketReporterWithConfigPreservesProtocolDefaultsWithoutConfigFile(t *testing.T) {
	origDial := wsDial
	defer func() { wsDial = origDial }()

	origWD, err := os.Getwd()
	if err != nil {
		t.Fatalf("get working directory: %v", err)
	}
	t.Cleanup(func() {
		_ = os.Chdir(origWD)
	})
	if err := os.Chdir(t.TempDir()); err != nil {
		t.Fatalf("change working directory: %v", err)
	}

	urls := make(chan string, 1)
	wsDial = func(_ *websocket.Dialer, rawURL string) (*websocket.Conn, *http.Response, error) {
		select {
		case urls <- rawURL:
		default:
		}
		return nil, nil, errors.New("dial failed")
	}

	reporter := StartWebSocketReporterWithConfig("panel.example.com:443", "abc", 1, 0, 1, "2.0.2")
	defer reporter.Stop()

	select {
	case rawURL := <-urls:
		if !strings.Contains(rawURL, "http=1&tls=0&socks=1") {
			t.Fatalf("expected reconnect URL to preserve startup protocol values, got %s", rawURL)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for websocket dial")
	}
}

func TestStartWebSocketReporterWithConfigLogsSanitizedURL(t *testing.T) {
	origDial := wsDial
	defer func() { wsDial = origDial }()

	ready := make(chan struct{}, 1)
	wsDial = func(_ *websocket.Dialer, rawURL string) (*websocket.Conn, *http.Response, error) {
		select {
		case ready <- struct{}{}:
		default:
		}
		return nil, nil, errors.New("dial failed")
	}

	output := captureStdout(t, func() {
		reporter := StartWebSocketReporterWithConfig("panel.example.com:443", "abc123", 1, 0, 1, "2.0.2")
		select {
		case <-ready:
		case <-time.After(2 * time.Second):
			t.Fatal("timed out waiting for websocket dial")
		}
		reporter.Stop()
	})

	if strings.Contains(output, "secret=abc123") {
		t.Fatalf("expected logged websocket URL to mask the node secret, got %s", output)
	}
	if !strings.Contains(output, "secret=%2A%2A%2A") {
		t.Fatalf("expected logged websocket URL to include masked secret, got %s", output)
	}
}