mux/router_serve.go

package mux

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
)

type ServeCB func(srv *http.Server) error

// Serve with graceful shutdown
func (r *Router) Serve(cb ServeCB) {
	// catch all options
	// lets get it thorugh all middlewares
	r.mux.Handle("OPTIONS /", optionsHandler{})

	srv := &http.Server{
		Handler: r,
	}

	idleConnsClosed := make(chan struct{})
	go func() {
		sigint := make(chan os.Signal, 1)
		signal.Notify(sigint, os.Interrupt)
		<-sigint

		// We received an interrupt signal, shut down.
		if err := srv.Shutdown(context.Background()); err != nil {
			// Error from closing listeners, or context timeout:
			slog.Error("server shutdown error", "error", err)
		} else {
			slog.Info("server shutdown")
		}
		close(idleConnsClosed)
	}()

	if err := cb(srv); !errors.Is(err, http.ErrServerClosed) {
		// Error starting or closing listener:
		slog.Error("start server error", "error", err)
	}

	<-idleConnsClosed
}

type optionsHandler struct{}

func (optionsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Length", "0")
	if r.ContentLength != 0 {
		// Read up to 4KB of OPTIONS body (as mentioned in the
		// spec as being reserved for future use), but anything
		// over that is considered a waste of server resources
		// (or an attack) and we abort and close the connection,
		// courtesy of MaxBytesReader's EOF behavior.
		mb := http.MaxBytesReader(w, r.Body, 4<<10)
		io.Copy(io.Discard, mb)
	}
}