llmgw/internal/control/adminapi/guard.go

package adminapi

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"slices"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
)

// GuardStore is where policy attachments are kept.
type GuardStore interface {
	GuardAttachments(ctx context.Context) ([]guard.Attachment, error)
	AttachGuard(ctx context.Context, a guard.Attachment, by string) error
	DetachGuard(ctx context.Context, a guard.Attachment) (bool, error)
}

// GuardAdmin is what the guard endpoints need: the library from server.yaml,
// the store, and a way to have the gate read attachments again at once.
type GuardAdmin struct {
	Library    guard.Config
	Store      GuardStore
	Invalidate func()
}

// SetGuard enables the guard endpoints.
func (a *API) SetGuard(g *GuardAdmin) { a.guard = g }

// GuardView is the library and where its policies are attached.
type GuardView struct {
	Policies    []string              `json:"policies"`
	Default     []string              `json:"default"`
	Attachments []GuardAttachmentView `json:"attachments"`
}

// GuardAttachmentView is one attachment, with the name people know its
// scope by.
type GuardAttachmentView struct {
	guard.Attachment
	Name string `json:"name,omitempty"`
}

// guardTarget reads where a policy is to be attached.
func (a *API) guardTarget(ctx context.Context, target string) (guard.Attachment, error) {
	kind, what, ok := strings.Cut(strings.TrimSpace(target), ":")
	what = strings.TrimSpace(what)
	if !ok || what == "" {
		return guard.Attachment{}, fmt.Errorf("%q is not org:, team:, project:, key: or user: and what", target)
	}
	switch kind {
	case "key", "user":
		return guard.Attachment{Scope: kind, ScopeID: what}, nil
	case "org", "team", "project":
		if a.quotas == nil {
			return guard.Attachment{Scope: kind, ScopeID: what}, nil
		}
		_, id, _, err := a.resolveScope(ctx, target)
		if err != nil {
			return guard.Attachment{}, err
		}
		return guard.Attachment{Scope: kind, ScopeID: id}, nil
	}
	return guard.Attachment{}, fmt.Errorf("%q: attach to org, team, project, key or user", kind)
}

func (a *API) registerGuard(mux *http.ServeMux) {
	mux.HandleFunc("GET /api/v1/guard", func(w http.ResponseWriter, r *http.Request) {
		if a.guard == nil {
			writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard: in server.yaml")
			return
		}
		att, err := a.guard.Store.GuardAttachments(r.Context())
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		v := GuardView{Default: a.guard.Library.Default, Attachments: []GuardAttachmentView{}}
		var names map[string]string
		if a.quotas != nil {
			names, _ = a.scopeNames(r.Context())
		}
		for _, at := range att {
			v.Attachments = append(v.Attachments, GuardAttachmentView{Attachment: at, Name: names[at.Scope+":"+at.ScopeID]})
		}
		for name := range a.guard.Library.Policies {
			v.Policies = append(v.Policies, name)
		}
		slices.Sort(v.Policies)
		writeJSON(w, http.StatusOK, v)
	})
	change := func(attach bool) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			if a.guard == nil {
				writeError(w, http.StatusNotImplemented, "no guard policies: define them under guard: in server.yaml")
				return
			}
			var at struct {
				// Target is org:NAME, team:ORG/TEAM, project:ORG/TEAM/PROJECT
				// (or the ID after the colon), key:KEY-ID or user:USER-ID.
				Target string `json:"target"`
				Policy string `json:"policy"`
			}
			if err := json.NewDecoder(r.Body).Decode(&at); err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			att, err := a.guardTarget(r.Context(), at.Target)
			if err != nil {
				writeError(w, http.StatusBadRequest, err.Error())
				return
			}
			att.Policy = at.Policy
			if _, ok := a.guard.Library.Policies[att.Policy]; !ok && attach {
				writeError(w, http.StatusBadRequest, fmt.Sprintf("no policy %q in server.yaml's guard.policies", att.Policy))
				return
			}
			changed := true
			if attach {
				err = a.guard.Store.AttachGuard(r.Context(), att, actorOf(r))
			} else {
				changed, err = a.guard.Store.DetachGuard(r.Context(), att)
			}
			if err != nil {
				writeError(w, http.StatusInternalServerError, err.Error())
				return
			}
			if !changed {
				writeError(w, http.StatusNotFound, "no such attachment")
				return
			}
			if a.guard.Invalidate != nil {
				a.guard.Invalidate()
			}
			writeJSON(w, http.StatusOK, att)
		}
	}
	mux.HandleFunc("POST /api/v1/guard/attachments", change(true))
	mux.HandleFunc("DELETE /api/v1/guard/attachments", change(false))
}