diff --git a/gemini.go b/gemini.go index 9908bbb..caffb1b 100644 --- a/gemini.go +++ b/gemini.go @@ -24,5 +24,9 @@ var ( // 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") + ErrAbortHandler = errors.New("gemini: abort Handler") + + // ErrHandlerTimeout is returned on ResponseWriter Write calls + // in handlers which have timed out. + ErrHandlerTimeout = errors.New("gemini: Handler timeout") ) diff --git a/timeout.go b/timeout.go new file mode 100644 index 0000000..cff5307 --- /dev/null +++ b/timeout.go @@ -0,0 +1,115 @@ +package gemini + +import ( + "bytes" + "context" + "sync" + "time" +) + +// TimeoutHandler returns a Handler that runs h with the given time limit. +// +// The new Handler calls h.ServeGemini to handle each request, but +// if a call runs for longer than its time limit, the handler responds with a +// 40 Temporary Failure error. After such a timeout, writes by h to its +// ResponseWriter will return ErrHandlerTimeout. +// +// TimeoutHandler does not support the Hijacker or Flusher interfaces. +func TimeoutHandler(h Handler, dt time.Duration) Handler { + return &timeoutHandler{ + h: h, + dt: dt, + } +} + +type timeoutHandler struct { + h Handler + dt time.Duration +} + +func (t *timeoutHandler) ServeGemini(w ResponseWriter, r *Request) { + ctx := r.Context + if ctx == nil { + ctx = context.Background() + } + var cancelCtx func() + ctx, cancelCtx = context.WithTimeout(ctx, t.dt) + defer cancelCtx() + + done := make(chan struct{}) + tw := &timeoutWriter{} + panicChan := make(chan interface{}, 1) + go func() { + defer func() { + if p := recover(); p != nil { + panicChan <- p + } + }() + t.h.ServeGemini(tw, r) + close(done) + }() + + select { + case p := <-panicChan: + panic(p) + 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()) + case <-ctx.Done(): + tw.mu.Lock() + defer tw.mu.Unlock() + w.WriteHeader(StatusTemporaryFailure, "Timeout") + tw.timedOut = true + } +} + +type timeoutWriter struct { + mu sync.Mutex + b bytes.Buffer + status int + meta string + mediatype string + wroteHeader bool + timedOut bool +} + +func (w *timeoutWriter) MediaType(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 int, meta string) { + w.mu.Lock() + defer w.mu.Unlock() + if w.timedOut { + return + } + w.writeHeaderLocked(status, meta) +} + +func (w *timeoutWriter) writeHeaderLocked(status int, meta string) { + if w.wroteHeader { + return + } + w.status = status + w.meta = meta + w.wroteHeader = true +}