cabinet/main.go

package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/peterbourgon/ff/v3"
	"github.com/peterbourgon/ff/v3/ffcli"
	"github.com/peterbourgon/ff/v3/fftoml"
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
	"go.jolheiser.com/cabinet/workspace"
)

func main() {
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
	zerolog.SetGlobalLevel(zerolog.InfoLevel)

	serveOpts := serveOpts{
		maxFileSize:   1024 * 1024 * 10,       // 10 MiB
		maxDiskSize:   1024 * 1024 * 1024 * 5, // 5 GiB
		sizePerMinute: 5 * 1024 * 1024,        // 5 MiB/min
		burstSize:     100 * 1024 * 1024,      // 100 MiB
		memPerRequest: 1024 * 1024,            // 1 MiB
	}
	serveFS := flag.NewFlagSet("serve", flag.ExitOnError)
	serveFS.BoolVar(&serveOpts.jsonMode, "json", false, "Log as JSON")
	serveFS.BoolVar(&serveOpts.debugMode, "debug", false, "Debug logging")
	serveFS.Func("max-file-size", "Max size of a file", fileSizeParse(&serveOpts.maxFileSize))
	serveFS.Func("max-disk-size", "Max size of all disk space usage", fileSizeParse(&serveOpts.maxDiskSize))
	serveFS.StringVar(&serveOpts.workspacePath, "workspace", ".cabinet", "Workspace for DB, files, tokens, etc.")
	serveFS.DurationVar(&serveOpts.gcInterval, "gc-interval", time.Hour, "GC interval for cleaning up files")
	serveFS.IntVar(&serveOpts.requestPerMinute, "request-per-minute", 12, "Request limit per-second")
	serveFS.Func("size-per-second", "File size limit per-second", fileSizeParse(&serveOpts.sizePerMinute))
	serveFS.Func("burst-size", "Burst size for files", fileSizeParse(&serveOpts.burstSize))
	serveFS.Func("mem-per-request", "Memory per request before storing on disk temporarily", fileSizeParse(&serveOpts.memPerRequest))
	serveFS.IntVar(&serveOpts.port, "port", 8080, "Port to serve on")
	serveFS.StringVar(&serveOpts.domain, "domain", "", "Domain the app is running on")

	var tokenOpts tokenOpts
	tokenFS := flag.NewFlagSet("token", flag.ExitOnError)
	tokenFS.StringVar(&tokenOpts.workspacePath, "workspace", ".cabinet", "Workspace for DB, files, tokens, etc.")
	tokenFS.Func("token", "Set token to bypass prompt (or empty to generate)", func(s string) error {
		if s == "" {
			var err error
			tokenOpts.token, err = randToken()
			return err
		}
		tokenOpts.token = s
		return nil
	})
	tokenFS.Func("permission", "Set the token permission to bypass prompt (or empty for all)", tokenPermParse(&tokenOpts.perm))
	tokenFS.StringVar(&tokenOpts.desc, "description", "", "Description for identifying the token usage/user")
	tokenFS.BoolVar(&tokenOpts.delete, "delete", false, "Delete a token by --token or prompt")

	cmd := ffcli.Command{
		Name:       "cabinet",
		ShortUsage: "Cabinet hosts your files and redirects",
		ShortHelp:  "Base command",
		Subcommands: []*ffcli.Command{
			{
				Name:       "serve",
				ShortUsage: "Serve the Cabinet server",
				ShortHelp:  "Listen for file and redirect requests",
				FlagSet:    serveFS,
				Options: []ff.Option{
					ff.WithEnvVarPrefix("CABINET"),
					ff.WithConfigFileFlag("config"),
					ff.WithAllowMissingConfigFile(true),
					ff.WithConfigFileParser(fftoml.New().Parse),
				},
				Exec: serveCmd(&serveOpts),
			},
			{
				Name:      "token",
				ShortHelp: "Generate or delete tokens",
				FlagSet:   tokenFS,
				Exec:      tokenCmd(&tokenOpts),
			},
		},
		Exec: func(_ context.Context, _ []string) error {
			return errors.New("cabinet [serve|token]")
		},
	}

	if err := cmd.ParseAndRun(context.Background(), os.Args[1:]); err != nil {
		log.Fatal().Err(err).Msg("could not run command")
	}
}

var fileSizePattern = regexp.MustCompile(`(?i)(\d+)(b|mi?b?|gi?b?)?`)

func fileSizeParse(out *int) func(string) error {
	return func(in string) error {
		match := fileSizePattern.FindStringSubmatch(in)
		if match == nil {
			return fmt.Errorf("pattern %q did not match a file size pattern", in)
		}
		b, err := strconv.Atoi(match[1])
		if err != nil {
			return fmt.Errorf("could not parse file size: %w", err)
		}
		switch strings.ToLower(match[2]) {
		case "", "b":
			*out = b
		case "m", "mb":
			*out = b * 1000 * 1000
		case "mi", "mib":
			*out = b * 1024 * 1024
		case "g", "gb":
			*out = b * 1000 * 1000 * 1000
		case "gi", "gib":
			*out = b * 1024 * 1024 * 1024
		default:
			return fmt.Errorf("pattern %q did not match a file size pattern", in)
		}
		return nil
	}
}

func tokenPermParse(out *workspace.TokenPermission) func(string) error {
	return func(in string) error {
		if in == "" {
			*out = workspace.TokenRedirect | workspace.TokenFile
		}
		var err error
		*out, err = workspace.ParseTokenPermission(in)
		return err
	}
}

func randToken() (string, error) {
	b := make([]byte, 10)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}