diff --git a/handler.go b/handler.go index e695c11..5c7d9aa 100644 --- a/handler.go +++ b/handler.go @@ -3,6 +3,7 @@ package gemini import ( "bytes" "context" + "io" "net/url" "strings" "time" @@ -96,22 +97,64 @@ func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *R ctx, cancel := context.WithTimeout(ctx, t.dt) defer cancel() - conn := w.Hijack() - - var b bytes.Buffer - w.reset(nopCloser{&b}) + buf := &bytes.Buffer{} + tw := &timeoutWriter{ + wc: &contextWriter{ + ctx: ctx, + cancel: cancel, + done: ctx.Done(), + wc: nopCloser{buf}, + }, + } done := make(chan struct{}) go func() { - t.h.ServeGemini(ctx, w, r) + t.h.ServeGemini(ctx, tw, r) close(done) }() select { case <-done: - conn.Write(b.Bytes()) + w.WriteHeader(tw.status, tw.meta) + w.Write(buf.Bytes()) case <-ctx.Done(): - w.reset(conn) w.WriteHeader(StatusTemporaryFailure, "Timeout") } } + +type timeoutWriter struct { + ResponseWriter + wc io.WriteCloser + status Status + meta string + mediatype string + wroteHeader bool +} + +func (w *timeoutWriter) SetMediaType(mediatype string) { + w.mediatype = mediatype +} + +func (w *timeoutWriter) Write(b []byte) (int, error) { + if !w.wroteHeader { + w.WriteHeader(StatusSuccess, w.mediatype) + } + return w.wc.Write(b) +} + +func (w *timeoutWriter) WriteHeader(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 w.wc.Close() +} diff --git a/response.go b/response.go index 6490579..2ab3bfd 100644 --- a/response.go +++ b/response.go @@ -175,8 +175,6 @@ type ResponseWriter interface { // or clear those deadlines as needed. Hijack() net.Conn - reset(io.WriteCloser) - // unexported method so we can extend this interface over time // without breaking existing code. Implementers must embed a concrete // type from elsewhere. @@ -200,14 +198,6 @@ func newResponseWriter(w io.WriteCloser) *responseWriter { } } -func (w *responseWriter) reset(wc io.WriteCloser) { - w.bw.Reset(wc) - *w = responseWriter{ - bw: w.bw, - cl: wc, - } -} - func (w *responseWriter) SetMediaType(mediatype string) { w.mediatype = mediatype }