ModelRouter/internal/config/store.go

package config

import (
	"encoding/json"
	"log"
	"os"
	"sync"
)

// Store persists web-UI editable runtime state (sources added/edited) to a
// JSON file so edits survive restarts. Base YAML sources are merged underneath.
type Store struct {
	mu   sync.Mutex
	path string
	data RuntimeConfig
	box  *SecretBox
}

func NewStore(path string) *Store {
	s := &Store{path: path}
	if box, err := NewSecretBox(path); err == nil {
		s.box = box
	} else {
		log.Printf("[config] secrets disabled: %v (sensitive values will be stored in plaintext)", err)
	}
	return s
}

// SecretBox exposes the box used for persist-time encryption of source
// api_key / header values read from other places (e.g. YAML). May be nil.
func (s *Store) SecretBox() *SecretBox {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.box
}

// Load reads the runtime file (missing file = empty state).
func (s *Store) Load() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	data, err := os.ReadFile(s.path)
	if err != nil {
		if os.IsNotExist(err) {
			s.data = RuntimeConfig{}
			return nil
		}
		return err
	}
	if err := json.Unmarshal(data, &s.data); err != nil {
		return err
	}
	if s.box != nil {
		s.decryptLocked()
	}
	return nil
}

// decryptLocked replaces persisted ciphertext fields with their plaintext in
// memory (the runtime always works with plaintext; only the file is sealed).
func (s *Store) decryptLocked() {
	for i := range s.data.Sources {
		s.data.Sources[i].APIKey = s.box.MustDecrypt(s.data.Sources[i].APIKey)
		for k, v := range s.data.Sources[i].Headers {
			s.data.Sources[i].Headers[k] = s.box.MustDecrypt(v)
		}
	}
	for i := range s.data.Keys {
		s.data.Keys[i].Key = s.box.MustDecrypt(s.data.Keys[i].Key)
	}
}

// List returns the runtime sources (those edited via web UI).
func (s *Store) List() []Source {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]Source, len(s.data.Sources))
	copy(out, s.data.Sources)
	return out
}

// Upsert adds or replaces a runtime source and persists.
func (s *Store) Upsert(src Source) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := range s.data.Sources {
		if s.data.Sources[i].Name == src.Name {
			s.data.Sources[i] = src
			s.data.DeletedSources = removeString(s.data.DeletedSources, src.Name)
			return s.persistLocked()
		}
	}
	s.data.Sources = append(s.data.Sources, src)
	s.data.DeletedSources = removeString(s.data.DeletedSources, src.Name)
	return s.persistLocked()
}

// Remove deletes a runtime source from the store and persists. Base YAML
// sources are handled (truly removed from the config file) by the caller.
func (s *Store) Remove(name string) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	kept := s.data.Sources[:0]
	removed := false
	for _, src := range s.data.Sources {
		if src.Name == name {
			removed = true
			continue
		}
		kept = append(kept, src)
	}
	s.data.Sources = kept
	return removed, s.persistLocked()
}

func (s *Store) DeletedSources() map[string]bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := map[string]bool{}
	for _, name := range s.data.DeletedSources {
		out[name] = true
	}
	return out
}

func (s *Store) persistLocked() error {
	if s.box != nil {
		s.encryptLocked()
	}
	b, err := json.MarshalIndent(s.data, "", "  ")
	if err != nil {
		return err
	}
	if s.box != nil {
		s.decryptLocked()
	}
	if err := os.WriteFile(s.path, b, 0644); err != nil {
		return err
	}
	return nil
}

// encryptLocked seals sensitive fields for the write; decryptLocked runs
// right after marshaling so in-memory data stays plaintext.
func (s *Store) encryptLocked() {
	for i := range s.data.Sources {
		if v, err := s.box.Encrypt(s.data.Sources[i].APIKey); err == nil {
			s.data.Sources[i].APIKey = v
		}
		h := s.data.Sources[i].Headers
		for k, v := range h {
			if v == "" {
				continue
			}
			if e, err := s.box.Encrypt(v); err == nil {
				h[k] = e
			}
		}
	}
	for i := range s.data.Keys {
		if v, err := s.box.Encrypt(s.data.Keys[i].Key); err == nil {
			s.data.Keys[i].Key = v
		}
	}
}

func containsString(list []string, s string) bool {
	for _, x := range list {
		if x == s {
			return true
		}
	}
	return false
}

func removeString(list []string, s string) []string {
	out := list[:0]
	for _, x := range list {
		if x != s {
			out = append(out, x)
		}
	}
	return out
}

// ListKeys returns the persisted gateway keys.
func (s *Store) ListKeys() []GWKey {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]GWKey, len(s.data.Keys))
	copy(out, s.data.Keys)
	return out
}

// KeyByValue looks up a gateway key record by its secret value.
func (s *Store) KeyByValue(key string) (GWKey, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, k := range s.data.Keys {
		if k.Key == key {
			return k, true
		}
	}
	return GWKey{}, false
}

// SaveKey upserts a gateway key record and persists.
func (s *Store) SaveKey(k GWKey) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i := range s.data.Keys {
		if s.data.Keys[i].Key == k.Key {
			s.data.Keys[i] = k
			return s.persistLocked()
		}
	}
	s.data.Keys = append(s.data.Keys, k)
	return s.persistLocked()
}

// DeleteKey removes a gateway key record and persists.
func (s *Store) DeleteKey(key string) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	kept := s.data.Keys[:0]
	removed := false
	for _, k := range s.data.Keys {
		if k.Key == key {
			removed = true
			continue
		}
		kept = append(kept, k)
	}
	if !removed {
		return false, nil
	}
	s.data.Keys = kept
	return true, s.persistLocked()
}

// AutoRules returns the persisted AUTO scheduling slots.
func (s *Store) AutoRules() []ModelScope {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]ModelScope, len(s.data.Auto))
	copy(out, s.data.Auto)
	return out
}

// SaveAutoRules persists the AUTO scheduling slots.
func (s *Store) SaveAutoRules(entries []ModelScope) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.data.Auto = entries
	return s.persistLocked()
}