ModelRouter/internal/lua/plugin_http_handler.go

package lua

// Plugin-registered HTTP endpoints.
//
// The gateway mounts what a plugin declares in plugin.http and calls back into
// the plugin's Lua state to produce each response. This is the piece billing
// never needed: it only receives hooks, so it has nothing to expose. A plugin
// that computes something an operator wants to query from outside the browser —
// a health probe, a CI check, another tool — needs a route, and hard-coding one
// per plugin would mean a gateway change for every such plugin.
//
// The request and response are both plain tables crossing the Go/Lua boundary:
//
//	request  { method, path, query{}, params{}, headers{}, body, role, key }
//	response { status, headers{}, body }  or  { status, json = <any table> }
//
// A handler returns (response, nil) or (nil, errorMessage). Errors never raise:
// an escaping Lua error inside a hook has taken the whole gateway down before
// (see hookGuardSrc), so the same discipline applies here.
//
// Two properties are enforced on the Go side rather than trusted to the plugin:
//
//   - Method and status are checked against an allowlist before the response is
//     written. A plugin returning status 0 or a redirect with an injected
//     Location must not be able to steer the gateway's own routing.
//   - Path traversal cannot escape the mount: params are substituted into a
//     fixed segment count, and the handler's own Path is rejected at
//     registration if it contains "..".

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
)

// allowedMethods is deliberately short. An arbitrary verb would need a case in
// the router for no benefit, and a plugin that wants every verb can register the
// five that exist.
var allowedMethods = map[string]bool{
	http.MethodGet: true, http.MethodPost: true, http.MethodPut: true,
	http.MethodDelete: true, http.MethodPatch: true,
}

// normalizeHandlerMethod validates and defaults a declared method.
func normalizeHandlerMethod(m string) (string, error) {
	m = strings.ToUpper(strings.TrimSpace(m))
	if m == "" {
		return http.MethodGet, nil
	}
	if !allowedMethods[m] {
		return "", fmt.Errorf("unsupported method %q (allowed: GET POST PUT DELETE PATCH)", m)
	}
	return m, nil
}

// validateHandlerPath rejects paths that could escape the mount or collide with
// the plugin state/call routes.
func validateHandlerPath(p string) error {
	p = strings.TrimSpace(p)
	if p == "" {
		return fmt.Errorf("empty path")
	}
	if strings.Contains(p, "..") {
		return fmt.Errorf("path %q contains .. ", p)
	}
	if strings.Contains(p, "//") {
		return fmt.Errorf("path %q contains an empty segment", p)
	}
	for _, seg := range strings.Split(strings.Trim(p, "/"), "/") {
		if seg == "" {
			continue
		}
		if strings.HasPrefix(seg, ":") {
			if len(seg) == 1 {
				return fmt.Errorf("path %q has a nameless param segment", p)
			}
			continue
		}
		// A literal segment must not look like a wildcard, or the route table
		// gets two patterns that match the same URL.
		if strings.ContainsAny(seg, "*?") {
			return fmt.Errorf("path %q uses a wildcard; use :name for a param", p)
		}
	}
	return nil
}

// handlerRequest is what a Lua handler sees.
type handlerRequest struct {
	Method  string
	Path    string
	Query   map[string]string
	Params  map[string]string
	Headers map[string]string
	Body    string
	Role    string
	Key     string
}

// handlerResponse is what a Lua handler returns.
type handlerResponse struct {
	Status  int
	Headers map[string]string
	Body    string
	// JSON is set when the handler returned {json = <table>}; it is encoded by
	// the kernel so a plugin does not hand-roll escaping.
	JSON interface{}
}

// writeHTTPHandlerResponse renders a handler's reply onto the gateway's writer.
//
// The status is validated even though it came from a plugin: a handler
// returning 0 or 999 would make net/http panic, and one returning 3xx with a
// Location the plugin chose could be used to bounce a caller at a URL of the
// plugin's picking. Anything outside 100..599 that is not a redirect becomes 500
// with an error body rather than a panic.
func writeHTTPHandlerResponse(w http.ResponseWriter, r *handlerResponse) {
	status := r.Status
	if status == 0 {
		status = http.StatusOK
	}
	if status < 100 || status > 599 {
		status = http.StatusInternalServerError
	}
	for k, v := range r.Headers {
		// A plugin cannot set Content-Length (the body is written by the kernel)
		// or hop-by-hop headers (Connection, Transfer-Encoding) which would
		// corrupt the gateway's own response framing.
		switch http.CanonicalHeaderKey(k) {
		case "Content-Length", "Connection", "Transfer-Encoding", "Upgrade":
			continue
		}
		w.Header().Set(k, v)
	}
	if r.JSON != nil {
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.WriteHeader(status)
		// Encoded here rather than by the gateway layer: this file is in the lua
		// package and importing the gateway would be a cycle.
		if enc, err := json.Marshal(r.JSON); err == nil {
			_, _ = w.Write(enc)
		} else {
			_, _ = w.Write([]byte(`{"error":"handler json could not be encoded"}`))
		}
		return
	}
	if w.Header().Get("Content-Type") == "" {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	}
	w.WriteHeader(status)
	if r.Body != "" {
		_, _ = w.Write([]byte(r.Body))
	}
}

// parseHandlerResponse reads a Lua handler's return value.
//
// It accepts three shapes, because a plugin should not have to think about it:
//
//	{ status = 200, body = "text" }
//	{ status = 201, json = { any = "table" } }
//	{ status = 204 }                       -- no body
//
// Anything else — a string, a number, a table without status — is a handler bug
// and becomes an error rather than an empty 200, so the plugin author sees it.
func parseHandlerResponse(v interface{}) (*handlerResponse, error) {
	m, ok := v.(map[string]interface{})
	if !ok {
		return nil, fmt.Errorf("handler must return a table, got %T", v)
	}
	out := &handlerResponse{Status: http.StatusOK}
	if s, ok := m["status"]; ok {
		switch n := s.(type) {
		case float64:
			out.Status = int(n)
		case int:
			out.Status = n
		default:
			return nil, fmt.Errorf("status must be a number, got %T", s)
		}
	}
	if b, ok := m["body"].(string); ok {
		out.Body = b
	}
	if h, ok := m["headers"].(map[string]interface{}); ok {
		out.Headers = make(map[string]string, len(h))
		for k, hv := range h {
			out.Headers[k] = asString(hv)
		}
	}
	if j, ok := m["json"]; ok && j != nil {
		out.JSON = j
	}
	// An empty table is a handler bug — it means the author forgot to say what to
	// return. A table carrying ONLY a status is not: `{status=204}` is the normal
	// way to answer with no body, and treating it as an error would make the
	// correct spelling unusable.
	if out.Body == "" && out.JSON == nil && out.Headers == nil && len(m) == 0 {
		return nil, fmt.Errorf("handler returned an empty table; return " +
			"{status=…, body=…}, {status=…, json=…}, or {status=…} for no body")
	}
	return out, nil
}

// InvokeHandler calls a plugin's HTTP handler function with a request table and
// returns the decoded response.
//
// It reuses the same guard-in-Lua contract as the hooks: the function is called
// through __llmsproxy_call_hook, so a plugin error comes back as (nil, message)
// and never as a C-level abort.
func (ps *Plugins) InvokeHandler(plugin, fnName string, req *handlerRequest) (*handlerResponse, error) {
	p := ps.find(plugin)
	if p == nil {
		return nil, fmt.Errorf("plugin %s not found", plugin)
	}
	p.mu.Lock()
	w := p.state
	p.mu.Unlock()
	if w == nil {
		return nil, fmt.Errorf("plugin %s has no state", plugin)
	}

	// The whole call happens under the plugin lock, network waits included —
	// same property as Call, and for the same reason: one Lua state.
	p.mu.Lock()
	defer p.mu.Unlock()
	L := w.L
	L.SetTop(0)
	defer L.SetTop(0)

	L.GetGlobal(pluginGlobal)
	if L.IsNil(-1) {
		return nil, fmt.Errorf("plugin table missing")
	}
	plug := L.GetTop()
	L.GetGlobal(hookGuardName)
	if L.IsNil(-1) {
		L.Pop(1)
	}
	L.GetField(plug, fnName)
	if !L.IsFunction(-1) {
		return nil, fmt.Errorf("handler %s not found in plugin %s", fnName, plugin)
	}

	pushGoValue(L, plainForLua(map[string]interface{}{
		"method":  req.Method,
		"path":    req.Path,
		"query":   stringMapToAny(req.Query),
		"params":  stringMapToAny(req.Params),
		"headers": stringMapToAny(req.Headers),
		"body":    req.Body,
		"role":    req.Role,
		"key":     req.Key,
	}))
	if err := L.Call(2, 2); err != nil {
		return nil, err
	}
	errMsg := ""
	if L.GetTop() >= 2 {
		if !L.IsNil(-1) {
			errMsg = L.ToString(-1)
		}
		L.Pop(1)
	}
	if errMsg != "" {
		return nil, fmt.Errorf("plugin %s: %s", plugin, errMsg)
	}
	if L.GetTop() < 1 || L.IsNil(-1) {
		return nil, fmt.Errorf("handler %s returned nothing", fnName)
	}
	var raw interface{}
	if err := luaToJSON(L, -1, &raw); err != nil {
		return nil, err
	}
	return parseHandlerResponse(raw)
}

func stringMapToAny(m map[string]string) map[string]interface{} {
	if len(m) == 0 {
		return map[string]interface{}{}
	}
	out := make(map[string]interface{}, len(m))
	for k, v := range m {
		out[k] = v
	}
	return out
}

// HandlerRequest is the request table handed to a plugin HTTP handler.
// Exported because the gateway package constructs it.
type HandlerRequest = handlerRequest

// HandlerResponse is what a handler returns, after decoding.
type HandlerResponse = handlerResponse

// WriteHandlerResponse renders a handler response onto the gateway's writer.
func WriteHandlerResponse(w http.ResponseWriter, r *HandlerResponse) {
	writeHTTPHandlerResponse(w, r)
}

// HTTPHandlers returns the registered endpoints across all plugins, each tagged
// with its owner so the dispatcher knows which Lua state to enter.
func (ps *Plugins) HTTPHandlers() []HTTPHandler {
	ps.mu.RLock()
	defer ps.mu.RUnlock()
	var out []HTTPHandler
	for _, p := range ps.plugins {
		if p.http == nil || p.LoadError != "" {
			continue
		}
		for _, declared := range p.http {
			// effectiveAccess applies the operator's per-endpoint override, so the
			// dispatcher enforces what the WebUI shows rather than what the .lua
			// declared.
			h := ps.effectiveAccess(p, declared)
			h.Plugin = p.Info.Name
			out = append(out, h)
		}
	}
	return out
}

// Allows is the exported form of the access rule, so the gateway dispatcher does
// not need a duplicate copy of it in its own package.
func (h HTTPHandler) Allows(role, keyName string) bool { return h.allows(role, keyName) }

// HTTPEndpointAccess is the operator's view of one endpoint's permissions,
// returned by /api/plugins so the WebUI can render what is exposed without
// parsing .lua files.
type HTTPEndpointAccess struct {
	Plugin      string   `json:"plugin"`
	Path        string   `json:"path"`
	FullPath    string   `json:"full_path"`
	Method      string   `json:"method"`
	Handler     string   `json:"handler"`
	Roles       []string `json:"roles,omitempty"`
	Keys        []string `json:"keys,omitempty"`
	Public      bool     `json:"public"`
	Description string   `json:"description,omitempty"`
	// OverrideRoles / OverrideKeys carry the operator's own restriction, which
	// REPLACES whatever the plugin declared. Empty means "use the plugin's".
	//
	// This is the mechanism the admin uses to narrow (or widen) an endpoint's
	// reachable keys without editing the plugin: a bundled .lua is overwritten on
	// the next upgrade, and a hand-installed one is not something the operator
	// should have to maintain.
	OverrideRoles []string `json:"override_roles,omitempty"`
	OverrideKeys  []string `json:"override_keys,omitempty"`
}

// endpointAccess returns the access table for every registered endpoint, with
// operator overrides folded in.
func (ps *Plugins) endpointAccess() []HTTPEndpointAccess {
	ps.mu.RLock()
	defer ps.mu.RUnlock()
	var out []HTTPEndpointAccess
	for _, p := range ps.plugins {
		if p.http == nil || p.LoadError != "" {
			continue
		}
		ov := ps.visibility[p.Info.Name]
		var ovRoles, ovKeys []string
		if ov != "" {
			for _, part := range strings.Split(ov, ",") {
				part = strings.TrimSpace(part)
				if part == "" {
					continue
				}
				// A key name is anything that is not a role: the override string
				// carries both, so "@name" marks a key and a bare word a role.
				if strings.HasPrefix(part, "@") {
					ovKeys = append(ovKeys, strings.TrimPrefix(part, "@"))
				} else {
					ovRoles = append(ovRoles, part)
				}
			}
		}
		for _, h := range p.http {
			out = append(out, HTTPEndpointAccess{
				Plugin: p.Info.Name, Path: h.Path, FullPath: h.routePath(),
				Method: h.Method, Handler: h.Handler,
				Roles: h.Roles, Keys: h.Keys, Public: h.Public,
				Description:   h.Description,
				OverrideRoles: ovRoles, OverrideKeys: ovKeys,
			})
		}
	}
	return out
}

// EndpointAccess exposes the table to the gateway package.
func (ps *Plugins) EndpointAccess() []HTTPEndpointAccess { return ps.endpointAccess() }

// SetEndpointAccess narrows or widens one endpoint's reachable keys/roles.
//
// It is an OVERRIDE, applied on top of what the plugin declared, and stored in
// the same per-plugin visibility map the UI audience uses. An empty value clears
// it. Overriding per endpoint rather than per plugin is what the operator
// actually wants: /health is fine for every key while /balance is not, and
// those two live in the same plugin.
func (ps *Plugins) SetEndpointAccess(plugin, path, spec string) error {
	p := ps.find(plugin)
	if p == nil {
		return fmt.Errorf("plugin %s not found", plugin)
	}
	found := false
	for _, h := range p.http {
		if h.Path == path {
			found = true
			break
		}
	}
	if !found {
		return fmt.Errorf("plugin %s has no endpoint %q", plugin, path)
	}
	ps.mu.Lock()
	defer ps.mu.Unlock()
	if ps.visibility == nil {
		ps.visibility = visibilityOverrides{}
	}
	key := endpointOverrideKey(plugin, path)
	if strings.TrimSpace(spec) == "" {
		delete(ps.visibility, key)
		return nil
	}
	ps.visibility[key] = spec
	return nil
}

// endpointOverrideKey namespaces a per-endpoint override so it cannot collide
// with a whole-plugin audience override of the same plugin name.
func endpointOverrideKey(plugin, path string) string {
	return plugin + "#" + path
}

// effectiveAccess applies an endpoint override to the declared handler.
func (ps *Plugins) effectiveAccess(p *Plugin, h HTTPHandler) HTTPHandler {
	ps.mu.RLock()
	ov := ps.visibility[endpointOverrideKey(p.Info.Name, h.Path)]
	ps.mu.RUnlock()
	if ov == "" {
		return h
	}
	out := h
	out.Roles, out.Keys = nil, nil
	out.Public = false
	for _, part := range strings.Split(ov, ",") {
		part = strings.TrimSpace(part)
		switch {
		case part == "":
		case part == "public":
			out.Public = true
		case strings.HasPrefix(part, "@"):
			out.Keys = append(out.Keys, strings.TrimPrefix(part, "@"))
		default:
			out.Roles = append(out.Roles, part)
		}
	}
	return out
}

// knownKeyNames is set by the gateway during wiring so ValidateAccessSpec can
// refuse an allowlist naming a key that does not exist.
//
// Without the check, an allowlist naming a typo ("homeagent" vs "home_agent")
// is stored, saved, shown in the WebUI as configured — and refuses every caller,
// including the one the operator meant. The failure looks like the endpoint being
// broken rather than the name being wrong.
var knownKeyNames func() []string

// SetKeyNameLookup registers the source of configured gateway key names.
func SetKeyNameLookup(fn func() []string) { knownKeyNames = fn }

// ValidateAccessSpec checks an access spec against the configured keys and roles.
//
// Roles are checked against the two that exist (admin, user) rather than left
// free-form: a role that does not exist produces an endpoint nobody can reach,
// which is the same silent failure as a bad key name.
func (ps *Plugins) ValidateAccessSpec(spec string) error {
	spec = strings.TrimSpace(spec)
	if spec == "" {
		return nil
	}
	known := map[string]bool{"public": true, "admin": true, "user": true}
	if knownKeyNames != nil {
		for _, n := range knownKeyNames() {
			if n != "" {
				known[strings.ToLower(n)] = true
			}
		}
	}
	for _, part := range strings.Split(spec, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		name := part
		if strings.HasPrefix(name, "@") {
			name = strings.TrimPrefix(name, "@")
		}
		if name == "" {
			return fmt.Errorf("access entry %q names no key", part)
		}
		if !known[strings.ToLower(name)] {
			return fmt.Errorf("access entry %q is neither a role (admin/user/public) "+
				"nor a configured key name", part)
		}
	}
	return nil
}