webui4frpc/internal/cluster/registry_test.go

package cluster

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
)

// writeFakeFrpc writes an executable file pretending to be frpc (any bytes
// work for tests that don't run --version).
func writeFakeFrpc(t *testing.T, dir, version string) string {
	t.Helper()
	binDir := filepath.Join(dir, "frpc-"+version)
	if err := os.MkdirAll(binDir, 0o755); err != nil {
		t.Fatal(err)
	}
	bin := filepath.Join(binDir, "frpc")
	if err := os.WriteFile(bin, []byte("#!/bin/sh\necho frpc "+version+"\n"), 0o755); err != nil {
		t.Fatal(err)
	}
	return bin
}

func TestCachedVersionsAndProvide(t *testing.T) {
	dir := t.TempDir()
	writeFakeFrpc(t, dir, "0.71.0")
	reg := NewRegistry(dir)
	vs := reg.CachedVersions()
	if len(vs) != 1 || vs[0] != "0.71.0" {
		t.Fatalf("cached versions = %+v", vs)
	}
	if !reg.HasVersion("0.71.0") || reg.HasVersion("0.99.0") {
		t.Fatalf("HasVersion wrong")
	}
	p, sum, err := reg.Provide("0.71.0")
	if err != nil {
		t.Fatal(err)
	}
	if p == "" || len(sum) != 64 {
		t.Fatalf("provide p=%s sum=%s", p, sum)
	}
}

func TestEnsureVersionPrefersLocalAndPeer(t *testing.T) {
	dir := t.TempDir()
	reg := NewRegistry(dir)

	peerBin := []byte("#!/bin/sh\necho fake-frpc\n")
	sum := sha256.Sum256(peerBin)
	peerSum := hex.EncodeToString(sum[:])

	peerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/frpc/0.71.0" {
			w.Header().Set("X-Frpc-SHA256", peerSum)
			_, _ = w.Write(peerBin)
			return
		}
		http.NotFound(w, r)
	}))
	defer peerSrv.Close()
	host := peerSrv.Listener.Addr().String()

	reg.RegisterPeer("peer1", host, "u", "p")
	if reg.HasVersion("0.71.0") {
		t.Fatal("should not be cached yet")
	}
	binPath, err := reg.EnsureVersion(context.Background(), "0.71.0")
	if err != nil {
		t.Fatalf("EnsureVersion: %v", err)
	}
	if !reg.HasVersion("0.71.0") {
		t.Fatal("EnsureVersion should have stored binary")
	}
	got, err := reg.EnsureVersion(context.Background(), "0.71.0")
	if err != nil || got != binPath {
		t.Fatalf("local hit = %q err=%v", got, err)
	}
}

func TestFetchFromPeerChecksumMismatch(t *testing.T) {
	peerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("corrupted"))
	}))
	defer peerSrv.Close()
	_, err := fetchFromPeer(context.Background(), &PeerTarget{Addr: peerSrv.Listener.Addr().String()}, "1.2.3", "0000000000000000000000000000000000000000000000000000000000000000")
	if err == nil {
		t.Fatal("expected checksum mismatch error")
	}
}

func TestRegisterAndDiscover(t *testing.T) {
	dir := t.TempDir()
	writeFakeFrpc(t, dir, "0.70.0")
	reg := NewRegistry(dir)

	peerSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/manager/cluster/nodes" {
			_, _ = w.Write([]byte(`{"nodes":[{"addr":"10.0.0.9:7500","cache":["0.71.0"],"version":"0.71.0"}]}`))
			return
		}
		http.NotFound(w, r)
	}))
	defer peerSrv.Close()

	reg.RegisterPeer("peerX", peerSrv.Listener.Addr().String(), "u", "p")
	ni, err := reg.queryNode(context.Background(), reg.PeerList()[0])
	if err != nil {
		t.Fatalf("queryNode: %v", err)
	}
	if ni.Version != "0.71.0" || len(ni.Cache) != 1 {
		t.Fatalf("node info = %+v", ni)
	}
	reg.UpdateNodeInfo(ni)
	if got := reg.NodeList(); len(got) != 1 || got[0].Addr != "10.0.0.9:7500" {
		t.Fatalf("node list = %+v", got)
	}
}