-
Notifications
You must be signed in to change notification settings - Fork 0
/
context.go
42 lines (35 loc) · 1.02 KB
/
context.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
// Package wrap provides wrappers for [http.Handler]s.
//
//spellchecker:words wrap
package wrap
//spellchecker:words context http
import (
"context"
"net/http"
)
// ContextFunc is a function that replaces contexts for a given request.
// A nil ContextFunc leaves the original context intact.
//
// - the returned context, if non-nil, is used to replace the context of the request.
// - the returned CancelFunc is called once the request ends.
type ContextFunc = func(r *http.Request) (context.Context, context.CancelFunc)
// Context wraps handler, replacing the context of any request using the given function.
func Context(handler http.Handler, f ContextFunc) http.Handler {
if f == nil {
return handler
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// compute the context
ctx, cancel := f(r)
// call cancel once the request ends
if cancel != nil {
defer cancel()
}
// use the new context
if ctx != nil {
r = r.WithContext(ctx)
}
// do the handling
handler.ServeHTTP(w, r)
})
}