ai-router/internal/server/admin.go

package server

import (
	"database/sql"
	"encoding/json"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type channel struct {
	ID       string            `json:"id"`
	Name     string            `json:"name"`
	BaseURL  string            `json:"base_url"`
	Protocol string            `json:"protocol"`
	Enabled  bool              `json:"enabled"`
	Models   map[string]string `json:"models"`
	APIKey   string            `json:"api_key,omitempty"`
}
type apiKey struct {
	Revoked     bool     `json:"revoked"`
	ID          string   `json:"id"`
	Name        string   `json:"name"`
	Prefix      string   `json:"prefix"`
	Enabled     bool     `json:"enabled"`
	AllChannels bool     `json:"all_channels"`
	Channels    []string `json:"channels"`
	ExpiresAt   string   `json:"expires_at"`
	CreatedAt   string   `json:"created_at"`
}

func (s *Server) loadChannels() ([]channel, error) {
	rows, e := s.db.Query("SELECT id,name,base_url,protocol,enabled,models,api_key FROM channels ORDER BY name")
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	out := []channel{}
	for rows.Next() {
		var c channel
		var m string
		if e = rows.Scan(&c.ID, &c.Name, &c.BaseURL, &c.Protocol, &c.Enabled, &m, &c.APIKey); e != nil {
			return nil, e
		}
		if e = json.Unmarshal([]byte(m), &c.Models); e != nil {
			return nil, e
		}
		out = append(out, c)
	}
	return out, rows.Err()
}
func (s *Server) channels(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		cs, e := s.loadChannels()
		if e != nil {
			errJSON(w, 500, "database error")
			return
		}
		for i := range cs {
			cs[i].APIKey = ""
		}
		writeJSON(w, cs)
		return
	}
	if r.Method != "POST" {
		errJSON(w, 405, "method not allowed")
		return
	}
	var c channel
	if decode(w, r, &c) != nil {
		errJSON(w, 400, "invalid channel")
		return
	}
	u, e := url.Parse(c.BaseURL)
	if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(c.BaseURL, "#") || c.Name == "" || strings.Contains(c.Name, "/") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
		errJSON(w, 400, "invalid channel")
		return
	}
	if c.Models == nil {
		c.Models = map[string]string{}
	}
	for a, m := range c.Models {
		if a == "" || m == "" {
			errJSON(w, 400, "invalid model mapping")
			return
		}
	}
	c.BaseURL = strings.TrimRight(c.BaseURL, "/")
	m, _ := json.Marshal(c.Models)
	if c.ID == "" {
		c.ID = randomID()
		_, e = s.db.Exec("INSERT INTO channels VALUES(?,?,?,?,?,?,?,?)", c.ID, c.Name, c.BaseURL, c.Protocol, c.Enabled, false, string(m), c.APIKey)
	} else {
		var old string
		if e = s.db.QueryRow("SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old); e == sql.ErrNoRows {
			errJSON(w, 404, "not found")
			return
		} else if e != nil {
			errJSON(w, 500, "database error")
			return
		}
		if c.APIKey == "" {
			c.APIKey = old
		}
		_, e = s.db.Exec("UPDATE channels SET name=?,base_url=?,protocol=?,enabled=?,models=?,api_key=? WHERE id=?", c.Name, c.BaseURL, c.Protocol, c.Enabled, string(m), c.APIKey, c.ID)
	}
	if e != nil {
		errJSON(w, 400, "channel could not be saved; names must be unique")
		return
	}
	c.APIKey = ""
	writeJSON(w, c)
}
func (s *Server) channelID(w http.ResponseWriter, r *http.Request) {
	if r.Method != "DELETE" {
		errJSON(w, 405, "method not allowed")
		return
	}
	res, e := s.db.Exec("DELETE FROM channels WHERE id=?", strings.TrimPrefix(r.URL.Path, "/admin/api/channels/"))
	mutation(w, res, e)
}
func mutation(w http.ResponseWriter, res sql.Result, e error) {
	if e != nil {
		errJSON(w, 500, "database error")
		return
	}
	n, e := res.RowsAffected()
	if e != nil {
		errJSON(w, 500, "database error")
		return
	}
	if n == 0 {
		errJSON(w, 404, "not found")
		return
	}
	writeJSON(w, map[string]bool{"ok": true})
}
func (s *Server) keys(w http.ResponseWriter, r *http.Request) {
	if r.Method == "GET" {
		rows, e := s.db.Query("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at,revoked FROM keys ORDER BY created_at DESC")
		if e != nil {
			errJSON(w, 500, "database error")
			return
		}
		defer rows.Close()
		out := []apiKey{}
		for rows.Next() {
			var k apiKey
			var ch string
			if rows.Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt, &k.Revoked) != nil {
				errJSON(w, 500, "database error")
				return
			}
			if json.Unmarshal([]byte(ch), &k.Channels) != nil {
				errJSON(w, 500, "database error")
				return
			}
			out = append(out, k)
		}
		if rows.Err() != nil {
			errJSON(w, 500, "database error")
			return
		}
		writeJSON(w, out)
		return
	}
	if r.Method != "POST" {
		errJSON(w, 405, "method not allowed")
		return
	}
	var k apiKey
	if decode(w, r, &k) != nil || k.Name == "" {
		errJSON(w, 400, "invalid key")
		return
	}
	if k.ExpiresAt != "" {
		t, e := time.Parse(time.RFC3339, k.ExpiresAt)
		if e != nil || !t.After(time.Now()) {
			errJSON(w, 400, "invalid expiry")
			return
		}
		k.ExpiresAt = t.UTC().Format(time.RFC3339Nano)
	}
	if k.Channels == nil {
		k.Channels = []string{}
	}
	for _, id := range k.Channels {
		var n int
		if s.db.QueryRow("SELECT count(*) FROM channels WHERE id=?", id).Scan(&n) != nil || n == 0 {
			errJSON(w, 400, "invalid channel permission")
			return
		}
	}
	raw := "sk-" + randomID()
	ch, _ := json.Marshal(k.Channels)
	_, e := s.db.Exec("INSERT INTO keys(id,name,prefix,hash,enabled,all_channels,channels,expires_at,created_at) VALUES(?,?,?,?,?,?,?,?,?)", randomID(), k.Name, raw[:11], hash(raw), true, k.AllChannels, string(ch), k.ExpiresAt, time.Now().UTC().Format(time.RFC3339Nano))
	if e != nil {
		errJSON(w, 500, "database error")
		return
	}
	writeJSON(w, map[string]string{"key": raw})
}
func (s *Server) keyID(w http.ResponseWriter, r *http.Request) {
	id := strings.TrimPrefix(r.URL.Path, "/admin/api/keys/")
	if purgeID, ok := strings.CutSuffix(id, "/purge"); ok {
		if r.Method != "DELETE" {
			errJSON(w, 405, "method not allowed")
			return
		}
		var revoked bool
		if e := s.db.QueryRow("SELECT revoked FROM keys WHERE id=?", purgeID).Scan(&revoked); e == sql.ErrNoRows {
			errJSON(w, 404, "not found")
			return
		} else if e != nil {
			errJSON(w, 500, "database error")
			return
		}
		if !revoked {
			errJSON(w, 409, "revoke the key before deleting")
			return
		}
		res, e := s.db.Exec("DELETE FROM keys WHERE id=?", purgeID)
		mutation(w, res, e)
		return
	}
	if r.Method == "DELETE" {
		res, e := s.db.Exec("UPDATE keys SET enabled=0,revoked=1 WHERE id=?", id)
		mutation(w, res, e)
		return
	}
	if r.Method == "PATCH" {
		var x struct {
			Enabled *bool `json:"enabled"`
		}
		if decode(w, r, &x) != nil || x.Enabled == nil {
			errJSON(w, 400, "invalid request")
			return
		}
		res, e := s.db.Exec("UPDATE keys SET enabled=? WHERE id=? AND revoked=0", *x.Enabled, id)
		mutation(w, res, e)
		return
	}
	errJSON(w, 405, "method not allowed")
}
func (s *Server) authenticate(r *http.Request) (apiKey, error) {
	var k apiKey
	raw := r.Header.Get("x-api-key")
	if a := r.Header.Get("Authorization"); a != "" {
		parts := strings.Fields(a)
		if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
			return k, sql.ErrNoRows
		}
		if raw != "" && raw != parts[1] {
			return k, sql.ErrNoRows
		}
		raw = parts[1]
	}
	if raw == "" {
		return k, sql.ErrNoRows
	}
	var ch string
	e := s.db.QueryRow("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at FROM keys WHERE hash=? AND enabled=1 AND revoked=0", hash(raw)).Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt)
	if e != nil {
		return k, e
	}
	if k.ExpiresAt != "" {
		t, e := time.Parse(time.RFC3339Nano, k.ExpiresAt)
		if e != nil || !time.Now().Before(t) {
			return apiKey{}, sql.ErrNoRows
		}
	}
	e = json.Unmarshal([]byte(ch), &k.Channels)
	return k, e
}
func (k apiKey) allows(id string) bool {
	if k.AllChannels {
		return true
	}
	for _, v := range k.Channels {
		if v == id {
			return true
		}
	}
	return false
}