From 991b18d5264c263a53788b61ec56609b605e596d Mon Sep 17 00:00:00 2001
From: Adnan Maolood <me@adnano.co>
Date: Wed, 24 Feb 2021 10:25:08 -0500
Subject: [PATCH] examples/server: Add logging middleware

---
 examples/server.go | 52 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/examples/server.go b/examples/server.go
index 7882d11..2f6ac88 100644
--- a/examples/server.go
+++ b/examples/server.go
@@ -26,7 +26,7 @@ func main() {
 	mux.Handle("/", gemini.FileServer(os.DirFS("/var/www")))
 
 	server := &gemini.Server{
-		Handler:        mux,
+		Handler:        logMiddleware(mux),
 		ReadTimeout:    30 * time.Second,
 		WriteTimeout:   1 * time.Minute,
 		GetCertificate: certificates.Get,
@@ -55,3 +55,53 @@ func main() {
 		}
 	}
 }
+
+func logMiddleware(h gemini.Handler) gemini.Handler {
+	return gemini.HandlerFunc(func(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
+		lw := &logResponseWriter{rw: w}
+		h.ServeGemini(ctx, lw, r)
+		host := r.TLS().ServerName
+		log.Printf("gemini: %s %q %d %d", host, r.URL, lw.status, lw.wrote)
+	})
+}
+
+type logResponseWriter struct {
+	rw          gemini.ResponseWriter
+	status      gemini.Status
+	meta        string
+	mediatype   string
+	wroteHeader bool
+	wrote       int
+}
+
+func (w *logResponseWriter) SetMediaType(mediatype string) {
+	w.mediatype = mediatype
+}
+
+func (w *logResponseWriter) Write(b []byte) (int, error) {
+	if !w.wroteHeader {
+		w.WriteHeader(gemini.StatusSuccess, w.mediatype)
+	}
+	n, err := w.rw.Write(b)
+	w.wrote += n
+	return n, err
+}
+
+func (w *logResponseWriter) WriteHeader(status gemini.Status, meta string) {
+	if w.wroteHeader {
+		return
+	}
+	w.status = status
+	w.meta = meta
+	w.wroteHeader = true
+	w.rw.WriteHeader(status, meta)
+	w.wrote += len(meta) + 5
+}
+
+func (w *logResponseWriter) Flush() error {
+	return nil
+}
+
+func (w *logResponseWriter) Close() error {
+	return nil
+}