1package http
2
3import (
4	"context"
5	"net/http"
6)
7
8// RequestFunc may take information from an HTTP request and put it into a
9// request context. In Servers, RequestFuncs are executed prior to invoking the
10// endpoint. In Clients, RequestFuncs are executed after creating the request
11// but prior to invoking the HTTP client.
12type RequestFunc func(context.Context, *http.Request) context.Context
13
14// ServerResponseFunc may take information from a request context and use it to
15// manipulate a ResponseWriter. ServerResponseFuncs are only executed in
16// servers, after invoking the endpoint but prior to writing a response.
17type ServerResponseFunc func(context.Context, http.ResponseWriter) context.Context
18
19// ClientResponseFunc may take information from an HTTP request and make the
20// response available for consumption. ClientResponseFuncs are only executed in
21// clients, after a request has been made, but prior to it being decoded.
22type ClientResponseFunc func(context.Context, *http.Response) context.Context
23
24// SetContentType returns a ServerResponseFunc that sets the Content-Type header
25// to the provided value.
26func SetContentType(contentType string) ServerResponseFunc {
27	return SetResponseHeader("Content-Type", contentType)
28}
29
30// SetResponseHeader returns a ServerResponseFunc that sets the given header.
31func SetResponseHeader(key, val string) ServerResponseFunc {
32	return func(ctx context.Context, w http.ResponseWriter) context.Context {
33		w.Header().Set(key, val)
34		return ctx
35	}
36}
37
38// SetRequestHeader returns a RequestFunc that sets the given header.
39func SetRequestHeader(key, val string) RequestFunc {
40	return func(ctx context.Context, r *http.Request) context.Context {
41		r.Header.Set(key, val)
42		return ctx
43	}
44}
45
46// PopulateRequestContext is a RequestFunc that populates several values into
47// the context from the HTTP request. Those values may be extracted using the
48// corresponding ContextKey type in this package.
49func PopulateRequestContext(ctx context.Context, r *http.Request) context.Context {
50	for k, v := range map[contextKey]string{
51		ContextKeyRequestMethod:          r.Method,
52		ContextKeyRequestURI:             r.RequestURI,
53		ContextKeyRequestPath:            r.URL.Path,
54		ContextKeyRequestProto:           r.Proto,
55		ContextKeyRequestHost:            r.Host,
56		ContextKeyRequestRemoteAddr:      r.RemoteAddr,
57		ContextKeyRequestXForwardedFor:   r.Header.Get("X-Forwarded-For"),
58		ContextKeyRequestXForwardedProto: r.Header.Get("X-Forwarded-Proto"),
59		ContextKeyRequestAuthorization:   r.Header.Get("Authorization"),
60		ContextKeyRequestReferer:         r.Header.Get("Referer"),
61		ContextKeyRequestUserAgent:       r.Header.Get("User-Agent"),
62		ContextKeyRequestXRequestID:      r.Header.Get("X-Request-Id"),
63		ContextKeyRequestAccept:          r.Header.Get("Accept"),
64	} {
65		ctx = context.WithValue(ctx, k, v)
66	}
67	return ctx
68}
69
70type contextKey int
71
72const (
73	// ContextKeyRequestMethod is populated in the context by
74	// PopulateRequestContext. Its value is r.Method.
75	ContextKeyRequestMethod contextKey = iota
76
77	// ContextKeyRequestURI is populated in the context by
78	// PopulateRequestContext. Its value is r.RequestURI.
79	ContextKeyRequestURI
80
81	// ContextKeyRequestPath is populated in the context by
82	// PopulateRequestContext. Its value is r.URL.Path.
83	ContextKeyRequestPath
84
85	// ContextKeyRequestProto is populated in the context by
86	// PopulateRequestContext. Its value is r.Proto.
87	ContextKeyRequestProto
88
89	// ContextKeyRequestHost is populated in the context by
90	// PopulateRequestContext. Its value is r.Host.
91	ContextKeyRequestHost
92
93	// ContextKeyRequestRemoteAddr is populated in the context by
94	// PopulateRequestContext. Its value is r.RemoteAddr.
95	ContextKeyRequestRemoteAddr
96
97	// ContextKeyRequestXForwardedFor is populated in the context by
98	// PopulateRequestContext. Its value is r.Header.Get("X-Forwarded-For").
99	ContextKeyRequestXForwardedFor
100
101	// ContextKeyRequestXForwardedProto is populated in the context by
102	// PopulateRequestContext. Its value is r.Header.Get("X-Forwarded-Proto").
103	ContextKeyRequestXForwardedProto
104
105	// ContextKeyRequestAuthorization is populated in the context by
106	// PopulateRequestContext. Its value is r.Header.Get("Authorization").
107	ContextKeyRequestAuthorization
108
109	// ContextKeyRequestReferer is populated in the context by
110	// PopulateRequestContext. Its value is r.Header.Get("Referer").
111	ContextKeyRequestReferer
112
113	// ContextKeyRequestUserAgent is populated in the context by
114	// PopulateRequestContext. Its value is r.Header.Get("User-Agent").
115	ContextKeyRequestUserAgent
116
117	// ContextKeyRequestXRequestID is populated in the context by
118	// PopulateRequestContext. Its value is r.Header.Get("X-Request-Id").
119	ContextKeyRequestXRequestID
120
121	// ContextKeyRequestAccept is populated in the context by
122	// PopulateRequestContext. Its value is r.Header.Get("Accept").
123	ContextKeyRequestAccept
124
125	// ContextKeyResponseHeaders is populated in the context whenever a
126	// ServerFinalizerFunc is specified. Its value is of type http.Header, and
127	// is captured only once the entire response has been written.
128	ContextKeyResponseHeaders
129
130	// ContextKeyResponseSize is populated in the context whenever a
131	// ServerFinalizerFunc is specified. Its value is of type int64.
132	ContextKeyResponseSize
133)
134