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	r.status = code
57	r.wroteHeader = true
58	r.ResponseWriter.WriteHeader(code)
59	if r.observeWriteHeader != nil {
60		r.observeWriteHeader(code)
61	}
62}
63
64func (r *responseWriterDelegator) Write(b []byte) (int, error) {
65	if !r.wroteHeader {
66		r.WriteHeader(http.StatusOK)
67	}
68	n, err := r.ResponseWriter.Write(b)
69	r.written += int64(n)
70	return n, err
71}
72
73type closeNotifierDelegator struct{ *responseWriterDelegator }
74type flusherDelegator struct{ *responseWriterDelegator }
75type hijackerDelegator struct{ *responseWriterDelegator }
76type readerFromDelegator struct{ *responseWriterDelegator }
77
78func (d closeNotifierDelegator) CloseNotify() <-chan bool {
79	return d.ResponseWriter.(http.CloseNotifier).CloseNotify()
80}
81func (d flusherDelegator) Flush() {
82	d.ResponseWriter.(http.Flusher).Flush()
83}
84func (d hijackerDelegator) Hijack() (net.Conn, *bufio.ReadWriter, error) {
85	return d.ResponseWriter.(http.Hijacker).Hijack()
86}
87func (d readerFromDelegator) ReadFrom(re io.Reader) (int64, error) {
88	if !d.wroteHeader {
89		d.WriteHeader(http.StatusOK)
90	}
91	n, err := d.ResponseWriter.(io.ReaderFrom).ReadFrom(re)
92	d.written += n
93	return n, err
94}
95
96var pickDelegator = make([]func(*responseWriterDelegator) delegator, 32)
97
98func init() {
99	// TODO(beorn7): Code generation would help here.
100	pickDelegator[0] = func(d *responseWriterDelegator) delegator { // 0
101		return d
102	}
103	pickDelegator[closeNotifier] = func(d *responseWriterDelegator) delegator { // 1
104		return closeNotifierDelegator{d}
105	}
106	pickDelegator[flusher] = func(d *responseWriterDelegator) delegator { // 2
107		return flusherDelegator{d}
108	}
109	pickDelegator[flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 3
110		return struct {
111			*responseWriterDelegator
112			http.Flusher
113			http.CloseNotifier
114		}{d, flusherDelegator{d}, closeNotifierDelegator{d}}
115	}
116	pickDelegator[hijacker] = func(d *responseWriterDelegator) delegator { // 4
117		return hijackerDelegator{d}
118	}
119	pickDelegator[hijacker+closeNotifier] = func(d *responseWriterDelegator) delegator { // 5
120		return struct {
121			*responseWriterDelegator
122			http.Hijacker
123			http.CloseNotifier
124		}{d, hijackerDelegator{d}, closeNotifierDelegator{d}}
125	}
126	pickDelegator[hijacker+flusher] = func(d *responseWriterDelegator) delegator { // 6
127		return struct {
128			*responseWriterDelegator
129			http.Hijacker
130			http.Flusher
131		}{d, hijackerDelegator{d}, flusherDelegator{d}}
132	}
133	pickDelegator[hijacker+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 7
134		return struct {
135			*responseWriterDelegator
136			http.Hijacker
137			http.Flusher
138			http.CloseNotifier
139		}{d, hijackerDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
140	}
141	pickDelegator[readerFrom] = func(d *responseWriterDelegator) delegator { // 8
142		return readerFromDelegator{d}
143	}
144	pickDelegator[readerFrom+closeNotifier] = func(d *responseWriterDelegator) delegator { // 9
145		return struct {
146			*responseWriterDelegator
147			io.ReaderFrom
148			http.CloseNotifier
149		}{d, readerFromDelegator{d}, closeNotifierDelegator{d}}
150	}
151	pickDelegator[readerFrom+flusher] = func(d *responseWriterDelegator) delegator { // 10
152		return struct {
153			*responseWriterDelegator
154			io.ReaderFrom
155			http.Flusher
156		}{d, readerFromDelegator{d}, flusherDelegator{d}}
157	}
158	pickDelegator[readerFrom+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 11
159		return struct {
160			*responseWriterDelegator
161			io.ReaderFrom
162			http.Flusher
163			http.CloseNotifier
164		}{d, readerFromDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
165	}
166	pickDelegator[readerFrom+hijacker] = func(d *responseWriterDelegator) delegator { // 12
167		return struct {
168			*responseWriterDelegator
169			io.ReaderFrom
170			http.Hijacker
171		}{d, readerFromDelegator{d}, hijackerDelegator{d}}
172	}
173	pickDelegator[readerFrom+hijacker+closeNotifier] = func(d *responseWriterDelegator) delegator { // 13
174		return struct {
175			*responseWriterDelegator
176			io.ReaderFrom
177			http.Hijacker
178			http.CloseNotifier
179		}{d, readerFromDelegator{d}, hijackerDelegator{d}, closeNotifierDelegator{d}}
180	}
181	pickDelegator[readerFrom+hijacker+flusher] = func(d *responseWriterDelegator) delegator { // 14
182		return struct {
183			*responseWriterDelegator
184			io.ReaderFrom
185			http.Hijacker
186			http.Flusher
187		}{d, readerFromDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}}
188	}
189	pickDelegator[readerFrom+hijacker+flusher+closeNotifier] = func(d *responseWriterDelegator) delegator { // 15
190		return struct {
191			*responseWriterDelegator
192			io.ReaderFrom
193			http.Hijacker
194			http.Flusher
195			http.CloseNotifier
196		}{d, readerFromDelegator{d}, hijackerDelegator{d}, flusherDelegator{d}, closeNotifierDelegator{d}}
197	}
198}
199