sift/cmd/sift/serve.go

package main

import (
	"context"
	"flag"
	"log/slog"
	"net"
	"net/netip"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"git.wntrmute.dev/kyle/sift/internal/api"
	"git.wntrmute.dev/kyle/sift/internal/auth"
	"git.wntrmute.dev/kyle/sift/internal/config"
	"git.wntrmute.dev/kyle/sift/internal/dcache"
	"git.wntrmute.dev/kyle/sift/internal/digest"
	"git.wntrmute.dev/kyle/sift/internal/fetch"
	"git.wntrmute.dev/kyle/sift/internal/fleetctx"
	"git.wntrmute.dev/kyle/sift/internal/guard"
	"git.wntrmute.dev/kyle/sift/internal/selfip"
)

type denyFlags []netip.Prefix

func (d *denyFlags) String() string {
	var b []string
	for _, p := range *d {
		b = append(b, p.String())
	}
	return strings.Join(b, ",")
}

func (d *denyFlags) Set(s string) error {
	p, err := netip.ParsePrefix(s)
	if err != nil {
		return err
	}
	*d = append(*d, p)
	return nil
}

func defaultConfigPath() string {
	if p := os.Getenv("SIFT_CONFIG"); p != "" {
		return p
	}
	return "/etc/sift/config.toml"
}

func serveMain() {
	fs := flag.NewFlagSet("serve", flag.ExitOnError)
	listen := fs.String("listen", "127.0.0.1:8890", "listen address")
	configPath := fs.String("config", defaultConfigPath(), "config file")
	keysPath := fs.String("keys", "/var/lib/sift/keys.toml", "caller keys file")
	ua := fs.String("user-agent", "sift/"+version, "outbound User-Agent")
	var deny denyFlags
	fs.Var(&deny, "deny-prefix", "extra CIDR to block (repeatable)")
	g := &guard.Guard{}
	registerTestFlags(fs, g)
	_ = fs.Parse(os.Args[1:])

	log := slog.New(slog.NewTextHandler(os.Stdout, nil))
	cfg, err := config.Load(*configPath)
	if err != nil {
		log.Error("config", "err", err, "path", *configPath)
		os.Exit(1)
	}
	if *listen == "127.0.0.1:8890" && cfg.Listen != "" {
		*listen = cfg.Listen
	}
	if *ua == "sift/"+version && cfg.UserAgent != "" {
		*ua = cfg.UserAgent
	}
	for _, s := range cfg.DenyPrefixes {
		p, err := netip.ParsePrefix(s)
		if err != nil {
			log.Error("deny_prefixes", "err", err, "value", s)
			os.Exit(1)
		}
		deny = append(deny, p)
	}
	g.DenyPrefixes = deny
	if n := len(deny); n > 0 {
		log.Info("deny prefixes loaded", "count", n)
	}

	store, err := auth.Load(*keysPath)
	if err != nil {
		log.Error("keys", "err", err, "path", *keysPath)
		os.Exit(1)
	}

	fetcher := &fetch.Fetcher{Guard: g, UA: *ua}
	parent, stop := context.WithCancel(context.Background())
	defer stop()

	srv := api.New(version, g, fetcher, store, log)
	if cfg.Dixie.Chat != "" {
		cachePath := cfg.Digest.Cache
		if cachePath == "" {
			cachePath = "/var/lib/sift/cache.db"
		}
		cache, err := dcache.Open(cachePath, 500<<20, log)
		if err != nil {
			log.Warn("digest cache", "err", err)
		}
		srv.Digest = &digest.Engine{
			ChatURL:    cfg.Dixie.Chat,
			EmbedURL:   cfg.Dixie.Embed,
			EmbedModel: cfg.Dixie.EmbedModel,
			Model:      cfg.Digest.Model,
			Thinking:   cfg.Digest.Thinking,
			Gate:       digest.NewGate(cfg.Dixie.MaxInflight, cfg.Dixie.QueueDepth, config.ParseDuration(cfg.Dixie.QueueWait, digest.DefaultQWait)),
			Cache:      cache,
			Log:        log,
		}
	}
	if cfg.Context.Honcho != "" {
		allow := cfg.Context.Workspaces
		if len(allow) == 0 {
			allow = []string{"hermes-fleet"}
		}
		srv.Ctx = &fleetctx.Client{Base: cfg.Context.Honcho, Allow: allow, Log: log}
	}
	if len(cfg.SelfIP.Sources) > 0 {
		disc := &selfip.Discoverer{
			Cfg: selfip.Config{
				Sources:      cfg.SelfIP.Sources,
				Interval:     config.ParseDuration(cfg.SelfIP.Interval, selfip.DefaultInterval),
				KeepPrevious: config.ParseDuration(cfg.SelfIP.KeepPrevious, selfip.DefaultKeepPrevious),
				Path:         cfg.SelfIP.Path,
			},
			Guard: g,
			Log:   log,
			Get: func(ctx context.Context, u string) ([]byte, error) {
				res, err := fetcher.Fetch(ctx, u)
				if err != nil {
					return nil, err
				}
				return res.Body, nil
			},
		}
		disc.Start(parent)
		srv.SelfIP = disc
	}

	hs := srv.HTTPServer(*listen)
	ln, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Error("listen", "err", err)
		os.Exit(1)
	}
	log.Info("sift listening", "addr", ln.Addr().String(), "version", version)

	go func() {
		ch := make(chan os.Signal, 1)
		signal.Notify(ch, syscall.SIGHUP)
		for range ch {
			if err := store.Reload(); err != nil {
				log.Error("keys reload", "err", err)
			} else {
				log.Info("keys reloaded")
			}
		}
	}()

	errCh := make(chan error, 1)
	go func() { errCh <- hs.Serve(ln) }()
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	select {
	case <-sig:
		stop()
		_ = hs.Close()
	case err := <-errCh:
		stop()
		if err != nil {
			log.Error("serve", "err", err)
			os.Exit(1)
		}
	}
}