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