-

internal/server/model_access.go

@@ -0,0 +1,179 @@
+package server
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"os"
+	"sort"
+	"strings"
+
+	"github.com/example/ollama-fair-gateway/internal/auth"
+	"github.com/example/ollama-fair-gateway/internal/config"
+	"github.com/example/ollama-fair-gateway/internal/proxy"
+)
+
+func cloneModelAccess(in config.ModelAccessConfig) config.ModelAccessConfig {
+	cloneRule := func(r config.ModelAccessRule) config.ModelAccessRule {
+		r.AllowedModels = append([]string(nil), r.AllowedModels...)
+		r.DeniedModels = append([]string(nil), r.DeniedModels...)
+		return r
+	}
+	out := config.ModelAccessConfig{Default: cloneRule(in.Default), Tenants: make(map[string]config.ModelAccessRule, len(in.Tenants))}
+	for name, r := range in.Tenants {
+		out.Tenants[name] = cloneRule(r)
+	}
+	return out
+}
+
+func (s *Server) modelAccessSnapshot() config.ModelAccessConfig {
+	if s == nil {
+		return config.ModelAccessConfig{Tenants: map[string]config.ModelAccessRule{}}
+	}
+	if v := s.modelAccess.Load(); v != nil {
+		return cloneModelAccess(v.(config.ModelAccessConfig))
+	}
+	return cloneModelAccess(s.cfg.ModelAccess)
+}
+
+func (s *Server) runtimeTenantModelAccessRule(tenant string) config.ModelAccessRule {
+	m := s.modelAccessSnapshot()
+	if r, ok := m.Tenants[tenant]; ok {
+		if r.Mode == "" {
+			r.Mode = "allow_all"
+		}
+		return r
+	}
+	r := m.Default
+	if r.Mode == "" {
+		r.Mode = "allow_all"
+	}
+	return r
+}
+
+func normalizeAccessRule(r config.ModelAccessRule) config.ModelAccessRule {
+	r.Mode = strings.TrimSpace(r.Mode)
+	if r.Mode == "" {
+		r.Mode = "allow_all"
+	}
+	r.AllowedModels = cleanStrings(r.AllowedModels)
+	r.DeniedModels = cleanStrings(r.DeniedModels)
+	return r
+}
+
+func validateModelAccessConfig(m config.ModelAccessConfig) error {
+	if err := config.ValidateModelAccessRule(normalizeAccessRule(m.Default)); err != nil {
+		return fmt.Errorf("default model access: %w", err)
+	}
+	for tenant, r := range m.Tenants {
+		if strings.TrimSpace(tenant) == "" {
+			return errors.New("tenant name is required")
+		}
+		if err := config.ValidateModelAccessRule(normalizeAccessRule(r)); err != nil {
+			return fmt.Errorf("tenant %q model access: %w", tenant, err)
+		}
+	}
+	return nil
+}
+
+func (s *Server) storeModelAccess(next config.ModelAccessConfig) error {
+	if s.configStore == nil {
+		return errors.New("persistent configuration store unavailable")
+	}
+	base := s.cfg
+	if loader, ok := s.configStore.(configOverrideLoader); ok {
+		if loaded, err := loader.LoadWithBootstrap(s.cfg); err == nil {
+			base = loaded
+		} else if !errors.Is(err, os.ErrNotExist) {
+			return err
+		}
+	}
+	b, err := json.Marshal(base)
+	if err != nil {
+		return err
+	}
+	var candidate config.Config
+	if err := json.Unmarshal(b, &candidate); err != nil {
+		return err
+	}
+	next = cloneModelAccess(next)
+	if next.Tenants == nil {
+		next.Tenants = map[string]config.ModelAccessRule{}
+	}
+	candidate.ModelAccess = next
+	if err := validateModelAccessConfig(next); err != nil {
+		return err
+	}
+	if err := s.configStore.Save(&candidate); err != nil {
+		return err
+	}
+	s.modelAccess.Store(cloneModelAccess(next))
+	return nil
+}
+
+func (s *Server) uiModelAccess(w http.ResponseWriter, r *http.Request, actor auth.Identity) {
+	if r.URL.Path == "/gateway/ui-api/model-access" {
+		if r.Method != http.MethodGet {
+			proxy.WriteJSONError(w, http.StatusMethodNotAllowed, "method_not_allowed", "GET required")
+			return
+		}
+		m := s.modelAccessSnapshot()
+		names := make([]string, 0, len(m.Tenants))
+		for n := range m.Tenants {
+			names = append(names, n)
+		}
+		sort.Strings(names)
+		writeJSON(w, http.StatusOK, map[string]any{"default": m.Default, "tenants": m.Tenants, "tenant_names": names, "runtime": true, "persistent": s.configStore != nil})
+		return
+	}
+	raw := strings.TrimPrefix(r.URL.Path, "/gateway/ui-api/model-access/")
+	tenant, err := url.PathUnescape(raw)
+	tenant = strings.TrimSpace(tenant)
+	if err != nil || tenant == "" || strings.Contains(tenant, "/") || len(tenant) > 256 {
+		proxy.WriteJSONError(w, 400, "bad_tenant", "invalid tenant")
+		return
+	}
+	if s.configStore == nil {
+		proxy.WriteJSONError(w, 503, "config_store", "persistent configuration store unavailable")
+		return
+	}
+	s.modelAccessMu.Lock()
+	defer s.modelAccessMu.Unlock()
+	next := s.modelAccessSnapshot()
+	switch r.Method {
+	case http.MethodPut:
+		var in config.ModelAccessRule
+		if err := decodeJSON(r, &in, 128<<10); err != nil {
+			proxy.WriteJSONError(w, 400, "bad_model_access", err.Error())
+			return
+		}
+		in = normalizeAccessRule(in)
+		if err := config.ValidateModelAccessRule(in); err != nil {
+			proxy.WriteJSONError(w, 400, "bad_model_access", err.Error())
+			return
+		}
+		next.Tenants[tenant] = in
+		if err := s.storeModelAccess(next); err != nil {
+			proxy.WriteJSONError(w, 400, "model_access_store", err.Error())
+			return
+		}
+		s.log.Info("tenant model access saved", "tenant", tenant, "mode", in.Mode, "admin_subject", actor.Subject)
+		writeJSON(w, 200, map[string]any{"tenant": tenant, "rule": in, "restart_required": false})
+	case http.MethodDelete:
+		if _, ok := next.Tenants[tenant]; !ok {
+			proxy.WriteJSONError(w, 404, "model_access_not_found", fmt.Sprintf("tenant model access %q not found", tenant))
+			return
+		}
+		delete(next.Tenants, tenant)
+		if err := s.storeModelAccess(next); err != nil {
+			proxy.WriteJSONError(w, 400, "model_access_store", err.Error())
+			return
+		}
+		s.log.Info("tenant model access reset", "tenant", tenant, "admin_subject", actor.Subject)
+		writeJSON(w, 200, map[string]any{"deleted": true, "tenant": tenant, "restart_required": false})
+	default:
+		proxy.WriteJSONError(w, http.StatusMethodNotAllowed, "method_not_allowed", "PUT or DELETE required")
+	}
+}