ModelRouter/internal/gateway/plugins_api.go

package gateway

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"llmsproxy/internal/lua"
)

var (
	errPluginName  = errors.New("plugin name must be non-empty and contain no path separator or dot")
	errNoPluginDir = errors.New("no plugin_dir configured; set plugin_dir in config.yaml to enable plugins")
)

// Plugin management API.
//
//	GET    /api/plugins                  list loaded plugins + their hooks/UI/errors
//	POST   /api/plugins                  upload/replace one plugin (.lua), hot-applied
//	DELETE /api/plugins/{name}           remove a plugin
//	GET    /api/plugins/{name}/state     the plugin's own published state
//	PUT    /api/plugins/{name}/state     replace that state (admin only)
//
// Why state is a first-class endpoint: request_end is the LAST stage, so a
// hook's return value has no downstream consumer inside the gateway. A plugin
// that accumulates numbers (the billing plugin does exactly this) therefore
// keeps them in its own Lua state and serves them here, which is what its UI
// component fetches. This keeps plugin data clearly separated from the
// gateway's own stats — see the accounting note in docs/plugins.md: the
// gateway's quota accounting stays authoritative, a plugin only reports.

// handlePluginUI serves the merged UI extensions the WebUI injects at boot.
//
// It is a single GET (not per-plugin) because the browser needs ALL extensions
// before it can build the sidebar: a page contributed by one plugin and an
// element contributed by another land in the same payload, and fetching them
// separately would mean the sidebar has to be rebuilt as each arrives.
//
// Served to any authenticated role: the WebUI is authenticated before it asks,
// and a plugin's own widgets need to render for the user whose key they show.
func (g *Gateway) handlePluginUI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	ps := g.core.Plugins()
	if ps == nil {
		writeJSON(w, http.StatusOK, map[string]interface{}{"ui": lua.UIExtension{}})
		return
	}
	// AllStages, not a hand-written list: an earlier version enumerated the
	// three stages literally here, so when chain_step was added it was silently
	// missing from this response — a plugin author reading the discovery
	// document would have believed the stage did not exist. Deriving it from
	// the one place that defines the order is the whole point of having it.
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ui":     ps.UI(),
		"stages": pipelineStageNames(),
	})
}

// stageNames returns the pipeline stage names in firing order, for the
// discovery payload and for tests.
func pipelineStageNames() []string {
	out := make([]string, 0, len(lua.AllStages))
	for _, s := range lua.AllStages {
		out = append(out, string(s))
	}
	return out
}

func (g *Gateway) handlePluginsAPI(w http.ResponseWriter, r *http.Request) {
	ps := g.core.Plugins()
	if ps == nil {
		writeJSON(w, http.StatusOK, map[string]interface{}{"plugins": []interface{}{}})
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/plugins")
	path = strings.Trim(path, "/")

	// /api/plugins/{name}/state
	if strings.HasSuffix(path, "/state") {
		name := strings.TrimSuffix(path, "/state")
		if name == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "plugin name required")
			return
		}
		g.handlePluginState(w, r, name)
		return
	}

	// /api/plugins/{name}/http/{path}/access — the operator's view of, and edit
	// surface for, which keys an endpoint answers.
	if i := strings.Index(path, "/http/"); i >= 0 {
		name := path[:i]
		rest := path[i+len("/http/"):]
		if strings.HasSuffix(rest, "/access") {
			epPath := strings.TrimSuffix(rest, "/access")
			g.handleEndpointAccess(w, r, name, epPath)
			return
		}
	}

	// /api/plugins/{name}/call/{fn}
	if i := strings.Index(path, "/call/"); i >= 0 {
		name := path[:i]
		fn := path[i+len("/call/"):]
		if name == "" || fn == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "call needs a plugin and a function name")
			return
		}
		g.handlePluginCall(w, r, name, fn)
		return
	}

	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}

	if path == "" {
		switch r.Method {
		case http.MethodGet:
			writeJSON(w, http.StatusOK, map[string]interface{}{
				"plugins":     ps.List(),
				"on_disk":     ps.OnDisk(),
				"hook_errors": ps.HookErrors(),
				"plugin_dir":  g.pluginDir(),
			})
		case http.MethodPost:
			var body struct {
				Name string `json:"name"`
				Code string `json:"code"`
			}
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
				return
			}
			if err := g.installPlugin(body.Name, body.Code); err != nil {
				writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "name": body.Name})
		default:
			writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
		}
		return
	}

	// /api/plugins/{name} — source read (GET), enable/disable (PUT), delete.
	// Multi-segment paths (…/state) are handled earlier and never reach here.
	if !strings.Contains(path, "/") {
		switch r.Method {
		case http.MethodGet:
			// Reading the SOURCE (not the runtime state) is what an editor
			// needs; the state endpoint is /state and returns accumulated data
			// instead. Mixing the two would make a "save what I read"
			// round-trip impossible.
			code, err := g.readPluginSource(path)
			if err != nil {
				writeError(w, http.StatusNotFound, "not_found", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{
				"name": path, "code": code,
				"enabled": g.core.Plugins().Enabled(path),
			})
		case http.MethodPut:
			// Enable / disable. The intent rides in the body rather than the
			// verb, because "enable" and "disable" are the same resource and
			// PUT /{name} is the one the docs advertise.
			var body struct {
				Enabled *bool `json:"enabled"`
			}
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
				return
			}
			if body.Enabled == nil {
				writeError(w, http.StatusBadRequest, "invalid_request", `body must be {"enabled":true|false}`)
				return
			}
			if err := g.core.Plugins().SetEnabled(path, *body.Enabled); err != nil {
				writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{
				"ok": true, "name": path, "enabled": *body.Enabled,
			})
		case http.MethodDelete:
			if err := g.removePlugin(path); err != nil {
				writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
		default:
			writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
		}
		return
	}

	writeError(w, http.StatusNotFound, "not_found", "unknown plugin sub-resource: "+path)
}

// readPluginSource returns a plugin's file contents for the editor.
func (g *Gateway) readPluginSource(name string) (string, error) {
	if err := validPluginName(name); err != nil {
		return "", err
	}
	if g.pluginDir() == "" {
		return "", errNoPluginDir
	}
	b, err := os.ReadFile(filepath.Join(g.pluginDir(), name+".lua"))
	if err != nil {
		return "", fmt.Errorf("no plugin source named %q", name)
	}
	return string(b), nil
}

// validPluginName rejects anything that could escape the plugin directory.
// Shared by install / remove / read so the check cannot drift between them.
func validPluginName(name string) error {
	if name == "" || strings.ContainsAny(name, `/\.`) {
		return errPluginName
	}
	return nil
}

// handlePluginState serves GET (read state) and PUT (replace state).
func (g *Gateway) handlePluginState(w http.ResponseWriter, r *http.Request, name string) {
	ps := g.core.Plugins()
	switch r.Method {
	case http.MethodGet:
		// The plugin's audience gates its DATA as well as its page. Hiding the UI
		// while leaving GET /state open would put the numbers one curl away, which
		// is the leak the audience field exists to close — and billing's state is
		// a global ledger, not the caller's own usage.
		//
		// A plugin that declares no audience is unchanged (every role may read
		// it), so a single-purpose widget keeps working without edits.
		role := reqRole(r.Context())
		if !ps.StateVisibleTo(name, role) {
			writeError(w, http.StatusForbidden, "forbidden",
				"plugin "+name+" is not visible to role "+role)
			return
		}
		// prices rides along: the billing rules editor must show the table
		// that is actually in effect, not a reconstruction from config. It is
		// absent for plugins with no prices (every plugin but billing).
		state := ps.State(name)
		// A plugin may SPLIT its state: global buckets stay deployment-wide while
		// per-caller buckets are narrowed to the authenticated key. Role
		// visibility alone cannot express this — the same role is allowed to read
		// the plugin, just not another caller's row.
		if scope := ps.StateScopeFor(name); scope != nil && role != "admin" {
			state = splitStateForKey(state, scope, reqKey(r.Context()))
		}
		payload := map[string]interface{}{
			"plugin": name,
			"state":  state,
		}
		if pr := ps.Prices(name); pr != nil {
			payload["prices"] = pr
		}
		writeJSON(w, http.StatusOK, payload)
	case http.MethodPut:
		if reqRole(r.Context()) != "admin" {
			writeError(w, http.StatusForbidden, "forbidden", "admin role required")
			return
		}
		var state interface{}
		if err := json.NewDecoder(r.Body).Decode(&state); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if err := ps.SetState(name, state); err != nil {
			writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

// Call invokes a named function on a plugin's Lua state and returns its result.
//
// This exists because plugin pages otherwise have NO way to reach their own
// plugin code: pluginAPI offers state read/write and a couple of callbacks, and
// that is enough for billing (whose numbers arrive on hooks during request
// forwarding) but useless for a plugin that must do work on demand. A quota
// page has to ask the vendor when the operator clicks Refresh, and no hook
// fires at that moment.
//
// It is deliberately narrow:
//   - The function must be one the plugin explicitly lists in plugin.exports.
//     Anything else is refused, so this endpoint cannot become a generic "run
//     arbitrary Lua with my credentials" oracle over a plugin's private state.
//   - The caller must be admin, and the result is JSON — no Lua value is ever
//     handed back to Go as a live closure.
//   - Errors come back as (nil, message) from the plugin, matching every other
//     plugin entry point, and surface as a 502 rather than raising.
func (g *Gateway) handlePluginCall(w http.ResponseWriter, r *http.Request, name, fn string) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
		return
	}
	var args map[string]interface{}
	if err := json.NewDecoder(r.Body).Decode(&args); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	result, err := g.core.Plugins().Call(name, fn, args)
	if err != nil {
		writeError(w, http.StatusBadGateway, "plugin_error", err.Error())
		return
	}
	writeJSON(w, http.StatusOK, map[string]interface{}{"result": result})
}

func (g *Gateway) pluginDir() string {
	if c := g.core.Config(); c != nil {
		return c.PluginDir
	}
	return ""
}

// installPlugin writes a plugin to disk and hot-loads it. A syntax error is
// returned to the caller AND the file is left on disk so the operator can fix
// it, matching how adapters behave (the file is authoritative once present).
func (g *Gateway) installPlugin(name, code string) error {
	if err := validPluginName(name); err != nil {
		return err
	}
	dir := g.pluginDir()
	if dir == "" {
		return errNoPluginDir
	}
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, name+".lua"), []byte(code), 0644); err != nil {
		return err
	}
	return g.core.Plugins().LoadSource(name, code)
}

func (g *Gateway) removePlugin(name string) error {
	if err := validPluginName(name); err != nil {
		return err
	}
	if g.pluginDir() == "" {
		return errNoPluginDir
	}
	_ = os.Remove(filepath.Join(g.pluginDir(), name+".lua"))
	return g.core.Plugins().Unload(name)
}


// handleEndpointAccess serves the operator's view and edit surface for one
// plugin endpoint's reachable keys.
//
//	GET    …/http/{path}/access   the declared rules plus any override
//	PUT    …/http/{path}/access   body {"access": "admin,@homeagent"}
//	DELETE …/http/{path}/access   clear the override
//
// The spec string is a comma list: a bare word is a role, "@name" is a gateway
// key's configured name, and the literal "public" opens it to any
// authenticated key. That is the same grammar the visibility override uses, and
// keeping one grammar means an operator learns it once.
//
// Key names are validated against the configured keys BEFORE being stored. An
// allowlist naming a key that does not exist is a silent no-op — every caller
// is refused — and the operator would have no way to tell that apart from the
// endpoint simply working.
func (g *Gateway) handleEndpointAccess(w http.ResponseWriter, r *http.Request, name, epPath string) {
	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}
	ps := g.core.Plugins()

	switch r.Method {
	case http.MethodGet:
		for _, e := range ps.EndpointAccess() {
			if e.Plugin == name && e.Path == epPath {
				writeJSON(w, http.StatusOK, map[string]interface{}{"endpoint": e})
				return
			}
		}
		writeError(w, http.StatusNotFound, "not_found", "no such endpoint")
	case http.MethodPut, http.MethodDelete:
		spec := ""
		if r.Method == http.MethodPut {
			var body struct {
				Access string `json:"access"`
			}
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
				return
			}
			spec = body.Access
		}
		if err := ps.ValidateAccessSpec(spec); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
			return
		}
		if err := ps.SetEndpointAccess(name, epPath, spec); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "access": spec})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}


// splitStateForKey rewrites a plugin state so that a non-admin caller sees only
// the buckets that belong to them.
//
// The plugin declares which buckets are global and which are per-key. A per-key
// bucket is keyed by the MASKED key ("***a1b2c3" — the same value the audit log
// uses), so the caller's own row is found by masking the presented secret and
// matching the last six characters.
//
// Keys that do not match are DROPPED, not zeroed: a zeroed row would render as
// "you spent nothing" and, worse, would leak that a row exists for someone else.
func splitStateForKey(state interface{}, scope *lua.StateScope, key string) interface{} {
	m, ok := state.(map[string]interface{})
	if !ok || scope == nil || key == "" {
		return state
	}
	masked := maskKey(key)
	out := make(map[string]interface{}, len(m))
	for k, v := range m {
		if !scope.IsGlobal(k) {
			// Per-key bucket: keep only the caller's rows, at every nesting level
			// the scope mentions (by_day_key is {day: {masked: bucket}}).
			out[k] = filterKeyBuckets(v, masked, scope.DepthOf(k))
			continue
		}
		out[k] = v
	}
	return out
}

// filterKeyBuckets narrows a keyed structure to one masked key.
//
// depth 1 is {masked: bucket}; depth 2 is {day: {masked: bucket}}. The level is
// declared rather than sniffed, because both shapes are dicts and sniffing would
// confuse an empty bucket with a nested one.
func filterKeyBuckets(v interface{}, masked string, depth int) interface{} {
	switch depth {
	case 0:
		return v // not a per-key bucket; leave alone
	case 1:
		m, ok := v.(map[string]interface{})
		if !ok {
			return v
		}
		if row, ok := m[masked]; ok {
			return map[string]interface{}{masked: row}
		}
		return map[string]interface{}{}
	default:
		m, ok := v.(map[string]interface{})
		if !ok {
			return v
		}
		out := make(map[string]interface{}, len(m))
		for day, inner := range m {
			im, ok := inner.(map[string]interface{})
			if !ok {
				continue
			}
			if row, ok := im[masked]; ok {
				out[day] = map[string]interface{}{masked: row}
			}
		}
		return out
	}
}