Warp/server/cmd/warp/commands.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log/slog"
	"os"

	"github.com/dariuszjaniak/warp/server/internal/buildinfo"
	"github.com/dariuszjaniak/warp/server/internal/config"
	"github.com/dariuszjaniak/warp/server/internal/logging"
	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/setup"
	"github.com/dariuszjaniak/warp/server/internal/storage"
	_ "github.com/dariuszjaniak/warp/server/internal/storage/drivers" // link in the storage drivers
)

// defaultConfigPath is used when neither --config nor WARP_CONFIG is given
// and the file happens to exist next to the binary's working directory.
const defaultConfigPath = "warp.yaml"

// newFlagSet builds a flag set that prints usage to stdout on --help.
func newFlagSet(name string) *flag.FlagSet {
	fs := flag.NewFlagSet(name, flag.ContinueOnError)
	fs.SetOutput(os.Stdout)
	return fs
}

// addConfigFlag registers --config and returns the resolver for it.
func addConfigFlag(fs *flag.FlagSet) func() (*config.Config, error) {
	path := fs.String("config", "", "path to warp.yaml (default: $WARP_CONFIG, then ./warp.yaml)")
	return func() (*config.Config, error) {
		resolved := resolveConfigPath(*path)
		cfg, err := config.Load(resolved)
		if err != nil {
			return nil, err
		}
		return cfg, nil
	}
}

// resolveConfigPath implements the --config / $WARP_CONFIG / ./warp.yaml
// precedence. An empty result means "defaults plus environment", which is
// how a container with only WARP_* variables starts.
func resolveConfigPath(flagValue string) string {
	if flagValue != "" {
		return flagValue
	}
	if env := os.Getenv("WARP_CONFIG"); env != "" {
		return env
	}
	if _, err := os.Stat(defaultConfigPath); err == nil {
		return defaultConfigPath
	}
	return ""
}

// newLogger builds the logger described by the config.
func newLogger(cfg *config.Config) (*slog.Logger, error) {
	return logging.New(os.Stderr, logging.Options{
		Level:     cfg.Logging.Level,
		Format:    cfg.Logging.Format,
		AddSource: cfg.Logging.AddSource,
	})
}

func runServe(ctx context.Context, args []string) error {
	fs := newFlagSet("serve")
	loadConfig := addConfigFlag(fs)
	if err := fs.Parse(args); err != nil {
		return err
	}

	cfg, err := loadConfig()
	if err != nil {
		return err
	}
	log, err := newLogger(cfg)
	if err != nil {
		return err
	}
	slog.SetDefault(log)

	if cfg.Path() == "" {
		log.Warn("no configuration file; using defaults and WARP_* environment variables")
	} else {
		log.Info("loaded configuration", "path", cfg.Path())
	}

	server, err := setup.New(ctx, cfg, setup.Options{Logger: log})
	if err != nil {
		return err
	}
	defer func() {
		if err := server.Close(); err != nil {
			log.Error("closing storage failed", "error", err)
		}
	}()

	return server.Run(ctx)
}

func runMigrate(ctx context.Context, args []string) error {
	fs := newFlagSet("migrate")
	loadConfig := addConfigFlag(fs)
	statusOnly := fs.Bool("status", false, "list pending migrations without applying them")
	if err := fs.Parse(args); err != nil {
		return err
	}

	cfg, err := loadConfig()
	if err != nil {
		return err
	}
	log, err := newLogger(cfg)
	if err != nil {
		return err
	}
	slog.SetDefault(log)

	db, err := openDatabase(ctx, cfg, log)
	if err != nil {
		return err
	}
	defer func() { _ = db.Close() }()

	pending, err := db.PendingMigrations(ctx)
	if err != nil {
		return err
	}
	if len(pending) == 0 {
		fmt.Printf("%s: schema is up to date\n", db.Driver())
		return nil
	}
	if *statusOnly {
		fmt.Printf("%s: %d pending migration(s):\n", db.Driver(), len(pending))
		for _, name := range pending {
			fmt.Printf("  %s\n", name)
		}
		return nil
	}
	if err := db.Migrate(ctx); err != nil {
		return err
	}
	fmt.Printf("%s: applied %d migration(s)\n", db.Driver(), len(pending))
	return nil
}

func runGenerateKeys(_ context.Context, args []string) error {
	fs := newFlagSet("generate-keys")
	loadConfig := addConfigFlag(fs)
	output := fs.String("output", "", "where to write the key (default: global.private_key from the config)")
	keyID := fs.String("key-id", "", "key identifier without the \"ed25519:\" prefix (default: random)")
	if err := fs.Parse(args); err != nil {
		return err
	}

	path := *output
	if path == "" {
		cfg, err := loadConfig()
		if err != nil {
			return fmt.Errorf("%w\n(pass --output to generate a key without a configuration file)", err)
		}
		path = cfg.Global.PrivateKeyPath
	}

	key, err := signing.Generate(*keyID)
	if err != nil {
		return err
	}
	if err := key.Save(path); err != nil {
		return err
	}
	fmt.Printf("wrote %s\n  key ID:     %s\n  public key: %s\n", path, key.KeyID(), key.PublicBase64())
	return nil
}

func runVersion(_ context.Context, args []string) error {
	fs := newFlagSet("version")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if fs.NArg() > 0 {
		return errors.New("version takes no arguments")
	}
	fmt.Printf("warp %s\n", buildinfo.String())
	fmt.Printf("storage drivers: %v\n", storage.Drivers())
	return nil
}