webui4frpc/internal/cluster/registry.go

// Package cluster implements peer-to-peer distribution of the frpc binary and
// an in-memory node registry for the cluster. Binary exchange is preferred
// between cluster members; the external GitHub URL is only a fallback. A new
// member joining asks its neighbors to transmit the binary automatically.
package cluster

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

// PeerTarget identifies a cluster node that can serve binaries.
type PeerTarget struct {
	Addr string // host:port of the peer's frpc binary endpoint
	User string
	Pass string
}

// NodeInfo is exchanged during discovery: each node reports which frpc
// versions it has cached.
type NodeInfo struct {
	ID      string   `json:"id"`
	Addr    string   `json:"addr"`
	Version string   `json:"version,omitempty"`
	Cache   []string `json:"cache,omitempty"`
	SeenAt  int64    `json:"seenAt"`
}

// Registry tracks peers and the locally cached frpc binaries.
type Registry struct {
	BinDir string

	mu      chan struct{} // simple mutex
	peers   map[string]*PeerTarget
	nodeInf map[string]*NodeInfo
}

// NewRegistry builds a Registry rooted at binDir (where frpc-<version> dirs live).
func NewRegistry(binDir string) *Registry {
	return &Registry{
		BinDir:  binDir,
		mu:      make(chan struct{}, 1),
		peers:   map[string]*PeerTarget{},
		nodeInf: map[string]*NodeInfo{},
	}
}

func (r *Registry) lock()   { r.mu <- struct{}{} }
func (r *Registry) unlock() { <-r.mu }

// platformName mirrors install.platformName (kept local to avoid import cycle).
func platformName() string {
	arch := runtime.GOARCH
	if arch == "x86_64" {
		arch = "amd64"
	}
	return runtime.GOOS + "_" + arch
}

// CachedVersions lists frpc versions present under binDir as frpc-<version>.
func (r *Registry) CachedVersions() []string {
	entries, err := os.ReadDir(r.BinDir)
	if err != nil {
		return nil
	}
	var out []string
	for _, e := range entries {
		if !e.IsDir() || !strings.HasPrefix(e.Name(), "frpc-") {
			continue
		}
		out = append(out, strings.TrimPrefix(e.Name(), "frpc-"))
	}
	sort.Strings(out)
	return out
}

// HasVersion reports whether a version is both cached and has a runnable frpc.
func (r *Registry) HasVersion(version string) bool {
	version = strings.TrimPrefix(version, "v")
	p := filepath.Join(r.BinDir, "frpc-"+version, "frpc")
	info, err := os.Stat(p)
	return err == nil && !info.IsDir() && (info.Mode()&0o111) != 0
}

// BinaryPath returns the local file path for a cached version.
func (r *Registry) BinaryPath(version string) string {
	return filepath.Join(r.BinDir, "frpc-"+strings.TrimPrefix(version, "v"), "frpc")
}

// Provide serves a cached binary to a peer. Returns the file path and a
// SHA-256 checksum for verification.
func (r *Registry) Provide(version string) (path, checksum string, err error) {
	version = strings.TrimPrefix(version, "v")
	p := r.BinaryPath(version)
	info, err := os.Stat(p)
	if err != nil || info.IsDir() {
		return "", "", fmt.Errorf("version %s not cached locally", version)
	}
	h := sha256.New()
	f, err := os.Open(p)
	if err != nil {
		return "", "", err
	}
	_, copyErr := io.Copy(h, f)
	f.Close()
	if copyErr != nil {
		return "", "", copyErr
	}
	return p, hex.EncodeToString(h.Sum(nil)), nil
}

// RegisterPeer remembers a reachable peer so binary exchange can prefer it.
func (r *Registry) RegisterPeer(id, addr, user, pass string) {
	r.lock()
	defer r.unlock()
	r.peers[id] = &PeerTarget{Addr: addr, User: user, Pass: pass}
}

// UpdateNodeInfo stores a discovered node's cached versions.
func (r *Registry) UpdateNodeInfo(ni *NodeInfo) {
	r.lock()
	defer r.unlock()
	ni.SeenAt = time.Now().Unix()
	r.nodeInf[ni.ID] = &NodeInfo{ID: ni.ID, Addr: ni.Addr, Version: ni.Version, Cache: append([]string(nil), ni.Cache...), SeenAt: ni.SeenAt}
}

// PeerList returns the currently registered peers (deep copy).
func (r *Registry) PeerList() []*PeerTarget {
	r.lock()
	defer r.unlock()
	out := make([]*PeerTarget, 0, len(r.peers))
	for _, p := range r.peers {
		out = append(out, &PeerTarget{Addr: p.Addr, User: p.User, Pass: p.Pass})
	}
	return out
}

// NodeList returns the discovered nodes (deep copies) including cached versions.
func (r *Registry) NodeList() []NodeInfo {
	r.lock()
	defer r.unlock()
	out := make([]NodeInfo, 0, len(r.nodeInf))
	for _, n := range r.nodeInf {
		out = append(out, NodeInfo{
			ID: n.ID, Addr: n.Addr, Version: n.Version,
			Cache: append([]string(nil), n.Cache...), SeenAt: n.SeenAt,
		})
	}
	return out
}

// fetchFromPeer downloads the frpc binary for a version from one peer over
// HTTP GET /frpc/{version}. Basic Auth is used; the body is capped and its
// SHA-256 must match the peer-provided checksum (defense against corruption).
func fetchFromPeer(ctx context.Context, peer *PeerTarget, version, wantChecksum string) ([]byte, error) {
	url := fmt.Sprintf("http://%s/frpc/%s", peer.Addr, strings.TrimPrefix(version, "v"))
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	if peer.User != "" || peer.Pass != "" {
		req.SetBasicAuth(peer.User, peer.Pass)
	}
	cli := &http.Client{Timeout: 30 * time.Second}
	resp, err := cli.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil, fmt.Errorf("peer GET %s -> HTTP %d", url, resp.StatusCode)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, 256<<20))
	if err != nil {
		return nil, err
	}
	if wantChecksum != "" {
		sum := sha256.Sum256(data)
		got := hex.EncodeToString(sum[:])
		if !strings.EqualFold(got, wantChecksum) {
			return nil, fmt.Errorf("peer binary checksum mismatch: got %s want %s", got, wantChecksum)
		}
	}
	return data, nil
}

// githubURL builds the official GitHub release download URL for a version.
func githubURL(version string) string {
	ver := strings.TrimPrefix(version, "v")
	return fmt.Sprintf("https://github.com/fatedier/frp/releases/download/v%s/frp_%s_%s.tar.gz", ver, ver, platformName())
}

// fetchFromURL downloads the frpc *release tarball* from GitHub and extracts
// the frpc binary (same shape as install.Install). Returns extracted bytes.
func fetchFromURL(ctx context.Context, version string) ([]byte, error) {
	url := githubURL(version)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	cli := &http.Client{Timeout: 120 * time.Second}
	resp, err := cli.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil, fmt.Errorf("URL %s -> HTTP %d", url, resp.StatusCode)
	}
	tarGz, err := io.ReadAll(io.LimitReader(resp.Body, 256<<20))
	if err != nil {
		return nil, err
	}
	return extractFrpcBytes(tarGz)
}

// EnsureVersion makes sure a given frpc version is cached locally. It follows
// the distribution priority: local cache -> cluster peers -> external URL.
// pathSafeVersion normalizes a requested version and rejects anything that
// could escape the cache dir (e.g. "../../x").
func pathSafeVersion(version string) (string, bool) {
	v := strings.TrimPrefix(version, "v")
	if v == "" || strings.ContainsAny(v, "/\\") || strings.Contains(v, "..") {
		return "", false
	}
	for _, c := range v {
		if (c < '0' || c > '9') && (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') && c != '.' && c != '-' && c != '_' {
			return "", false
		}
	}
	return v, true
}

func (r *Registry) EnsureVersion(ctx context.Context, version string) (string, error) {
	safe, ok := pathSafeVersion(version)
	if !ok {
		return "", fmt.Errorf("invalid version %q", version)
	}
	version = safe
	if r.HasVersion(version) {
		return r.BinaryPath(version), nil
	}

	// 1) Try peers first (cluster-internal exchange is preferred).
	for _, peer := range r.PeerList() {
		data, err := fetchFromPeer(ctx, peer, version, "")
		if err != nil {
			continue
		}
		if err := r.storeBinary(version, data); err != nil {
			continue
		}
		if r.HasVersion(version) {
			return r.BinaryPath(version), nil
		}
	}

	// 2) Fall back to the external GitHub URL.
	if data, err := fetchFromURL(ctx, version); err == nil {
		if err := r.storeBinary(version, data); err == nil && r.HasVersion(version) {
			return r.BinaryPath(version), nil
		}
	}

	return "", fmt.Errorf("unable to obtain frpc %s from cache, peers, or URL", version)
}

// storeBinary persists freshly fetched frpc bytes under binDir/frpc-<version>
// and verifies the binary runs with --version.
func (r *Registry) storeBinary(version string, data []byte) error {
	if err := os.MkdirAll(r.BinDir, 0o755); err != nil {
		return err
	}
	version = strings.TrimPrefix(version, "v")
	dir := filepath.Join(r.BinDir, "frpc-"+version)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	bin := filepath.Join(dir, "frpc")
	if err := os.WriteFile(bin, data, 0o755); err != nil {
		return err
	}
	if out, err := exec.Command(bin, "--version").CombinedOutput(); err != nil {
		_ = os.RemoveAll(dir)
		return fmt.Errorf("peer-provided frpc failed --version: %v (%s)", err, bytes.TrimSpace(out))
	}
	return nil
}

// extractFrpcBytes extracts the frpc executable bytes from a GitHub release
// tarball (frp_<v>_<platform>.tar.gz). Only the top-level frpc is taken.
func extractFrpcBytes(tarGz []byte) ([]byte, error) {
	gzr, err := gzip.NewReader(bytes.NewReader(tarGz))
	if err != nil {
		return nil, err
	}
	defer gzr.Close()
	tr := tar.NewReader(gzr)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			return nil, fmt.Errorf("frpc not found in archive")
		}
		if err != nil {
			return nil, err
		}
		name := filepath.Base(hdr.Name)
		if hdr.Typeflag == tar.TypeReg && name == "frpc" {
			data, err := io.ReadAll(io.LimitReader(tr, 256<<20))
			if err != nil {
				return nil, err
			}
			return data, nil
		}
	}
}

// Discover periodically polls registered peers for their cached versions and
// refreshes the node registry (M6 stage B heartbeats). Run in a goroutine.
func (r *Registry) Discover(ctx context.Context, interval time.Duration) {
	tick := time.NewTicker(interval)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
			peers := r.PeerList()
			for _, p := range peers {
				ni, err := r.queryNode(ctx, p)
				if err != nil {
					continue // peer unreachable; keep last known info
				}
				r.UpdateNodeInfo(ni)
			}
		}
	}
}

// queryNode asks one peer for its node info (cached frpc versions).
func (r *Registry) queryNode(ctx context.Context, p *PeerTarget) (*NodeInfo, error) {
	url := fmt.Sprintf("http://%s/api/manager/cluster/nodes", p.Addr)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	if p.User != "" || p.Pass != "" {
		req.SetBasicAuth(p.User, p.Pass)
	}
	cli := &http.Client{Timeout: 10 * time.Second}
	resp, err := cli.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil, fmt.Errorf("node query HTTP %d", resp.StatusCode)
	}
	var payload struct {
		Nodes []struct {
			Addr    string   `json:"addr"`
			Cache   []string `json:"cache"`
			Version string   `json:"version"`
		} `json:"nodes"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
		return nil, err
	}
	if len(payload.Nodes) == 0 {
		return nil, fmt.Errorf("peer returned no nodes")
	}
	self := payload.Nodes[0]
	return &NodeInfo{ID: self.Addr, Addr: self.Addr, Version: self.Version, Cache: self.Cache}, nil
}

// SyncFromPeers ensures the locally cached frpc set covers the latest version
// each known peer offers (new-member auto propagation). It leaves version
// moves to EnsureVersion on demand; here we just warm the cache.
func (r *Registry) SyncFromPeers(ctx context.Context) error {
	// First refresh node info from registered peers so the cache set reflects
	// what neighbors actually offer (a fresh join may have no info yet).
	for _, p := range r.PeerList() {
		if ni, err := r.queryNode(ctx, p); err == nil {
			r.UpdateNodeInfo(ni)
		}
	}
	for _, n := range r.NodeList() {
		if n.Version == "" {
			continue
		}
		if r.HasVersion(n.Version) {
			continue
		}
		// EnsureVersion internally iterates registered peers (and falls back
		// to the external URL), so no peer selection is needed here.
		if _, err := r.EnsureVersion(ctx, n.Version); err != nil {
			continue
		}
	}
	return nil
}

// SelfInfo returns this node's own info for reporting.
func (r *Registry) SelfInfo() NodeInfo {
	return NodeInfo{Addr: "self", Version: "", Cache: r.CachedVersions()}
}

// CacheEntry describes one locally cached frpc version (for UI + prune).
type CacheEntry struct {
	Version string `json:"version"`
	Path    string `json:"path"`
	Size    int64  `json:"size"`
	ModTime int64  `json:"modTime"`
}

// CacheInfo lists all locally cached frpc versions with metadata.
func (r *Registry) CacheInfo() []CacheEntry {
	entries, err := os.ReadDir(r.BinDir)
	if err != nil {
		return nil
	}
	var out []CacheEntry
	for _, e := range entries {
		if !e.IsDir() || !strings.HasPrefix(e.Name(), "frpc-") {
			continue
		}
		ver := strings.TrimPrefix(e.Name(), "frpc-")
		bin := filepath.Join(r.BinDir, e.Name(), "frpc")
		info, err := os.Stat(bin)
		if err != nil {
			continue
		}
		out = append(out, CacheEntry{Version: ver, Path: bin, Size: info.Size(), ModTime: info.ModTime().Unix()})
	}
	return out
}

// PruneCache removes all cached versions except the newest keep (LRU-ish).
func (r *Registry) PruneCache(keep int) []string {
	entries := r.CacheInfo()
	sort.Slice(entries, func(i, j int) bool { return entries[i].ModTime > entries[j].ModTime })
	var removed []string
	for i := keep; i < len(entries); i++ {
		dir := filepath.Join(r.BinDir, "frpc-"+entries[i].Version)
		if err := os.RemoveAll(dir); err == nil {
			removed = append(removed, entries[i].Version)
		}
	}
	return removed
}