Add some middleware

cmd/server/server.go

@@ -57,6 +57,7 @@ func ActionServe(c *cli.Context) error {
 	// Setup loggers
 	rootLogger := getRootLogger(cfg.LogLevel)
 	serverLogger := rootLogger.Named("SERV")
+	accessLogger := rootLogger.Named("ACCS")
 
 	// Setup contexts for clean shutdown
 	rootCtx, rootCancel := signal.NotifyContext(context.Background(), os.Interrupt)
@@ -70,6 +71,7 @@ func ActionServe(c *cli.Context) error {
 		srv := gpaste.NewHTTPServer(cfg)
 		srv.Addr = cfg.ListenAddr
 		srv.Logger = serverLogger
+		srv.AccessLogger = accessLogger
 
 		// Wait for cancel
 		go func() {

filestore_fs_test.go

@@ -14,4 +14,13 @@ func TestFSFileStore(t *testing.T) {
 	}
 
 	RunFilestoreTest(s, t)
+	persistentDir := t.TempDir()
+	newFunc := func() gpaste.FileStore {
+		s, err := gpaste.NewFSFileStore(persistentDir)
+		if err != nil {
+			t.Fatalf("Error creating store: %s", err)
+		}
+		return s
+	}
+	RunPersistentFilestoreTest(newFunc, t)
 }

filestore_test.go

@@ -3,9 +3,12 @@ package gpaste_test
 import (
 	"bytes"
 	"io"
+	"strings"
 	"testing"
+	"time"
 
 	"git.t-juice.club/torjus/gpaste"
+	"github.com/google/go-cmp/cmp"
 	"github.com/google/uuid"
 )
 
@@ -74,3 +77,81 @@ func RunFilestoreTest(s gpaste.FileStore, t *testing.T) {
 		}
 	})
 }
+
+func RunPersistentFilestoreTest(newStoreFunc func() gpaste.FileStore, t *testing.T) {
+	s := newStoreFunc()
+
+	files := []struct {
+		File         *gpaste.File
+		ExpectedData string
+	}{
+		{
+			File: &gpaste.File{
+				ID:               uuid.NewString(),
+				OriginalFilename: "testfile.txt",
+				MaxViews:         5,
+				ExpiresOn:        time.Now().Add(10 * time.Minute),
+				Body:             io.NopCloser(strings.NewReader("cocks!")),
+			},
+			ExpectedData: "cocks!",
+		},
+		{
+			File: &gpaste.File{
+				ID:               uuid.NewString(),
+				OriginalFilename: "testfile2.txt",
+				MaxViews:         5,
+				ExpiresOn:        time.Now().Add(10 * time.Minute),
+				Body:             io.NopCloser(strings.NewReader("derps!")),
+			},
+			ExpectedData: "derps!",
+		},
+	}
+
+	for _, f := range files {
+		err := s.Store(f.File)
+		if err != nil {
+			t.Fatalf("Error storing file: %s", err)
+		}
+	}
+	for _, f := range files {
+		retrieved, err := s.Get(f.File.ID)
+		if err != nil {
+			t.Fatalf("Unable to retrieve file: %s", err)
+		}
+
+		ignoreBody := cmp.FilterPath(func(p cmp.Path) bool { return p.String() == "Body" }, cmp.Ignore())
+		if !cmp.Equal(retrieved, f.File, ignoreBody) {
+			t.Errorf("Mismatch: %s", cmp.Diff(retrieved, f.File))
+		}
+		buf := new(strings.Builder)
+		if _, err := io.Copy(buf, retrieved.Body); err != nil {
+			t.Fatalf("Error reading from body: %s", err)
+		}
+		retrieved.Body.Close()
+		if buf.String() != f.ExpectedData {
+			t.Fatalf("Data does not match. %s", cmp.Diff(buf.String(), f.ExpectedData))
+		}
+	}
+
+	// Reopen store, and fetch again
+	s = newStoreFunc()
+	for _, f := range files {
+		retrieved, err := s.Get(f.File.ID)
+		if err != nil {
+			t.Fatalf("Unable to retrieve file: %s", err)
+		}
+
+		ignoreBody := cmp.FilterPath(func(p cmp.Path) bool { return p.String() == "Body" }, cmp.Ignore())
+		if !cmp.Equal(retrieved, f.File, ignoreBody) {
+			t.Errorf("Mismatch: %s", cmp.Diff(retrieved, f.File))
+		}
+		buf := new(strings.Builder)
+		if _, err := io.Copy(buf, retrieved.Body); err != nil {
+			t.Fatalf("Error reading from body: %s", err)
+		}
+		retrieved.Body.Close()
+		if buf.String() != f.ExpectedData {
+			t.Fatalf("Data does not match. %s", cmp.Diff(buf.String(), f.ExpectedData))
+		}
+	}
+}

http.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/go-chi/chi/v5"
+	"github.com/go-chi/chi/v5/middleware"
 	"github.com/google/uuid"
 	"go.uber.org/zap"
 )
@@ -15,6 +16,7 @@ type HTTPServer struct {
 	store        FileStore
 	config       *ServerConfig
 	Logger       *zap.SugaredLogger
+	AccessLogger *zap.SugaredLogger
 	http.Server
 }
 
@@ -22,10 +24,14 @@ func NewHTTPServer(cfg *ServerConfig) *HTTPServer {
 	srv := &HTTPServer{
 		config:       cfg,
 		Logger:       zap.NewNop().Sugar(),
+		AccessLogger: zap.NewNop().Sugar(),
 	}
 	srv.store = NewMemoryFileStore()
 
 	r := chi.NewRouter()
+	r.Use(middleware.RealIP)
+	r.Use(middleware.RequestID)
+	r.Use(srv.MiddlewareAccessLogger)
 	r.Get("/", srv.HandlerIndex)
 	r.Post("/api/file", srv.HandlerAPIFilePost)
 	r.Get("/api/file/{id}", srv.HandlerAPIFileGet)
@@ -43,6 +49,7 @@ func (s *HTTPServer) HandlerAPIFilePost(w http.ResponseWriter, r *http.Request)
 		ID:   uuid.Must(uuid.NewRandom()).String(),
 		Body: r.Body,
 	}
+	reqID := middleware.GetReqID(r.Context())
 
 	// Check if multipart form
 	ct := r.Header.Get("Content-Type")
@@ -53,10 +60,10 @@ func (s *HTTPServer) HandlerAPIFilePost(w http.ResponseWriter, r *http.Request)
 	err := s.store.Store(f)
 	if err != nil {
 		w.WriteHeader(http.StatusInternalServerError)
-		s.Logger.Warnw("Error storing file.", "erorr", err, "id", f.ID, "remote_addr", r.RemoteAddr)
+		s.Logger.Warnw("Error storing file.", "req_id", reqID, "error", err, "id", f.ID, "remote_addr", r.RemoteAddr)
 		return
 	}
-	s.Logger.Infow("Stored file.", "id", f.ID, "remote_addr", r.RemoteAddr)
+	s.Logger.Infow("Stored file.", "req_id", reqID, "id", f.ID, "remote_addr", r.RemoteAddr)
 	var resp = struct {
 		Message string `json:"message"`
 		ID      string `json:"id"`
@@ -69,7 +76,7 @@ func (s *HTTPServer) HandlerAPIFilePost(w http.ResponseWriter, r *http.Request)
 	w.WriteHeader(http.StatusAccepted)
 	encoder := json.NewEncoder(w)
 	if err := encoder.Encode(&resp); err != nil {
-		s.Logger.Warnw("Error encoding response to client.", "error", err, "remote_addr", r.RemoteAddr)
+		s.Logger.Warnw("Error encoding response to client.", "req_id", reqID, "error", err, "remote_addr", r.RemoteAddr)
 	}
 }
 
@@ -89,12 +96,13 @@ func (s *HTTPServer) HandlerAPIFileGet(w http.ResponseWriter, r *http.Request) {
 
 	w.WriteHeader(http.StatusOK)
 	if _, err := io.Copy(w, f.Body); err != nil {
-		s.Logger.Warnw("Error writing file to client.", "error", err, "remote_addr", r.RemoteAddr)
+		reqID := middleware.GetReqID(r.Context())
+		s.Logger.Warnw("Error writing file to client.", "req_id", reqID, "error", err, "remote_addr", r.RemoteAddr)
 	}
 }
 
 func (s *HTTPServer) processMultiPartFormUpload(w http.ResponseWriter, r *http.Request) {
-	s.Logger.Debugw("Processing multipart form.")
+	reqID := middleware.GetReqID(r.Context())
 	type resp struct {
 		Message string `json:"message"`
 		ID      string `json:"id"`
@@ -104,12 +112,12 @@ func (s *HTTPServer) processMultiPartFormUpload(w http.ResponseWriter, r *http.R
 	var responses []resp
 
 	if err := r.ParseMultipartForm(1024 * 1024 * 10); err != nil {
-		s.Logger.Warnw("Error parsing multipart form.", "err", err)
+		s.Logger.Warnw("Error parsing multipart form.", "req_id", reqID, "err", err)
 	}
 	for k := range r.MultipartForm.File {
 		ff, fh, err := r.FormFile(k)
 		if err != nil {
-			s.Logger.Warnw("Error reading file from multipart form.", "error", err)
+			s.Logger.Warnw("Error reading file from multipart form.", "req_id", reqID, "error", err)
 			return
 		}
 		f := &File{
@@ -120,10 +128,10 @@ func (s *HTTPServer) processMultiPartFormUpload(w http.ResponseWriter, r *http.R
 
 		if err := s.store.Store(f); err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
-			s.Logger.Warnw("Error storing file.", "erorr", err, "id", f.ID, "remote_addr", r.RemoteAddr)
+			s.Logger.Warnw("Error storing file.", "req_id", reqID, "error", err, "id", f.ID, "remote_addr", r.RemoteAddr)
 			return
 		}
-		s.Logger.Infow("Stored file.", "id", f.ID, "filename", f.OriginalFilename, "remote_addr", r.RemoteAddr)
+		s.Logger.Infow("Stored file.", "req_id", reqID, "id", f.ID, "filename", f.OriginalFilename, "remote_addr", r.RemoteAddr)
 
 		responses = append(responses, resp{Message: "OK", ID: f.ID, URL: "TODO"})
 
@@ -132,6 +140,6 @@ func (s *HTTPServer) processMultiPartFormUpload(w http.ResponseWriter, r *http.R
 	w.WriteHeader(http.StatusAccepted)
 	encoder := json.NewEncoder(w)
 	if err := encoder.Encode(&responses); err != nil {
-		s.Logger.Warnw("Error encoding response to client.", "error", err, "remote_addr", r.RemoteAddr)
+		s.Logger.Warnw("Error encoding response to client.", "req_id", reqID, "error", err, "remote_addr", r.RemoteAddr)
 	}
 }

middleware.go

@@ -0,0 +1,30 @@
+package gpaste
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/go-chi/chi/v5/middleware"
+)
+
+func (s *HTTPServer) MiddlewareAccessLogger(next http.Handler) http.Handler {
+	fn := func(w http.ResponseWriter, r *http.Request) {
+		ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
+		t1 := time.Now()
+
+		reqID := middleware.GetReqID(r.Context())
+
+		defer func() {
+			s.AccessLogger.Infow(r.Method,
+				"path", r.URL.Path,
+				"status", ww.Status(),
+				"written", ww.BytesWritten(),
+				"remote_addr", r.RemoteAddr,
+				"processing_time_ms", time.Since(t1).Milliseconds(),
+				"req_id", reqID)
+		}()
+
+		next.ServeHTTP(ww, r)
+	}
+	return http.HandlerFunc(fn)
+}