@@ -6,12 +6,15 @@ package httptest
6
6
7
7
import (
8
8
"bytes"
9
+ "errors"
9
10
"fmt"
10
11
"io"
11
12
"net/http"
12
13
"net/textproto"
14
+ "os"
13
15
"strconv"
14
16
"strings"
17
+ "time"
15
18
16
19
"golang.org/x/net/http/httpguts"
17
20
)
@@ -42,9 +45,18 @@ type ResponseRecorder struct {
42
45
// Flushed is whether the Handler called Flush.
43
46
Flushed bool
44
47
48
+ // ReadDeadline is the write deadline that has been set using
49
+ // "net/http".ResponseController
50
+ ReadDeadline time.Time
51
+
52
+ // WriteDeadline is the write deadline that has been set using
53
+ // "net/http".ResponseController
54
+ WriteDeadline time.Time
55
+
45
56
result * http.Response // cache of Result's return value
46
57
snapHeader http.Header // snapshot of HeaderMap at first Write
47
58
wroteHeader bool
59
+ requestBody * deadlineBodyReader
48
60
}
49
61
50
62
// NewRecorder returns an initialized ResponseRecorder.
@@ -56,6 +68,21 @@ func NewRecorder() *ResponseRecorder {
56
68
}
57
69
}
58
70
71
+ // NewRecorderWithDeadlineAwareRequest returns an initialized ResponseRecorder
72
+ // and wraps the body of the HTTP request passed as parameter in a special "io".ReadCloser
73
+ // that supports read deadlines.
74
+ // The request read deadline can be set using ResponseRecorder.SetReadDeadline
75
+ // and "http".ResponseController.
76
+ // The body of returned the HTTP request returns an error when the read deadline is reached.
77
+ // The read deadline can be inspected by reading ResponseRecorder.ReadDeadline.
78
+ func NewRecorderWithDeadlineAwareRequest (r * http.Request ) (* ResponseRecorder , * http.Request ) {
79
+ rw := NewRecorder ()
80
+ rw .requestBody = & deadlineBodyReader {r .Body , time.Time {}}
81
+ r .Body = rw .requestBody
82
+
83
+ return rw , r
84
+ }
85
+
59
86
// DefaultRemoteAddr is the default remote address to return in RemoteAddr if
60
87
// an explicit DefaultRemoteAddr isn't set on ResponseRecorder.
61
88
const DefaultRemoteAddr = "1.2.3.4"
@@ -105,6 +132,10 @@ func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
105
132
// Write implements http.ResponseWriter. The data in buf is written to
106
133
// rw.Body, if not nil.
107
134
func (rw * ResponseRecorder ) Write (buf []byte ) (int , error ) {
135
+ if ! rw .WriteDeadline .IsZero () && time .Now ().After (rw .WriteDeadline ) {
136
+ return 0 , os .ErrDeadlineExceeded
137
+ }
138
+
108
139
rw .writeHeader (buf , "" )
109
140
if rw .Body != nil {
110
141
rw .Body .Write (buf )
@@ -115,6 +146,10 @@ func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
115
146
// WriteString implements io.StringWriter. The data in str is written
116
147
// to rw.Body, if not nil.
117
148
func (rw * ResponseRecorder ) WriteString (str string ) (int , error ) {
149
+ if ! rw .WriteDeadline .IsZero () && time .Now ().After (rw .WriteDeadline ) {
150
+ return 0 , os .ErrDeadlineExceeded
151
+ }
152
+
118
153
rw .writeHeader (nil , str )
119
154
if rw .Body != nil {
120
155
rw .Body .WriteString (str )
@@ -154,13 +189,54 @@ func (rw *ResponseRecorder) WriteHeader(code int) {
154
189
rw .snapHeader = rw .HeaderMap .Clone ()
155
190
}
156
191
157
- // Flush implements http.Flusher. To test whether Flush was
192
+ // FlushError allows using "net/http".ResponseController.Flush()
193
+ // with the recorder. To test whether Flush was
158
194
// called, see rw.Flushed.
159
- func (rw * ResponseRecorder ) Flush () {
195
+ func (rw * ResponseRecorder ) FlushError () error {
196
+ if ! rw .WriteDeadline .IsZero () && time .Now ().After (rw .WriteDeadline ) {
197
+ return os .ErrDeadlineExceeded
198
+ }
199
+
160
200
if ! rw .wroteHeader {
161
201
rw .WriteHeader (200 )
162
202
}
163
203
rw .Flushed = true
204
+
205
+ return nil
206
+ }
207
+
208
+ // Flush implements http.Flusher. To test whether Flush was
209
+ // called, see rw.Flushed.
210
+ func (rw * ResponseRecorder ) Flush () {
211
+ rw .FlushError ()
212
+ }
213
+
214
+ // SetReadDeadline allows using "net/http".ResponseController.SetReadDeadline()
215
+ // with the recorder.
216
+ // To retrieve the deadline, use rw.ReadDeadline.
217
+ // To use this method, be sure NewRecorderWithDeadlineAwareRequest
218
+ func (rw * ResponseRecorder ) SetReadDeadline (deadline time.Time ) error {
219
+ if rw .requestBody == nil {
220
+ return errors .New ("The request has not been created using NewRecorderWithDeadlineAwareRequest()" )
221
+ }
222
+
223
+ if deadline .After (rw .ReadDeadline ) {
224
+ rw .ReadDeadline = deadline
225
+ rw .requestBody .deadline = deadline
226
+ }
227
+
228
+ return nil
229
+ }
230
+
231
+ // SetWriteDeadline allows using "net/http".ResponseController.SetWriteDeadline()
232
+ // with the recorder.
233
+ // To retrieve the deadline, use rw.WriteDeadline.
234
+ func (rw * ResponseRecorder ) SetWriteDeadline (deadline time.Time ) error {
235
+ if deadline .After (rw .WriteDeadline ) {
236
+ rw .WriteDeadline = deadline
237
+ }
238
+
239
+ return nil
164
240
}
165
241
166
242
// Result returns the response generated by the handler.
@@ -253,3 +329,16 @@ func parseContentLength(cl string) int64 {
253
329
}
254
330
return int64 (n )
255
331
}
332
+
333
+ type deadlineBodyReader struct {
334
+ io.ReadCloser
335
+ deadline time.Time
336
+ }
337
+
338
+ func (r * deadlineBodyReader ) Read (p []byte ) (n int , err error ) {
339
+ if time .Now ().After (r .deadline ) {
340
+ return 0 , os .ErrDeadlineExceeded
341
+ }
342
+
343
+ return r .ReadCloser .Read (p )
344
+ }
0 commit comments