ollama-proxy/internal/config/accounts.go

// Package config holds runtime configuration for ollama-proxy: the list of
// Ollama Cloud accounts and the server runtime options.
package config

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"time"
)

// Account is a single Ollama Cloud API key with an optional alias.
type Account struct {
	ID      string    `json:"id"`
	Name    string    `json:"name"`
	APIKey  string    `json:"api_key"`
	BaseURL string    `json:"base_url,omitempty"` // overrides ServerConfig.BaseURL; empty = use default
	Created time.Time `json:"created"`
}

// AccountsFile is the on-disk JSON structure.
type AccountsFile struct {
	BaseURL  string    `json:"base_url,omitempty"` // default upstream URL
	Accounts []Account `json:"accounts"`
}

// AccountsPath returns the canonical accounts.json path under XDG_CONFIG_HOME
// (or ~/.config when XDG_CONFIG_HOME is unset). The parent directory must exist.
func AccountsPath() (string, error) {
	dir := os.Getenv("XDG_CONFIG_HOME")
	if dir == "" {
		home, err := os.UserHomeDir()
		if err != nil {
			return "", err
		}
		dir = filepath.Join(home, ".config")
	}
	return filepath.Join(dir, "ollama-proxy", "accounts.json"), nil
}

// EnsureAccountsDir creates ~/.config/ollama-proxy with 0700 perms if missing.
func EnsureAccountsDir() error {
	p, err := AccountsPath()
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil {
		return fmt.Errorf("create config dir: %w", err)
	}
	return nil
}

// LoadAccounts reads accounts.json from the default path. Returns an empty
// AccountsFile (not an error) when the file does not exist.
func LoadAccounts() (*AccountsFile, error) {
	p, err := AccountsPath()
	if err != nil {
		return nil, err
	}
	data, err := os.ReadFile(p)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return &AccountsFile{}, nil
		}
		return nil, fmt.Errorf("read accounts: %w", err)
	}
	var af AccountsFile
	if err := json.Unmarshal(data, &af); err != nil {
		return nil, fmt.Errorf("parse accounts: %w", err)
	}
	return &af, nil
}

// Save writes accounts.json atomically with 0600 permissions. The temp file is
// written in the same directory and renamed, so a crash mid-write never leaves
// a truncated file. The parent directory must already exist (see EnsureAccountsDir).
func (af *AccountsFile) Save() error {
	p, err := AccountsPath()
	if err != nil {
		return err
	}
	if af.Accounts == nil {
		af.Accounts = []Account{}
	}
	data, err := json.MarshalIndent(af, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal accounts: %w", err)
	}
	dir := filepath.Dir(p)
	tmp, err := os.CreateTemp(dir, ".accounts.json.*")
	if err != nil {
		return fmt.Errorf("create temp: %w", err)
	}
	tmpName := tmp.Name()
	defer func() {
		if err != nil {
			_ = os.Remove(tmpName)
		}
	}()
	if err := tmp.Chmod(0o600); err != nil {
		tmp.Close()
		return fmt.Errorf("chmod temp: %w", err)
	}
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		return fmt.Errorf("write temp: %w", err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("close temp: %w", err)
	}
	if err := os.Rename(tmpName, p); err != nil {
		return fmt.Errorf("rename temp: %w", err)
	}
	if err := os.Chmod(p, 0o600); err != nil {
		return fmt.Errorf("chmod final: %w", err)
	}
	return nil
}

// NewID returns a random 16-byte hex string suitable as a stable account ID.
func NewID() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		panic("crypto/rand failed: " + err.Error())
	}
	return hex.EncodeToString(b)
}

// shortID returns the first 8 chars of an ID, for display only.
func shortID(id string) string {
	if len(id) < 8 {
		return id
	}
	return id[:8]
}

// maskKey renders an API key as "****ABCD" (last 4 chars) for safe display.
func maskKey(key string) string {
	if len(key) <= 4 {
		return "****"
	}
	return "****" + key[len(key)-4:]
}

// FindByName returns the first account matching name (case-sensitive). Returns
// nil when no account matches.
func (af *AccountsFile) FindByName(name string) *Account {
	for i := range af.Accounts {
		if af.Accounts[i].Name == name {
			return &af.Accounts[i]
		}
	}
	return nil
}

// FindByID returns the account with the given ID prefix (matches short ID too).
func (af *AccountsFile) FindByID(idPrefix string) *Account {
	for i := range af.Accounts {
		if af.Accounts[i].ID == idPrefix || hasPrefix(af.Accounts[i].ID, idPrefix) {
			return &af.Accounts[i]
		}
	}
	return nil
}

func hasPrefix(s, prefix string) bool {
	if len(prefix) > len(s) {
		return false
	}
	return s[:len(prefix)] == prefix
}

// Find returns the account matching the given identifier, which may be either
// a name or an ID/ID-prefix. Name match takes precedence.
func (af *AccountsFile) Find(ident string) *Account {
	if a := af.FindByName(ident); a != nil {
		return a
	}
	return af.FindByID(ident)
}

// Remove deletes the account with the given ID or name. Returns ErrNotFound
// when no account matched.
func (af *AccountsFile) Remove(ident string) error {
	idx := -1
	for i := range af.Accounts {
		if af.Accounts[i].Name == ident || af.Accounts[i].ID == ident || hasPrefix(af.Accounts[i].ID, ident) {
			idx = i
			break
		}
	}
	if idx < 0 {
		return ErrNotFound
	}
	af.Accounts = append(af.Accounts[:idx], af.Accounts[idx+1:]...)
	return nil
}

// ErrNotFound is returned when an account lookup or removal misses.
var ErrNotFound = errors.New("account not found")

// ValidateName returns an error when name contains characters outside
// [A-Za-z0-9_-]. Names must be safe for CLI display and log output.
func ValidateName(name string) error {
	ok, err := regexp.MatchString(`^[A-Za-z0-9_-]+$`, name)
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("name must contain only letters, digits, '_' or '-'")
	}
	return nil
}