byre/cmd/byre/serve.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"g.fc.run/ai/byre/internal/server"
	"g.fc.run/ai/byre/internal/store"
	"tailscale.com/tsnet"
)

const (
	defaultListen = "127.0.0.1:8270"
	loopbackURL   = "http://" + defaultListen
)

func serveCmd(args []string) {
	configDir, err := os.UserConfigDir()
	if err != nil {
		fatal("user config dir", "err", err)
	}

	flags := flag.NewFlagSet("serve", flag.ExitOnError)
	dataDir := flags.String("data", filepath.Join(configDir, "byre"), "directory holding byre's database, worktrees, leads and blobs")
	listenAddr := flags.String("listen", defaultListen, "loopback address serving GET /api/tailnet")
	ports := flags.Int("ports", store.DefaultPortFirst, "first port of the tasks' port blocks")
	hostname := flags.String("ts-hostname", "byre-"+macSlug(), "hostname of byre's tailnet node")
	flags.Parse(args)

	host, _, err := net.SplitHostPort(*listenAddr)
	if err != nil {
		fatal("parse --listen", "err", err)
	}
	ip := net.ParseIP(host)
	if ip == nil || !ip.IsLoopback() {
		fatal("--listen must be a loopback address", "addr", *listenAddr)
	}

	// Protect local credentials before tsnet can write state or a login URL.
	if err := os.MkdirAll(*dataDir, 0o700); err != nil {
		fatal("create data directory", "err", err)
	}
	if err := os.Chmod(*dataDir, 0o700); err != nil {
		fatal("protect data directory", "err", err)
	}
	// Keep ownership until shutdown; never unlink the lock file, since a new
	// inode would let another process lock the same data directory concurrently.
	lock, err := os.OpenFile(filepath.Join(*dataDir, "serve.lock"), os.O_CREATE|os.O_RDWR, 0o600)
	if err != nil {
		fatal("open service lock", "err", err)
	}
	defer lock.Close()
	if err := syscall.Flock(int(lock.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil {
		fatal("acquire service lock", "data", *dataDir, "err", err)
	}

	token := rand.Text()
	if err := os.WriteFile(filepath.Join(*dataDir, "loopback-token"), []byte(token), 0o600); err != nil {
		fatal("write loopback credential", "err", err)
	}
	home, err := os.UserHomeDir()
	if err != nil {
		fatal("home directory", "err", err)
	}
	logDir := filepath.Join(home, "Library", "Logs")
	if err := os.MkdirAll(logDir, 0o755); err != nil {
		fatal("create log directory", "err", err)
	}
	logFile, err := os.OpenFile(filepath.Join(logDir, "byre.log"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	if err != nil {
		fatal("open service log", "err", err)
	}
	if err := logFile.Chmod(0o600); err != nil {
		fatal("protect service log", "err", err)
	}
	logFile.Close()

	ts := &tsnet.Server{
		Dir:      filepath.Join(*dataDir, "tsnet"),
		Hostname: *hostname,
		UserLogf: tsnetLogf(),
	}
	err = ts.Start()
	if err != nil {
		fatal("start tsnet", "err", err)
	}
	defer ts.Close()

	srv, err := server.New(*dataDir, *ports, ts)
	if err != nil {
		fatal("start server", "err", err)
	}
	defer srv.Close()

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	ctx, replaced := context.WithCancel(ctx)
	defer replaced()
	go watchExecutable(ctx, replaced)
	go updateDaily(ctx)

	base := func(net.Listener) context.Context { return ctx }
	loopback := &http.Server{Addr: *listenAddr, Handler: srv.LoopbackHandler(token), BaseContext: base}
	tailnet := &http.Server{Handler: srv.Handler(), BaseContext: base}

	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		loopback.Shutdown(shutdownCtx)
		tailnet.Shutdown(shutdownCtx)
	}()

	go serveTailnet(ctx, ts, tailnet)

	slog.Info("listening", "addr", *listenAddr, "hostname", *hostname, "data", *dataDir, "ports", *ports)
	err = loopback.ListenAndServe()
	if !errors.Is(err, http.ErrServerClosed) {
		fatal("serve", "err", err)
	}
}

func serveTailnet(ctx context.Context, ts *tsnet.Server, hs *http.Server) {
	st, err := ts.Up(ctx)
	if ctx.Err() != nil {
		return
	}
	if err != nil {
		fatal("tsnet up", "err", err)
	}
	if len(st.CertDomains) == 0 {
		slog.Error("tailnet has no https certificates; turn on HTTPS on the DNS page of the tailscale admin console, then restart byre")
		return
	}
	lc, err := ts.LocalClient()
	if err != nil {
		fatal("tsnet local client", "err", err)
	}
	ln, err := ts.Listen("tcp", ":443")
	if err != nil {
		fatal("tsnet listen", "err", err)
	}
	tlsLn := tls.NewListener(ln, &tls.Config{GetCertificate: lc.GetCertificate, NextProtos: []string{"h2", "http/1.1"}})
	slog.Info("listening on tailnet", "origin", "https://"+st.CertDomains[0])
	err = hs.Serve(tlsLn)
	if !errors.Is(err, http.ErrServerClosed) {
		fatal("serve tailnet", "err", err)
	}
}

func tsnetLogf() func(string, ...any) {
	mu := sync.Mutex{}
	last := ""
	return func(format string, args ...any) {
		msg := fmt.Sprintf(format, args...)
		mu.Lock()
		repeat := msg == last
		last = msg
		mu.Unlock()
		if repeat {
			return
		}
		slog.Info("tsnet", "msg", msg)
	}
}

func macSlug() string {
	out, err := exec.Command("scutil", "--get", "ComputerName").Output()
	if err != nil {
		fatal("computer name", "err", err)
	}
	slug := strings.Builder{}
	dash := false
	for _, r := range strings.ToLower(strings.TrimSpace(string(out))) {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') {
			if dash && slug.Len() > 0 {
				slug.WriteByte('-')
			}
			slug.WriteRune(r)
			dash = false
			continue
		}
		dash = true
	}
	return slug.String()
}

func watchExecutable(ctx context.Context, replaced context.CancelFunc) {
	exe, err := os.Executable()
	if err != nil {
		fatal("executable path", "err", err)
	}
	started, err := os.Stat(exe)
	if err != nil {
		fatal("stat executable", "path", exe, "err", err)
	}
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		now, err := os.Stat(exe)
		if err != nil {
			continue
		}
		if os.SameFile(started, now) {
			continue
		}
		slog.Info("executable replaced, exiting", "path", exe)
		replaced()
		return
	}
}

func updateDaily(ctx context.Context) {
	for {
		update(ctx)
		select {
		case <-ctx.Done():
			return
		case <-time.After(24 * time.Hour):
		}
	}
}

func update(ctx context.Context) {
	bin, err := exec.LookPath("byre-update")
	if err != nil {
		slog.Error("find byre-update", "err", err)
		return
	}
	ctx, cancel := context.WithTimeout(ctx, 30*time.Minute)
	defer cancel()
	out, err := exec.CommandContext(ctx, bin).CombinedOutput()
	if err != nil {
		slog.Error("byre-update", "err", err, "output", strings.TrimSpace(string(out)))
		return
	}
	slog.Info("byre-update", "output", strings.TrimSpace(string(out)))
}