webui4frpc/internal/httpapi/handlers_forwards.go

package httpapi

import (
	"encoding/json"
	"net/http"

	"webui4frpc/internal/process"
	"webui4frpc/internal/store"
)

// forwardsReq is the {local,remote,remotePort} natural key identifying a single
// forward. The triple is stable across canvas re-saves (unlike Link.ID, which
// ReplaceLinks wholesale-replaces) and matches HasTask's keying.
type forwardsReq struct {
	Local      string `json:"local"`
	Remote     string `json:"remote"`
	RemotePort int    `json:"remotePort"`
}

// groupReq selects a group for one-click start/stop on the forwards page.
type groupReq struct {
	Group string `json:"group"`
}

// findLinkByTriple returns the link matching the (local,remote,remotePort)
// natural key, or ok=false.
func findLinkByTriple(s *store.Store, local, remote string, port int) (store.Link, bool) {
	links, err := s.ListLinks()
	if err != nil {
		return store.Link{}, false
	}
	for _, l := range links {
		if l.Local == local && l.Remote == remote && l.RemotePort == port {
			return l, true
		}
	}
	return store.Link{}, false
}

// findLinkFromTopology looks up a forward's local/remote/link data from the
// ring topology when it's not in the local SQLite store (forward owned by
// another node). Returns the store.Local, store.Remote, store.Link, and ok.
func (h *Handler) findLinkFromTopology(local, remote string, port int) (store.Local, store.Remote, store.Link, bool) {
	if h.Ring == nil {
		return store.Local{}, store.Remote{}, store.Link{}, false
	}
	snap := h.Ring.Snapshot()
	for _, t := range snap.Topology {
		if t.Local.Name == local && t.Remote.Name == remote && t.Link.RemotePort == port {
			return t.Local, t.Remote, t.Link, true
		}
	}
	return store.Local{}, store.Remote{}, store.Link{}, false
}

// startForward flips a forward to enabled and brings it up. A local-only
// forward restarts/starts its local frpc worker so the re-enabled proxy is
// rendered back in; a cluster (non-localOnly) forward is submitted to the ring
// so the lowest-load member claims and spawns it (plan §令牌环协议). SubmitTask
// is idempotent via HasTask. Falls back to ring topology when the forward
// exists only in the cluster (owned by another node, not in local store).
func (h *Handler) startForward(local, remote string, port int) error {
	ln, ok := findLinkByTriple(h.Store, local, remote, port)
	loc, lok := h.Store.GetLocal(local)
	rem, rok := h.Store.GetRemote(remote)
	if !ok || !lok || !rok {
		// Fall back to ring topology for forwards owned by other nodes.
		tLoc, tRem, tLn, tok := h.findLinkFromTopology(local, remote, port)
		if !tok {
			return store.ErrNotFound
		}
		ln, loc, rem = tLn, tLoc, tRem
	}
	_ = h.Store.SetLinkDisabled(local, remote, port, false)
	ln.Disabled = false
	if loc.LocalOnly {
		if h.Process != nil {
			key := process.WorkerKey(local, remote, port)
			if _, has := h.Process.Status(key); has {
				_ = h.Process.Restart(key)
			} else {
				_ = h.Process.Start(key)
			}
		}
		return nil
	}
	if h.Ring != nil {
		// Publish the enable into the topology so it rides the ring (a peer that
		// still holds the stale "stopped" copy learns about it on adoption).
		h.Ring.UpdateTopologyDisabled(local, remote, port, false)
	}
	// A cluster forward whose topology entry still exists needs its OWNER to
	// bring the worker back, and neither of the normal channels can do it:
	//   - SubmitTask is idempotency-guarded, and the entry still exists (a stop
	//     keeps it as the flag's carrier), so the submission is deduped away;
	//   - the claim path drops any task for a forward that already has an
	//     owner, as a defence against duplicate-claim collisions.
	// So a re-enable of an existing entry is published as a dedicated RESTART
	// task, which the owner applies unconditionally. Without it a forward the
	// user can stop but not restart — which is what marking-instead-of-removing
	// would otherwise have produced.
	if h.Ring != nil && h.Ring.HasTopologyEntry(local, remote, port) {
		h.Ring.SubmitRestart(loc, rem, ln)
		return nil
	}
	if h.Ring != nil {
		h.Ring.SubmitTask(loc, rem, ln)
	}
	return nil
}

// stopForward flips a forward to disabled and tears it down. A local-only
// forward restarts its worker so renderRemote omits the proxy (sibling forwards
// on the same remote keep running — true per-forward stop); a cluster forward
// is revoked so the owning node cancels its worker and drops it from the
// topology (plan §任务撤销). RevokeTask is idempotent. Falls back to ring
// topology when the forward exists only in the cluster (owned by another node).
func (h *Handler) stopForward(local, remote string, port int) error {
	ln, ok := findLinkByTriple(h.Store, local, remote, port)
	loc, lok := h.Store.GetLocal(local)
	rem, rok := h.Store.GetRemote(remote)
	if !ok || !lok || !rok {
		tLoc, tRem, tLn, tok := h.findLinkFromTopology(local, remote, port)
		if !tok {
			return store.ErrNotFound
		}
		ln, loc, rem = tLn, tLoc, tRem
	} else {
		_ = h.Store.SetLinkDisabled(local, remote, port, true)
	}
	// The revoke task travels to whichever node OWNS the forward, and that node
	// re-reads the disabled flag from its own store before starting a worker —
	// so the flag has to be set on every node that has a copy of this link, not
	// just the one handling this request. Propagating Disabled on the task lets
	// the owner's RevokeFn stop the worker even if its own store row is stale.
	//
	// This also fixes a latent inconsistency: `ln` was read BEFORE the
	// SetLinkDisabled(true) above, so the link published into the token still
	// carried disabled=false and got copied into the topology entry verbatim.
	ln.Disabled = true
	// Publish the stop into the topology BEFORE revoking: the revoke retires the
	// own-side worker/entry, so the flag must already exist somewhere that
	// survives it and travels the ring. This is the send half of cluster-wide
	// disabled propagation (see store.ReconcileLinkDisabled for the receive
	// half).
	if h.Ring != nil {
		h.Ring.UpdateTopologyDisabled(local, remote, port, true)
	}
	if loc.LocalOnly {
		if h.Process != nil {
			key := process.WorkerKey(local, remote, port)
			if _, has := h.Process.Status(key); has {
				_ = h.Process.Restart(key)
			}
		}
		return nil
	}
	if h.Ring != nil {
		h.Ring.RevokeTask(loc, rem, ln)
	}
	return nil
}

// handleForwardsStart toggles one forward on.
func (h *Handler) handleForwardsStart(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w)
		return
	}
	var req forwardsReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "parse json: "+err.Error(), http.StatusBadRequest)
		return
	}
	if req.Local == "" || req.Remote == "" || req.RemotePort <= 0 {
		http.Error(w, "local/remote/remotePort required", http.StatusBadRequest)
		return
	}
	if err := h.startForward(req.Local, req.Remote, req.RemotePort); err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

// handleForwardsStop toggles one forward off.
func (h *Handler) handleForwardsStop(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w)
		return
	}
	var req forwardsReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "parse json: "+err.Error(), http.StatusBadRequest)
		return
	}
	if req.Local == "" || req.Remote == "" || req.RemotePort <= 0 {
		http.Error(w, "local/remote/remotePort required", http.StatusBadRequest)
		return
	}
	if err := h.stopForward(req.Local, req.Remote, req.RemotePort); err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

// handleForwardsGroupStart starts every forward in a group with one click.
func (h *Handler) handleForwardsGroupStart(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w)
		return
	}
	var req groupReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "parse json: "+err.Error(), http.StatusBadRequest)
		return
	}
	links, err := h.Store.ListLinks()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	for _, ln := range links {
		if ln.Group != req.Group {
			continue
		}
		_ = h.startForward(ln.Local, ln.Remote, ln.RemotePort)
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

// handleForwardsGroupStop stops every forward in a group with one click.
func (h *Handler) handleForwardsGroupStop(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w)
		return
	}
	var req groupReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "parse json: "+err.Error(), http.StatusBadRequest)
		return
	}
	links, err := h.Store.ListLinks()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	for _, ln := range links {
		if ln.Group != req.Group {
			continue
		}
		_ = h.stopForward(ln.Local, ln.Remote, ln.RemotePort)
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

// assignReq carries the group label to assign to a single forward (status
// page group chip edit). Empty Group clears the assignment (移出分组).
type assignReq struct {
	Local      string `json:"local"`
	Remote     string `json:"remote"`
	RemotePort int    `json:"remotePort"`
	Group      string `json:"group"`
}

// handleForwardsAssign changes the group label of a single forward. The
// status page group chip is the quick entry; the canvas port editor also
// carries a group field (both persist via the same store column). The local
// DB row is updated if present; the ring topology entry is always updated so
// the group label propagates to all nodes via the next token cycle.
func (h *Handler) handleForwardsAssign(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w)
		return
	}
	var req assignReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "parse json: "+err.Error(), http.StatusBadRequest)
		return
	}
	if req.Local == "" || req.Remote == "" || req.RemotePort <= 0 {
		http.Error(w, "local/remote/remotePort required", http.StatusBadRequest)
		return
	}
	found := false
	if _, ok := findLinkByTriple(h.Store, req.Local, req.Remote, req.RemotePort); ok {
		found = true
		if err := h.Store.SetLinkGroup(req.Local, req.Remote, req.RemotePort, req.Group); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	if h.Ring != nil {
		if h.Ring.UpdateTopologyGroup(req.Local, req.Remote, req.RemotePort, req.Group) {
			found = true
		}
	}
	if !found {
		http.Error(w, store.ErrNotFound.Error(), http.StatusNotFound)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

// handleForwardsGroupDelete dissolves a group: every forward in the named
// group is moved to 未分组 (grp=""). The group itself is not a stored entity
// — it exists only as a label on links — so clearing all members is the
// complete "delete". Clears both local store links and ring topology entries
// so the change syncs to all nodes.
func (h *Handler) handleForwardsGroupDelete(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		methodNotAllowed(w)
		return
	}
	var req groupReq
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "parse json: "+err.Error(), http.StatusBadRequest)
		return
	}
	if req.Group == "" {
		http.Error(w, "group required", http.StatusBadRequest)
		return
	}
	links, err := h.Store.ListLinks()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	for _, ln := range links {
		if ln.Group != req.Group {
			continue
		}
		_ = h.Store.SetLinkGroup(ln.Local, ln.Remote, ln.RemotePort, "")
	}
	// Also clear group on ring topology entries (forwards owned by other nodes).
	if h.Ring != nil {
		snap := h.Ring.Snapshot()
		for _, t := range snap.Topology {
			if t.Link.Group == req.Group {
				h.Ring.UpdateTopologyGroup(t.Local.Name, t.Remote.Name, t.Link.RemotePort, "")
			}
		}
	}
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}