Use ffcli, tidy up usage message

cmd/hetty/config.go

@@ -0,0 +1,23 @@
+package main
+
+import (
+	"flag"
+
+	"go.uber.org/zap"
+)
+
+// Config represents the global configuration shared amongst all commands.
+type Config struct {
+	verbose  bool
+	jsonLogs bool
+	logger   *zap.Logger
+}
+
+// RegisterFlags registers the flag fields into the provided flag.FlagSet. This
+// helper function allows subcommands to register the root flags into their
+// flagsets, creating "global" flags that can be passed after any subcommand at
+// the commandline.
+func (cfg *Config) RegisterFlags(fs *flag.FlagSet) {
+	fs.BoolVar(&cfg.verbose, "verbose", false, "Enable verbose logging.")
+	fs.BoolVar(&cfg.jsonLogs, "json", false, "Encode logs as JSON, instead of pretty/human readable output.")
+}

cmd/hetty/hetty.go

@@ -0,0 +1,288 @@
+package main
+
+import (
+	"context"
+	"crypto/tls"
+	"embed"
+	"errors"
+	"flag"
+	"fmt"
+	"io/fs"
+	"net"
+	"net/http"
+	"os"
+	"os/exec"
+	"os/signal"
+	"strings"
+
+	"github.com/chromedp/chromedp"
+	badgerdb "github.com/dgraph-io/badger/v3"
+	"github.com/gorilla/mux"
+	"github.com/mitchellh/go-homedir"
+	"github.com/peterbourgon/ff/v3/ffcli"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+
+	"github.com/dstotijn/hetty/pkg/api"
+	"github.com/dstotijn/hetty/pkg/chrome"
+	"github.com/dstotijn/hetty/pkg/db/badger"
+	"github.com/dstotijn/hetty/pkg/proj"
+	"github.com/dstotijn/hetty/pkg/proxy"
+	"github.com/dstotijn/hetty/pkg/reqlog"
+	"github.com/dstotijn/hetty/pkg/scope"
+	"github.com/dstotijn/hetty/pkg/sender"
+)
+
+var version = "0.0.0"
+
+//go:embed admin
+//go:embed admin/_next/static
+//go:embed admin/_next/static/chunks/pages/*.js
+//go:embed admin/_next/static/*/*.js
+var adminContent embed.FS
+
+var hettyUsage = `
+Usage:
+    hetty [flags]
+
+Options:
+    --cert         Path to root CA certificate. Creates file if it doesn't exist. (Default: "~/.hetty/hetty_cert.pem")
+    --key          Path to root CA private key. Creates file if it doesn't exist. (Default: "~/.hetty/hetty_key.pem")
+    --db           Database directory path. (Default: "~/.hetty/db")
+    --addr         TCP address for HTTP server to listen on, in the form \"host:port\". (Default: ":8080")
+    --chrome       Launch Chrome with proxy settings applied and certificate errors ignored. (Default: false)
+    --verbose      Enable verbose logging.
+    --json         Encode logs as JSON, instead of pretty/human readable output.
+    --version, -v  Output version.
+    --help, -h     Output this usage text.
+
+Visit https://hetty.xyz to learn more about Hetty.
+`
+
+type HettyCommand struct {
+	config *Config
+
+	cert    string
+	key     string
+	db      string
+	addr    string
+	chrome  bool
+	version bool
+}
+
+func NewHettyCommand() (*ffcli.Command, *Config) {
+	cmd := HettyCommand{
+		config: &Config{},
+	}
+
+	fs := flag.NewFlagSet("hetty", flag.ExitOnError)
+
+	fs.StringVar(&cmd.cert, "cert", "~/.hetty/hetty_cert.pem",
+		"Path to root CA certificate. Creates a new certificate if file doesn't exist.")
+	fs.StringVar(&cmd.key, "key", "~/.hetty/hetty_key.pem",
+		"Path to root CA private key. Creates a new private key if file doesn't exist.")
+	fs.StringVar(&cmd.db, "db", "~/.hetty/db", "Database directory path.")
+	fs.StringVar(&cmd.addr, "addr", ":8080", "TCP address to listen on, in the form \"host:port\".")
+	fs.BoolVar(&cmd.chrome, "chrome", false, "Launch Chrome with proxy settings applied and certificate errors ignored.")
+	fs.BoolVar(&cmd.version, "version", false, "Output version.")
+	fs.BoolVar(&cmd.version, "v", false, "Output version.")
+
+	cmd.config.RegisterFlags(fs)
+
+	return &ffcli.Command{
+		Name:       "hetty",
+		ShortUsage: "hetty [-v] [subcommand] [flags] [<arg>...]",
+		FlagSet:    fs,
+		Exec:       cmd.Exec,
+		UsageFunc: func(*ffcli.Command) string {
+			return hettyUsage
+		},
+	}, cmd.config
+}
+
+func (cmd *HettyCommand) Exec(ctx context.Context, _ []string) error {
+	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
+	defer stop()
+
+	if cmd.version {
+		fmt.Fprint(os.Stdout, version+"\n")
+		return nil
+	}
+
+	mainLogger := cmd.config.logger.Named("main")
+
+	listenHost, listenPort, err := net.SplitHostPort(cmd.addr)
+	if err != nil {
+		mainLogger.Fatal("Failed to parse listening address.", zap.Error(err))
+	}
+
+	url := fmt.Sprintf("http://%v:%v", listenHost, listenPort)
+	if listenHost == "" || listenHost == "0.0.0.0" || listenHost == "127.0.0.1" || listenHost == "::1" {
+		url = fmt.Sprintf("http://localhost:%v", listenPort)
+	}
+
+	// Expand `~` in filepaths.
+	caCertFile, err := homedir.Expand(cmd.cert)
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to parse CA certificate filepath.", zap.Error(err))
+	}
+
+	caKeyFile, err := homedir.Expand(cmd.key)
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to parse CA private key filepath.", zap.Error(err))
+	}
+
+	dbPath, err := homedir.Expand(cmd.db)
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to parse database path.", zap.Error(err))
+	}
+
+	// Load existing CA certificate and key from disk, or generate and write
+	// to disk if no files exist yet.
+	caCert, caKey, err := proxy.LoadOrCreateCA(caKeyFile, caCertFile)
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to load or create CA key pair.", zap.Error(err))
+	}
+
+	// BadgerDB logs some verbose entries with `INFO` level, so unless
+	// we're running in debug mode, bump the minimal level to `WARN`.
+	dbLogger := cmd.config.logger.Named("badgerdb").WithOptions(zap.IncreaseLevel(zapcore.WarnLevel))
+
+	dbSugaredLogger := dbLogger.Sugar()
+
+	badger, err := badger.OpenDatabase(
+		badgerdb.DefaultOptions(dbPath).WithLogger(badger.NewLogger(dbSugaredLogger)),
+	)
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to open database.", zap.Error(err))
+	}
+	defer badger.Close()
+
+	scope := &scope.Scope{}
+
+	reqLogService := reqlog.NewService(reqlog.Config{
+		Scope:      scope,
+		Repository: badger,
+		Logger:     cmd.config.logger.Named("reqlog").Sugar(),
+	})
+
+	senderService := sender.NewService(sender.Config{
+		Repository:    badger,
+		ReqLogService: reqLogService,
+	})
+
+	projService, err := proj.NewService(proj.Config{
+		Repository:    badger,
+		ReqLogService: reqLogService,
+		SenderService: senderService,
+		Scope:         scope,
+	})
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to create new projects service.", zap.Error(err))
+	}
+
+	proxy, err := proxy.NewProxy(proxy.Config{
+		CACert: caCert,
+		CAKey:  caKey,
+		Logger: cmd.config.logger.Named("proxy").Sugar(),
+	})
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to create new proxy.", zap.Error(err))
+	}
+
+	proxy.UseRequestModifier(reqLogService.RequestModifier)
+	proxy.UseResponseModifier(reqLogService.ResponseModifier)
+
+	fsSub, err := fs.Sub(adminContent, "admin")
+	if err != nil {
+		cmd.config.logger.Fatal("Failed to construct file system subtree from admin dir.", zap.Error(err))
+	}
+
+	adminHandler := http.FileServer(http.FS(fsSub))
+	router := mux.NewRouter().SkipClean(true)
+	adminRouter := router.MatcherFunc(func(req *http.Request, match *mux.RouteMatch) bool {
+		hostname, _ := os.Hostname()
+		host, _, _ := net.SplitHostPort(req.Host)
+
+		// Serve local admin routes when either:
+		// - The `Host` is well-known, e.g. `hetty.proxy`, `localhost:[port]`
+		//   or the listen addr `[host]:[port]`.
+		// - The request is not for TLS proxying (e.g. no `CONNECT`) and not
+		//   for proxying an external URL. E.g. Request-Line (RFC 7230, Section 3.1.1)
+		//   has no scheme.
+		return strings.EqualFold(host, hostname) ||
+			req.Host == "hetty.proxy" ||
+			req.Host == fmt.Sprintf("%v:%v", "localhost", listenPort) ||
+			req.Host == fmt.Sprintf("%v:%v", listenHost, listenPort) ||
+			req.Method != http.MethodConnect && !strings.HasPrefix(req.RequestURI, "http://")
+	}).Subrouter().StrictSlash(true)
+
+	// GraphQL server.
+	gqlEndpoint := "/api/graphql/"
+	adminRouter.Path(gqlEndpoint).Handler(api.HTTPHandler(&api.Resolver{
+		ProjectService:    projService,
+		RequestLogService: reqLogService,
+		SenderService:     senderService,
+	}, gqlEndpoint))
+
+	// Admin interface.
+	adminRouter.PathPrefix("").Handler(adminHandler)
+
+	// Fallback (default) is the Proxy handler.
+	router.PathPrefix("").Handler(proxy)
+
+	httpServer := &http.Server{
+		Addr:         cmd.addr,
+		Handler:      router,
+		TLSNextProto: map[string]func(*http.Server, *tls.Conn, http.Handler){}, // Disable HTTP/2
+		ErrorLog:     zap.NewStdLog(cmd.config.logger.Named("http")),
+	}
+
+	go func() {
+		mainLogger.Info(fmt.Sprintf("Hetty (v%v) is running on %v ...", version, cmd.addr))
+		mainLogger.Info(fmt.Sprintf("\x1b[%dm%s\x1b[0m", uint8(32), "Get started at "+url))
+
+		err := httpServer.ListenAndServe()
+		if err != http.ErrServerClosed {
+			mainLogger.Fatal("HTTP server closed unexpected.", zap.Error(err))
+		}
+	}()
+
+	if cmd.chrome {
+		ctx, cancel := chrome.NewExecAllocator(ctx, chrome.Config{
+			ProxyServer:      url,
+			ProxyBypassHosts: []string{url},
+		})
+		defer cancel()
+
+		taskCtx, cancel := chromedp.NewContext(ctx)
+		defer cancel()
+
+		err = chromedp.Run(taskCtx, chromedp.Navigate(url))
+
+		switch {
+		case errors.Is(err, exec.ErrNotFound):
+			mainLogger.Info("Chrome executable not found.")
+		case err != nil:
+			mainLogger.Error(fmt.Sprintf("Failed to navigate to %v.", url), zap.Error(err))
+		default:
+			mainLogger.Info("Launched Chrome.")
+		}
+	}
+
+	// Wait for interrupt signal.
+	<-ctx.Done()
+	// Restore signal, allowing "force quit".
+	stop()
+
+	mainLogger.Info("Shutting down HTTP server. Press Ctrl+C to force quit.")
+
+	// Note: We expect httpServer.Handler to handle timeouts, thus, we don't
+	// need a context value with deadline here.
+	err = httpServer.Shutdown(context.Background())
+	if err != nil {
+		return fmt.Errorf("failed to shutdown HTTP server: %w", err)
+	}
+
+	return nil
+}