1// Copyright 2017 The Prometheus Authors
2// Licensed under the Apache License, Version 2.0 (the "License");
3// you may not use this file except in compliance with the License.
4// You may obtain a copy of the License at
5//
6// http://www.apache.org/licenses/LICENSE-2.0
7//
8// Unless required by applicable law or agreed to in writing, software
9// distributed under the License is distributed on an "AS IS" BASIS,
10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11// See the License for the specific language governing permissions and
12// limitations under the License.
13
14package promhttp
15
16import (
17	"bufio"
18	"io"
19	"net"
20	"net/http"
21)
22
23const (
24	closeNotifier = 1 << iota
25	flusher
26	hijacker
27	readerFrom
28	pusher
29)
30
31type delegator interface {
32	http.ResponseWriter
33
34	Status() int
35	Written() int64
36}
37
38type responseWriterDelegator struct {
39	http.ResponseWriter
40
41	status             int
42	written            int64
43	wroteHeader        bool
44	observeWriteHeader func(int)
45}
46
47func (r *responseWriterDelegator) Status() int {
48	return r.status
49}
50
51func (r *responseWriterDelegator) Written() int64 {
52	return r.written
53}
54
55func (r *responseWriterDelegator) WriteHeader(code int) {
56	if r.observeWriteHeader != nil && !r.wroteHeader {
57		// Only call observeWriteHeader for the 1st time. It's a bug if
58		// WriteHeader is called more than once, but we want to protect
59		// against it here. Note that we still delegate the WriteHeader
60		// to the original ResponseWriter to not mask the bug from it.
61		r.observeWriteHeader(code)
62	}
63	r.status = code
64	r.wroteHeader = true
65	r.ResponseWriter.WriteHeader(code)
66}
67
68func (r *responseWriterDelegator) Write(b []byte) (int, error) {
69	// If applicable, call WriteHeader here so that observeWriteHeader is
70	// handled appropriately.
71	if !r.wroteHeader {
72		r.WriteHeader(http.StatusOK)
73	}
74	n, err := r.ResponseWriter.Write(b)
75	r.written += int64(n)
76	return n, err
77}
78
79type closeNotifierDelegator struct{ *responseWriterDelegator }
80type flusherDelegator struct{ *responseWriterDelegator }
81type hijackerDelegator struct{ *responseWriterDelegator }
82type readerFromDelegator struct{ *responseWriterDelegator }
83type pusherDelegator struct{ *responseWriterDelegator }
84
85func (d closeNotifierDelegator) CloseNotify() <-chan bool {
86	//nolint:staticcheck // Ignore SA1019. http.CloseNotifier is deprecated but we keep it here to not break existing users.
87	return d.ResponseWriter.(http.CloseNotifier).CloseNotify()
88}
89func (d flusherDelegator) Flush() {
90	// If applicable, call WriteHeader here so that observeWriteHeader is
91	// handled appropriately.
92	if !d.wroteHeader {
93		d.WriteHeader(http.StatusOK)
94	}
95	d.ResponseWriter.(http.Flusher).Flush()
96}
97func (d hijackerDelegator) Hijack() (net.Conn, *bufio.ReadWriter, error) {
98	return d.ResponseWriter.(http.Hijacker).Hijack()
99}
100func (d readerFromDelegator) ReadFrom(re io.Reader) (int64, error) {
101	// If applicable, call WriteHeader here so that observeWriteHeader is
102	// handled appropriately.
103	if !d.wroteHeader {
104		d.WriteHeader(http.StatusOK)
105	}
106	n, err := d.ResponseWriter.(io.ReaderFrom).ReadFrom(re)
107	d.written += n
108	return n, err
109}
110func (d pusherDelegator) Push(target string, opts *http.PushOptions) error {
111	return d.ResponseWriter.(http.Pusher).Push(target, opts)
112}
113
114var pickDelegator = make([]func(*responseWriterDelegator) delegator, 32)
115
116func init() {
117	// TODO(beorn7): Code generation would help here.
118	pickDelegator[0] = func(d *responseWriterDelegator) delegator { // 0
119		return d
120	}
121	pickDelegator[closeNotifier] = func(d *responseWriterDelegator) delegator { // 1
122		return closeNotifierDelegator{d}
123	}
124	pickDelegator[flusher] = func(d *responseWriterDelegator) delegator { // 2
125		return flusherDelegator{d}
126	}
127	pickDelegator[flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 3
128		return struct {
129			*responseWriterDelegator
130			http.Flusher
131			http.CloseNotifier
132		}{d, flusherDelegator{d}, closeNotifierDelegator{d}}
133	}
134	pickDelegator[hijacker] = func(d *responseWriterDelegator) delegator { // 4
135		return hijackerDelegator{d}
136	}
137	pickDelegator[hijacker+closeNotifier] = func(d *responseWriterDelegator) delegator { // 5
138		return struct {
139			*responseWriterDelegator
140			http.Hijacker
141			http.CloseNotifier
142		}{d, hijackerDelegator{d}, closeNotifierDelegator{d}}
143	}
144	pickDelegator[hijacker+flusher] = func(d *responseWriterDelegator) delegator { // 6
145		return struct {
146			*responseWriterDelegator
147			http.Hijacker
148			http.Flusher
149		}{d, hijackerDelegator{d}, flusherDelegator{d}}
150	}
151	pickDelegator[hijacker+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 7
152		return struct {
153			*responseWriterDelegator
154			http.Hijacker
155			http.Flusher
156			http.CloseNotifier
157		}{d, hijackerDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
158	}
159	pickDelegator[readerFrom] = func(d *responseWriterDelegator) delegator { // 8
160		return readerFromDelegator{d}
161	}
162	pickDelegator[readerFrom+closeNotifier] = func(d *responseWriterDelegator) delegator { // 9
163		return struct {
164			*responseWriterDelegator
165			io.ReaderFrom
166			http.CloseNotifier
167		}{d, readerFromDelegator{d}, closeNotifierDelegator{d}}
168	}
169	pickDelegator[readerFrom+flusher] = func(d *responseWriterDelegator) delegator { // 10
170		return struct {
171			*responseWriterDelegator
172			io.ReaderFrom
173			http.Flusher
174		}{d, readerFromDelegator{d}, flusherDelegator{d}}
175	}
176	pickDelegator[readerFrom+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 11
177		return struct {
178			*responseWriterDelegator
179			io.ReaderFrom
180			http.Flusher
181			http.CloseNotifier
182		}{d, readerFromDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
183	}
184	pickDelegator[readerFrom+hijacker] = func(d *responseWriterDelegator) delegator { // 12
185		return struct {
186			*responseWriterDelegator
187			io.ReaderFrom
188			http.Hijacker
189		}{d, readerFromDelegator{d}, hijackerDelegator{d}}
190	}
191	pickDelegator[readerFrom+hijacker+closeNotifier] = func(d *responseWriterDelegator) delegator { // 13
192		return struct {
193			*responseWriterDelegator
194			io.ReaderFrom
195			http.Hijacker
196			http.CloseNotifier
197		}{d, readerFromDelegator{d}, hijackerDelegator{d}, closeNotifierDelegator{d}}
198	}
199	pickDelegator[readerFrom+hijacker+flusher] = func(d *responseWriterDelegator) delegator { // 14
200		return struct {
201			*responseWriterDelegator
202			io.ReaderFrom
203			http.Hijacker
204			http.Flusher
205		}{d, readerFromDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}}
206	}
207	pickDelegator[readerFrom+hijacker+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 15
208		return struct {
209			*responseWriterDelegator
210			io.ReaderFrom
211			http.Hijacker
212			http.Flusher
213			http.CloseNotifier
214		}{d, readerFromDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
215	}
216	pickDelegator[pusher] = func(d *responseWriterDelegator) delegator { // 16
217		return pusherDelegator{d}
218	}
219	pickDelegator[pusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 17
220		return struct {
221			*responseWriterDelegator
222			http.Pusher
223			http.CloseNotifier
224		}{d, pusherDelegator{d}, closeNotifierDelegator{d}}
225	}
226	pickDelegator[pusher+flusher] = func(d *responseWriterDelegator) delegator { // 18
227		return struct {
228			*responseWriterDelegator
229			http.Pusher
230			http.Flusher
231		}{d, pusherDelegator{d}, flusherDelegator{d}}
232	}
233	pickDelegator[pusher+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 19
234		return struct {
235			*responseWriterDelegator
236			http.Pusher
237			http.Flusher
238			http.CloseNotifier
239		}{d, pusherDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
240	}
241	pickDelegator[pusher+hijacker] = func(d *responseWriterDelegator) delegator { // 20
242		return struct {
243			*responseWriterDelegator
244			http.Pusher
245			http.Hijacker
246		}{d, pusherDelegator{d}, hijackerDelegator{d}}
247	}
248	pickDelegator[pusher+hijacker+closeNotifier] = func(d *responseWriterDelegator) delegator { // 21
249		return struct {
250			*responseWriterDelegator
251			http.Pusher
252			http.Hijacker
253			http.CloseNotifier
254		}{d, pusherDelegator{d}, hijackerDelegator{d}, closeNotifierDelegator{d}}
255	}
256	pickDelegator[pusher+hijacker+flusher] = func(d *responseWriterDelegator) delegator { // 22
257		return struct {
258			*responseWriterDelegator
259			http.Pusher
260			http.Hijacker
261			http.Flusher
262		}{d, pusherDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}}
263	}
264	pickDelegator[pusher+hijacker+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { //23
265		return struct {
266			*responseWriterDelegator
267			http.Pusher
268			http.Hijacker
269			http.Flusher
270			http.CloseNotifier
271		}{d, pusherDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
272	}
273	pickDelegator[pusher+readerFrom] = func(d *responseWriterDelegator) delegator { // 24
274		return struct {
275			*responseWriterDelegator
276			http.Pusher
277			io.ReaderFrom
278		}{d, pusherDelegator{d}, readerFromDelegator{d}}
279	}
280	pickDelegator[pusher+readerFrom+closeNotifier] = func(d *responseWriterDelegator) delegator { // 25
281		return struct {
282			*responseWriterDelegator
283			http.Pusher
284			io.ReaderFrom
285			http.CloseNotifier
286		}{d, pusherDelegator{d}, readerFromDelegator{d}, closeNotifierDelegator{d}}
287	}
288	pickDelegator[pusher+readerFrom+flusher] = func(d *responseWriterDelegator) delegator { // 26
289		return struct {
290			*responseWriterDelegator
291			http.Pusher
292			io.ReaderFrom
293			http.Flusher
294		}{d, pusherDelegator{d}, readerFromDelegator{d}, flusherDelegator{d}}
295	}
296	pickDelegator[pusher+readerFrom+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 27
297		return struct {
298			*responseWriterDelegator
299			http.Pusher
300			io.ReaderFrom
301			http.Flusher
302			http.CloseNotifier
303		}{d, pusherDelegator{d}, readerFromDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
304	}
305	pickDelegator[pusher+readerFrom+hijacker] = func(d *responseWriterDelegator) delegator { // 28
306		return struct {
307			*responseWriterDelegator
308			http.Pusher
309			io.ReaderFrom
310			http.Hijacker
311		}{d, pusherDelegator{d}, readerFromDelegator{d}, hijackerDelegator{d}}
312	}
313	pickDelegator[pusher+readerFrom+hijacker+closeNotifier] = func(d *responseWriterDelegator) delegator { // 29
314		return struct {
315			*responseWriterDelegator
316			http.Pusher
317			io.ReaderFrom
318			http.Hijacker
319			http.CloseNotifier
320		}{d, pusherDelegator{d}, readerFromDelegator{d}, hijackerDelegator{d}, closeNotifierDelegator{d}}
321	}
322	pickDelegator[pusher+readerFrom+hijacker+flusher] = func(d *responseWriterDelegator) delegator { // 30
323		return struct {
324			*responseWriterDelegator
325			http.Pusher
326			io.ReaderFrom
327			http.Hijacker
328			http.Flusher
329		}{d, pusherDelegator{d}, readerFromDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}}
330	}
331	pickDelegator[pusher+readerFrom+hijacker+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 31
332		return struct {
333			*responseWriterDelegator
334			http.Pusher
335			io.ReaderFrom
336			http.Hijacker
337			http.Flusher
338			http.CloseNotifier
339		}{d, pusherDelegator{d}, readerFromDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
340	}
341}
342
343func newDelegator(w http.ResponseWriter, observeWriteHeaderFunc func(int)) delegator {
344	d := &responseWriterDelegator{
345		ResponseWriter:     w,
346		observeWriteHeader: observeWriteHeaderFunc,
347	}
348
349	id := 0
350	//nolint:staticcheck // Ignore SA1019. http.CloseNotifier is deprecated but we keep it here to not break existing users.
351	if _, ok := w.(http.CloseNotifier); ok {
352		id += closeNotifier
353	}
354	if _, ok := w.(http.Flusher); ok {
355		id += flusher
356	}
357	if _, ok := w.(http.Hijacker); ok {
358		id += hijacker
359	}
360	if _, ok := w.(io.ReaderFrom); ok {
361		id += readerFrom
362	}
363	if _, ok := w.(http.Pusher); ok {
364		id += pusher
365	}
366
367	return pickDelegator[id](d)
368}
369