cloudflared-mirror/vendor/github.com/mitchellh/go-server-timing/middleware.go

85 lines
2.5 KiB
Go
Raw Normal View History

package servertiming
import (
"net/http"
"github.com/felixge/httpsnoop"
)
// MiddlewareOpts are options for the Middleware.
type MiddlewareOpts struct {
// Nothing currently, reserved for the future.
}
// Middleware wraps an http.Handler and provides a *Header in the request
// context that can be used to set Server-Timing headers. The *Header can be
// extracted from the context using FromContext.
//
// The options supplied to this can be nil to use defaults.
//
// The Server-Timing header will be written when the status is written
// only if there are non-empty number of metrics.
//
// To control when Server-Timing is sent, the easiest approach is to wrap
// this middleware and only call it if the request should send server timings.
// For examples, see the README.
func Middleware(next http.Handler, _ *MiddlewareOpts) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var (
// Create the Server-Timing headers struct
h Header
// Remember if the timing header were added to the response headers
headerWritten bool
)
// This places the *Header value into the request context. This
// can be extracted again with FromContext.
r = r.WithContext(NewContext(r.Context(), &h))
// Get the header map. This is a reference and shouldn't change.
headers := w.Header()
// Hook the response writer we pass upstream so we can modify headers
// before they write them to the wire, but after we know what status
// they are writing.
hooks := httpsnoop.Hooks{
WriteHeader: func(original httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
// Return a function with same signature as
// http.ResponseWriter.WriteHeader to be called in it's place
return func(code int) {
// Write the headers
writeHeader(headers, &h)
// Remember that headers were written
headerWritten = true
// Call the original WriteHeader function
original(code)
}
},
}
w = httpsnoop.Wrap(w, hooks)
next.ServeHTTP(w, r)
// In case that next did not called WriteHeader function, add timing header to the response headers
if !headerWritten {
writeHeader(headers, &h)
}
})
}
func writeHeader(headers http.Header, h *Header) {
// Grab the lock just in case there is any ongoing concurrency that
// still has a reference and may be modifying the value.
h.Lock()
defer h.Unlock()
// If there are no metrics set, do nothing
if len(h.Metrics) == 0 {
return
}
headers.Set(HeaderKey, h.String())
}