From 2c7f8273e9e45e498d7214e660a5553acc3c1a75 Mon Sep 17 00:00:00 2001
From: Adnan Maolood <me@adnano.co>
Date: Mon, 15 Feb 2021 00:36:08 -0500
Subject: [PATCH] server: Recover from ServeGemini panics

---
 gemini.go |  6 ++++++
 server.go | 15 +++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/gemini.go b/gemini.go
index 0462dfb..9908bbb 100644
--- a/gemini.go
+++ b/gemini.go
@@ -19,4 +19,10 @@ var (
 	// ErrServerClosed is returned by the Server's Serve and ListenAndServe
 	// methods after a call to Shutdown or Close.
 	ErrServerClosed = errors.New("gemini: server closed")
+
+	// ErrAbortHandler is a sentinel panic value to abort a handler.
+	// While any panic from ServeGemini aborts the response to the client,
+	// panicking with ErrAbortHandler also suppresses logging of a stack
+	// trace to the server's error log.
+	ErrAbortHandler = errors.New("net/http: abort Handler")
 )
diff --git a/server.go b/server.go
index daeb097..38040f4 100644
--- a/server.go
+++ b/server.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"log"
 	"net"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -348,6 +349,15 @@ func (srv *Server) deleteConn(conn *net.Conn) {
 func (srv *Server) respond(conn net.Conn) {
 	defer conn.Close()
 
+	defer func() {
+		if err := recover(); err != nil && err != ErrAbortHandler {
+			const size = 64 << 10
+			buf := make([]byte, size)
+			buf = buf[:runtime.Stack(buf, false)]
+			srv.logf("gemini: panic serving %v: %v\n%s", conn.RemoteAddr(), err, buf)
+		}
+	}()
+
 	srv.trackConn(&conn)
 	defer srv.deleteConn(&conn)
 
@@ -359,11 +369,11 @@ func (srv *Server) respond(conn net.Conn) {
 	}
 
 	w := NewResponseWriter(conn)
-	defer w.Flush()
 
 	req, err := ReadRequest(conn)
 	if err != nil {
 		w.Status(StatusBadRequest)
+		w.Flush()
 		return
 	}
 
@@ -379,10 +389,12 @@ func (srv *Server) respond(conn net.Conn) {
 	resp := srv.responder(req)
 	if resp == nil {
 		w.Status(StatusNotFound)
+		w.Flush()
 		return
 	}
 
 	resp.ServeGemini(w, req)
+	w.Flush()
 }
 
 func (srv *Server) responder(r *Request) Handler {
@@ -418,7 +430,6 @@ func (srv *Server) logf(format string, args ...interface{}) {
 //
 // Handlers should not modify the provided Request.
 //
-// TODO:
 // If ServeGemini panics, the server (the caller of ServeGemini) assumes that
 // the effect of the panic was isolated to the active request. It recovers
 // the panic, logs a stack trace to the server error log, and closes the