diff --git a/context/http.go b/context/http.go index 1808c6b44..98ab436d9 100644 --- a/context/http.go +++ b/context/http.go @@ -15,7 +15,8 @@ import ( // Common errors used with this package. var ( - ErrNoRequestContext = errors.New("no http request in context") + ErrNoRequestContext = errors.New("no http request in context") + ErrNoResponseWriterContext = errors.New("no http response in context") ) func parseIP(ipStr string) net.IP { @@ -110,6 +111,20 @@ func WithResponseWriter(ctx Context, w http.ResponseWriter) (Context, http.Respo return irw, irw } +// GetResponseWriter returns the http.ResponseWriter from the provided +// context. If not present, ErrNoResponseWriterContext is returned. The +// returned instance provides instrumentation in the context. +func GetResponseWriter(ctx Context) (http.ResponseWriter, error) { + v := ctx.Value("http.response") + + rw, ok := v.(http.ResponseWriter) + if !ok || rw == nil { + return nil, ErrNoResponseWriterContext + } + + return rw, nil +} + // getVarsFromRequest let's us change request vars implementation for testing // and maybe future changes. var getVarsFromRequest = mux.Vars