Fix server

cmd/restic/cmd_serve.go

@@ -2,7 +2,10 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"net"
 	"net/http"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -35,6 +38,8 @@ func init() {
 	cmdFlags.StringVarP(&serveOptions.Listen, "listen", "l", "localhost:3080", "set the listen host name and `address`")
 }
 
+const serverShutdownTimeout = 30 * time.Second
+
 func runWebServer(ctx context.Context, opts ServeOptions, gopts GlobalOptions, args []string) error {
 	if len(args) > 0 {
 		return errors.Fatal("this command does not accept additional arguments")
@@ -57,10 +62,42 @@ func runWebServer(ctx context.Context, opts ServeOptions, gopts GlobalOptions, a
 		return err
 	}
 
-	srv := server.New(repo, snapshotLister, TimeFormat)
+	srv := http.Server{
+		BaseContext: func(l net.Listener) context.Context {
+			// just return the global context
+			return ctx
+		},
+		Handler: server.New(repo, snapshotLister, TimeFormat),
+	}
+
+	listener, err := net.Listen("tcp", opts.Listen)
+	if err != nil {
+		return fmt.Errorf("start listener: %v", err)
+	}
+
+	// wait until context is cancelled, then close listener
+	go func() {
+		<-ctx.Done()
+		Printf("gracefully shutting down server\n")
+
+		ctxTimeout, cancel := context.WithTimeout(context.Background(), serverShutdownTimeout)
+		defer cancel()
+
+		_ = srv.Shutdown(ctxTimeout)
+	}()
 
 	Printf("Now serving the repository at http://%s\n", opts.Listen)
 	Printf("When finished, quit with Ctrl-c here.\n")
 
-	return http.ListenAndServe(opts.Listen, srv)
+	err = srv.Serve(listener)
+
+	if errors.Is(err, http.ErrServerClosed) {
+		err = nil
+	}
+
+	if err != nil {
+		return fmt.Errorf("serve: %v", err)
+	}
+
+	return nil
 }

internal/server/server.go

@@ -4,7 +4,6 @@ package server
 import (
 	"context"
 	"fmt"
-	"io/fs"
 	"net/http"
 	"sort"
 	"strings"
@@ -109,11 +108,12 @@ func New(repo restic.Repository, snapshotLister restic.Lister, timeFormat string
 		}
 	})
 
-	http.HandleFunc("/", func(rw http.ResponseWriter, req *http.Request) {
+	mux.HandleFunc("/", func(rw http.ResponseWriter, req *http.Request) {
 		if req.URL.Path != "/" {
 			http.NotFound(rw, req)
 			return
 		}
+
 		var rows []indexPageRow
 		for sn := range findFilteredSnapshots(req.Context(), snapshotLister, repo, &restic.SnapshotFilter{}, nil) {
 			rows = append(rows, indexPageRow{
@@ -125,23 +125,29 @@ func New(repo restic.Repository, snapshotLister restic.Lister, timeFormat string
 				Paths: sn.Paths,
 			})
 		}
+
 		sort.Slice(rows, func(i, j int) bool {
 			return rows[i].Time.After(rows[j].Time)
 		})
+
 		if err := indexPage.Execute(rw, indexPageData{"Snapshots", rows}); err != nil {
 			http.Error(rw, err.Error(), http.StatusInternalServerError)
 		}
 	})
 
-	http.HandleFunc("/style.css", func(rw http.ResponseWriter, _ *http.Request) {
-		rw.Header().Set("Cache-Control", "max-age=300")
-		buf, err := fs.ReadFile(assets.FS, "style.css")
-		if err == nil {
+	mux.HandleFunc("/style.css", func(rw http.ResponseWriter, req *http.Request) {
+		buf, err := assets.FS.ReadFile("style.css")
+		if err != nil {
 			rw.WriteHeader(http.StatusInternalServerError)
-			fmt.Fprintf(rw, "error: %v", err)
+
+			fmt.Fprintf(rw, "error reading embedded style.css: %v\n", err)
+
 			return
 		}
 
+		rw.Header().Set("Cache-Control", "max-age=300")
+		rw.Header().Set("Content-Type", "text/css")
+
 		_, _ = rw.Write(buf)
 	})