nixos-router/pkgs/router-webui/tokens.go

package main

import (
	"bufio"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"time"
)

// API tokens for programmatic access (the MCP endpoint). A token is shown once
// when it is made; only its SHA-256 is stored. Scopes:
//
//	read    state and settings
//	change  also edit the settings (they stay unapplied until a person applies
//	        them in the web UI)
//
// An "apply" scope is reserved for later.
const (
	scopeRead   = "read"
	scopeChange = "change"
	tokenPrefix = "rtr_"
	maxTokens   = 32
)

type Token struct {
	ID       string    `json:"id"`
	Name     string    `json:"name"`
	Scope    string    `json:"scope"`
	Hash     string    `json:"hash"` // hex SHA-256 of the secret
	Created  time.Time `json:"created"`
	LastUsed time.Time `json:"lastUsed,omitzero"`
}

func (s System) tokensPath() string { return filepath.Join(s.StateDir, "webui", "tokens.json") }
func (s System) auditPath() string  { return filepath.Join(s.StateDir, "webui", "audit.log") }

// tokenStore keeps the tokens in a file; every call reads it, so edits and
// restarts need no care.
type tokenStore struct {
	mu       sync.Mutex
	path     string
	lastSave map[string]time.Time // when LastUsed of a token was last written
}

func newTokenStore(path string) *tokenStore {
	return &tokenStore{path: path, lastSave: map[string]time.Time{}}
}

func (t *tokenStore) load() []Token {
	data, err := os.ReadFile(t.path)
	if err != nil {
		return nil
	}
	var toks []Token
	if json.Unmarshal(data, &toks) != nil {
		return nil
	}
	return toks
}

func (t *tokenStore) save(toks []Token) error {
	data, err := json.MarshalIndent(toks, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(t.path), 0o700); err != nil {
		return err
	}
	return writeFileAtomic(t.path, append(data, '\n'), 0o600)
}

func hashSecret(secret string) string {
	h := sha256.Sum256([]byte(secret))
	return hex.EncodeToString(h[:])
}

func (t *tokenStore) List() []Token {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.load()
}

// Create makes a token and returns its secret, which is not stored.
func (t *tokenStore) Create(name, scope string) (Token, string, error) {
	name = strings.TrimSpace(name)
	switch {
	case name == "" || len(name) > 40:
		return Token{}, "", errors.New("give the token a name of up to 40 characters")
	case scope != scopeRead && scope != scopeChange:
		return Token{}, "", fmt.Errorf("unknown scope %q", scope)
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	toks := t.load()
	if len(toks) >= maxTokens {
		return Token{}, "", errors.New("too many tokens; revoke one first")
	}
	raw := make([]byte, 32)
	if _, err := rand.Read(raw); err != nil {
		return Token{}, "", err
	}
	id := make([]byte, 4)
	if _, err := rand.Read(id); err != nil {
		return Token{}, "", err
	}
	secret := tokenPrefix + base64.RawURLEncoding.EncodeToString(raw)
	tok := Token{ID: hex.EncodeToString(id), Name: name, Scope: scope, Hash: hashSecret(secret), Created: time.Now().UTC().Truncate(time.Second)}
	return tok, secret, t.save(append(toks, tok))
}

func (t *tokenStore) Revoke(id string) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	toks := t.load()
	n := len(toks)
	toks = slices.DeleteFunc(toks, func(k Token) bool { return k.ID == id })
	if len(toks) == n {
		return errors.New("no such token")
	}
	return t.save(toks)
}

// Lookup finds the token for a secret and notes that it was used.
func (t *tokenStore) Lookup(secret string) (Token, bool) {
	if !strings.HasPrefix(secret, tokenPrefix) {
		return Token{}, false
	}
	want := hashSecret(secret)
	t.mu.Lock()
	defer t.mu.Unlock()
	toks := t.load()
	found := -1
	for i, k := range toks {
		if subtle.ConstantTimeCompare([]byte(k.Hash), []byte(want)) == 1 {
			found = i
		}
	}
	if found < 0 {
		return Token{}, false
	}
	now := time.Now().UTC().Truncate(time.Second)
	tok := toks[found]
	tok.LastUsed = now
	// Writing on every call would wear the disk; once in a few minutes is enough.
	if now.Sub(t.lastSave[tok.ID]) > 5*time.Minute {
		toks[found] = tok
		if t.save(toks) == nil {
			t.lastSave[tok.ID] = now
		}
	}
	return tok, true
}

func (k Token) allows(scope string) bool {
	return scope == scopeRead || k.Scope == scopeChange
}

// Audit log: one JSON line per call of a tool, kept to the last 1 MB (and the
// file before it).
type AuditEntry struct {
	Time  time.Time `json:"time"`
	Token string    `json:"token"`
	Tool  string    `json:"tool"`
	Args  string    `json:"args,omitempty"`
	Error string    `json:"error,omitempty"`
}

var auditMu sync.Mutex

func (s System) audit(e AuditEntry) {
	auditMu.Lock()
	defer auditMu.Unlock()
	if len(e.Args) > 300 {
		e.Args = e.Args[:300] + "…"
	}
	path := s.auditPath()
	if fi, err := os.Stat(path); err == nil && fi.Size() > 1<<20 {
		_ = os.Rename(path, path+".1")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return
	}
	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
	if err != nil {
		return
	}
	defer f.Close()
	line, _ := json.Marshal(e)
	fmt.Fprintf(f, "%s\n", line)
}

// RecentAudit returns the last n entries, newest first.
func (s System) RecentAudit(n int) []AuditEntry {
	auditMu.Lock()
	defer auditMu.Unlock()
	f, err := os.Open(s.auditPath())
	if err != nil {
		return nil
	}
	defer f.Close()
	var all []AuditEntry
	sc := bufio.NewScanner(f)
	sc.Buffer(make([]byte, 0, 64*1024), 1<<20)
	for sc.Scan() {
		var e AuditEntry
		if json.Unmarshal(sc.Bytes(), &e) == nil {
			all = append(all, e)
		}
	}
	slices.Reverse(all)
	if len(all) > n {
		all = all[:n]
	}
	return all
}