snice/main.go

package main

import (
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/lmittmann/tint"
	"github.com/urfave/cli/v2"
)

type (
	// struct for holding response details
	responseData struct {
		status int
		size   int
	}

	// our http.ResponseWriter implementation
	loggingResponseWriter struct {
		http.ResponseWriter // compose original http.ResponseWriter
		responseData        *responseData
	}
)

func (r *loggingResponseWriter) Write(b []byte) (int, error) {
	size, err := r.ResponseWriter.Write(b) // write response using original http.ResponseWriter
	r.responseData.size += size            // capture size
	return size, err
}

func (r *loggingResponseWriter) WriteHeader(statusCode int) {
	r.ResponseWriter.WriteHeader(statusCode) // write status code using original http.ResponseWriter
	r.responseData.status = statusCode       // capture status code
}

func WithLogging(h http.Handler) http.Handler {
	loggingFn := func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		remoteAddr := r.RemoteAddr
		fwdAddress := r.Header.Get("X-Forwarded-For")
		if fwdAddress != "" {
			remoteAddr = fwdAddress
			ips := strings.Split(fwdAddress, ", ")
			if len(ips) > 1 {
				remoteAddr = ips[0]
			}
		}
		remoteAddr, _, err := net.SplitHostPort(remoteAddr)
		if err != nil {
			slog.Error("Failed to format remoteAddr", slog.Any("err", err))
		}

		responseData := &responseData{
			status: 0,
			size:   0,
		}
		lrw := loggingResponseWriter{
			ResponseWriter: w, // compose original http.ResponseWriter
			responseData:   responseData,
		}
		h.ServeHTTP(&lrw, r) // inject our implementation of http.ResponseWriter

		duration := time.Since(start)

		// t=2023-10-27T18:08:47.231895532+13:00
		slog.Info("Request Completed:",
			slog.String("method", r.Method),
			slog.String("path", r.RequestURI),
			slog.String("url", r.URL.Path),
			slog.String("host", r.Host),
			slog.String("referer", r.Referer()),
			slog.String("remote_addr", remoteAddr),
			slog.Int("status", responseData.status),
			slog.String("duration", duration.String()),
			slog.Int("size", responseData.size),
		)

	}
	return http.HandlerFunc(loggingFn)
}

func pingHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte("OK"))
}

func main() {
	logLevel := &slog.LevelVar{} // INFO

	opts := tint.Options{
		Level: logLevel,
	}
	logLevel.Set(slog.LevelDebug)
	handler := tint.NewHandler(os.Stdout, &opts)
	logger := slog.New(handler)
	slog.SetDefault(logger)

	var host string
	var port string
	var directory string

	app := &cli.App{
		Name:                 "snice",
		Usage:                "Serve Static Files",
		Version:              "v0.2.0",
		EnableBashCompletion: true,
		Compiled:             time.Time{},
		Authors: []*cli.Author{
			{Name: "Hadley Rich", Email: "hads@nice.net.nz"},
		},
		SliceFlagSeparator:     "",
		UseShortOptionHandling: true,
		Suggest:                true,
		Commands: []*cli.Command{
			{
				Name:    "serve",
				Aliases: []string{"s"},
				Usage:   "Serve directory",
				Action: func(cCtx *cli.Context) error {
					var addr string = host + ":" + port

					mux := http.NewServeMux()
					mux.Handle("/ping", http.HandlerFunc(pingHandler))
					fileHandler := http.FileServer(http.Dir(directory))
					mux.Handle("/", fileHandler)

					srv := &http.Server{
						Addr:         addr,
						ReadTimeout:  5 * time.Second,
						WriteTimeout: 10 * time.Second,
						IdleTimeout:  120 * time.Second,
					}
					srv.Handler = WithLogging(mux)

					listener, err := net.Listen("tcp", addr)
					if err != nil {
						slog.Error("Listen error", err)
					}
					slog.Info("Starting server",
						slog.String("dir", directory),
						slog.String("addr", fmt.Sprintf(":%s", listener.Addr())),
					)
					err = srv.Serve(listener)
					if err != nil {
						slog.Error("Serve error", err)
					}
					return nil
				},
				Flags: []cli.Flag{
					&cli.StringFlag{
						Name:        "directory",
						Aliases:     []string{"dir", "d"},
						EnvVars:     []string{"DIRECTORY"},
						Value:       "/srv",
						Usage:       "Directory to serve",
						Destination: &directory,
					},
					&cli.StringFlag{
						Name:        "host",
						EnvVars:     []string{"HOST"},
						Value:       "0.0.0.0",
						Usage:       "Host to listen on",
						Destination: &host,
					}},
			},
			{
				Name:    "healthcheck",
				Aliases: []string{"hc"},
				Usage:   "Call healthcheck endpoint",
				Action: func(cCtx *cli.Context) error {
					url := fmt.Sprintf("http://127.0.0.1:%s/ping", port)
					slog.Debug("Healthcheck: ", slog.String("url", url))
					client := http.Client{
						Timeout: 1 * time.Second,
					}
					res, err := client.Get(url)
					if err != nil {
						return cli.Exit("FAIL", 1)
					}
					if res.StatusCode == 200 {
						return cli.Exit("OK", 0)
					}
					slog.Debug(fmt.Sprintf("Status: %d\n", res.StatusCode))
					return cli.Exit("FAIL", 1)
				}}},
		Flags: []cli.Flag{
			&cli.BoolFlag{Name: "quiet", Aliases: []string{"q"}},
			&cli.StringFlag{
				Name:        "port",
				Aliases:     []string{"p"},
				EnvVars:     []string{"PORT"},
				Value:       "3000",
				Usage:       "Port to serve on",
				Destination: &port,
			},
		},
	}

	if err := app.Run(os.Args); err != nil {
		os.Exit(1)
	}
}