Add message argument to TimeoutHandler

This commit is contained in:
Adnan Maolood 2021-02-28 22:07:24 -05:00
parent ad5d78f08f
commit ae7d58549d
1 changed files with 11 additions and 8 deletions

View File

@ -79,18 +79,21 @@ func StripPrefix(prefix string, h Handler) Handler {
// //
// The new Handler calls h.ServeGemini to handle each request, but // 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 // 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 // 40 Temporary Failure status code and the given message in its response meta.
// its ResponseWriter will return context.DeadlineExceeded. // After such a timeout, writes by h to its ResponseWriter will return
func TimeoutHandler(h Handler, dt time.Duration) Handler { // context.DeadlineExceeded.
func TimeoutHandler(h Handler, dt time.Duration, message string) Handler {
return &timeoutHandler{ return &timeoutHandler{
h: h, h: h,
dt: dt, dt: dt,
msg: message,
} }
} }
type timeoutHandler struct { type timeoutHandler struct {
h Handler h Handler
dt time.Duration dt time.Duration
msg string
} }
func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) { func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
@ -118,7 +121,7 @@ func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *R
w.WriteHeader(tw.status, tw.meta) w.WriteHeader(tw.status, tw.meta)
w.Write(buf.Bytes()) w.Write(buf.Bytes())
case <-ctx.Done(): case <-ctx.Done():
w.WriteHeader(StatusTemporaryFailure, "Timeout") w.WriteHeader(StatusTemporaryFailure, t.msg)
} }
} }