gocustomurls/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var errorLog *log.Logger = log.New(os.Stderr, "", log.LstdFlags)

const (
	DEFAULT_RULES_FILE string = "/var/lib/gocustomurls/rules.json"
	DEFAULT_LOG_FILE   string = "/var/log/gocustomurls/app.log"
)

// flagsSet returns a set of all the flags what were actually set on the
// command line.
func flagsSet(flags *flag.FlagSet) map[string]bool {
	s := make(map[string]bool)
	flags.Visit(func(f *flag.Flag) {
		s[f.Name] = true
	})
	return s
}

func main() {
	programName := os.Args[0]
	// errorLog = log.New(os.Stderr, "", log.LstdFlags)

	flags := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	flags.Usage = func() {
		out := flags.Output()
		fmt.Fprintf(out, "Usage: %v [flags]\n\n", programName)
		fmt.Fprint(out, "  This utility serves vanity urls for the go get/install command.\n")
		fmt.Fprint(out, "  By default, the server listens on localhost:7070.\n")
		flags.PrintDefaults()
	}

	portFlag := flags.String("port", "7070", "port to listen to")
	rulesFileFlag := flags.String("config", DEFAULT_RULES_FILE, "contains go-import mapping")
	logFileFlag := flags.String("logfile", DEFAULT_LOG_FILE, "default log file")
	flags.Parse(os.Args[1:])

	if len(flags.Args()) > 1 {
		errorLog.Println("Error: too many command-line arguments")
		flags.Usage()
		os.Exit(1)
	}

	allSetFlags := flagsSet(flags)

	var port string
	if allSetFlags["port"] {
		port = *portFlag
	}

	var rulesFile string
	if allSetFlags["config"] {
		rulesFile = *rulesFileFlag
	}

	var logFile string
	if allSetFlags["logFile"] {
		logFile = *logFileFlag
	}

	// load rules mapping
	c := &Config{}
	err := c.LoadMappingFile(rulesFile)
	if err != nil {
		errorLog.Println(err)
		os.Exit(1)
	}
	l, err := newFileLogger(logFile)
	if err != nil {
		errorLog.Println(err)
		os.Exit(1)
	}

	app := &Application{
		Config: c,
		Log:    l,
	}
	srv := app.Setup(port)

	// For graceful shutdowns
	go func() {
		err := srv.ListenAndServe()
		if !errors.Is(err, http.ErrServerClosed) {
			errorLog.Printf("HTTP Server error: %+v\n", err)
			os.Exit(1)
		}
		app.Log.logger.Println("Stopped serving new connections.")
	}()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
	<-sigChan

	shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), 10*time.Second)
	defer shutdownRelease()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		errorLog.Printf("HTTP shutdown error: %+v\n", err)
		os.Exit(1)
	}
	app.Log.logger.Println("Graceful shutdown complete.")
}