From 9eae88f00c317bb1c221da8309e84c932ab4d7fa Mon Sep 17 00:00:00 2001
From: Adnan Maolood <me@adnano.co>
Date: Tue, 23 Feb 2021 21:59:16 -0500
Subject: [PATCH] Reimplement TimeoutHandler

---
 response.go |  8 ++++++
 timeout.go  | 79 ++++++-----------------------------------------------
 2 files changed, 17 insertions(+), 70 deletions(-)

diff --git a/response.go b/response.go
index 1d7a870..b939e5b 100644
--- a/response.go
+++ b/response.go
@@ -138,6 +138,14 @@ func NewResponseWriter(w io.WriteCloser) *ResponseWriter {
 	}
 }
 
+func (w *ResponseWriter) reset(wc io.WriteCloser) {
+	w.bw.Reset(wc)
+	*w = ResponseWriter{
+		bw: w.bw,
+		cl: wc,
+	}
+}
+
 // SetMediaType sets the media type that will be sent by Write for a
 // successful response. If no media type is set, a default of
 // "text/gemini; charset=utf-8" will be used.
diff --git a/timeout.go b/timeout.go
index 17c78e4..d879a84 100644
--- a/timeout.go
+++ b/timeout.go
@@ -1,11 +1,8 @@
-// +build ignore
-
 package gemini
 
 import (
 	"bytes"
 	"context"
-	"sync"
 	"time"
 )
 
@@ -27,84 +24,26 @@ type timeoutHandler struct {
 	dt time.Duration
 }
 
-func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
+func (t *timeoutHandler) ServeGemini(ctx context.Context, w *ResponseWriter, r *Request) {
 	ctx, cancel := context.WithTimeout(ctx, t.dt)
 	defer cancel()
 
+	conn := w.Hijack()
+
+	var b bytes.Buffer
+	w.reset(nopCloser{&b})
+
 	done := make(chan struct{})
-	tw := &timeoutWriter{}
 	go func() {
-		t.h.ServeGemini(ctx, tw, r)
+		t.h.ServeGemini(ctx, w, r)
 		close(done)
 	}()
 
 	select {
 	case <-done:
-		tw.mu.Lock()
-		defer tw.mu.Unlock()
-		if !tw.wroteHeader {
-			tw.status = StatusSuccess
-		}
-		w.WriteHeader(tw.status, tw.meta)
-		w.Write(tw.b.Bytes())
+		conn.Write(b.Bytes())
 	case <-ctx.Done():
-		tw.mu.Lock()
-		defer tw.mu.Unlock()
+		w.reset(conn)
 		w.WriteHeader(StatusTemporaryFailure, "Timeout")
-		tw.timedOut = true
 	}
 }
-
-type timeoutWriter struct {
-	mu          sync.Mutex
-	b           bytes.Buffer
-	status      Status
-	meta        string
-	mediatype   string
-	wroteHeader bool
-	timedOut    bool
-}
-
-func (w *timeoutWriter) SetMediaType(mediatype string) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	w.mediatype = mediatype
-}
-
-func (w *timeoutWriter) Write(b []byte) (int, error) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	if w.timedOut {
-		return 0, ErrHandlerTimeout
-	}
-	if !w.wroteHeader {
-		w.writeHeaderLocked(StatusSuccess, w.mediatype)
-	}
-	return w.b.Write(b)
-}
-
-func (w *timeoutWriter) WriteHeader(status Status, meta string) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	if w.timedOut {
-		return
-	}
-	w.writeHeaderLocked(status, meta)
-}
-
-func (w *timeoutWriter) writeHeaderLocked(status Status, meta string) {
-	if w.wroteHeader {
-		return
-	}
-	w.status = status
-	w.meta = meta
-	w.wroteHeader = true
-}
-
-func (w *timeoutWriter) Flush() error {
-	return nil
-}
-
-func (w *timeoutWriter) Close() error {
-	return nil
-}