1package middleware
2
3import (
4 "bufio"
5 "context"
6 "errors"
7 "io"
8 "net"
9 "net/http"
10 "sync"
11 "time"
12)
13
14// errTransferFinished is returned by writer calls made after the request
15// finished (or was hijacked). The controller is never touched in that state.
16var errTransferFinished = errors.New("transfer idle: request already finished")
17
18// idleWriteChunk is the progress unit of response writes. Each chunk is handed
19// to the connection separately and counts as progress once it has been fully
20// accepted, so a client that keeps consuming data is never cut while a client
21// that stops reading is cut after the idle timeout. The residual granularity is
22// one chunk: a client must accept idleWriteChunk bytes within one idle period,
23// which is a floor of idleWriteChunk/idle bytes per second (about 68 B/s for
24// the 2 minute default).
25const idleWriteChunk = 8 * 1024
26
27// TransferIdleTimeout returns a middleware that replaces fixed per-request
28// read/write timeouts with a request-scoped inactivity timeout. The server it
29// wraps must not set ReadTimeout or WriteTimeout.
30//
31// A single timer covers the request. Progress renews it: a response chunk
32// accepted by the connection, a flush, or request body bytes read. When idle
33// elapses without progress, the middleware cancels the request context (so a
34// ReverseProxy aborts its upstream request) and expires the connection's read
35// and write deadlines, which unblocks pending I/O and makes the server drop the
36// connection once the handler returns. The timer also covers a handler that is
37// silent and not blocked in any I/O, which a plain write deadline cannot.
38//
39// The timer is armed when the request starts, so the wait for the first
40// response byte is bounded too. It is stopped when the handler returns and when
41// the connection is hijacked: upgraded streams (WebSocket) are exempt.
42// Deadlines go through http.ResponseController, so HTTP/1.x and HTTP/2 both
43// work; Flush, Hijack and Unwrap are preserved.
44//
45// net/http does work after the handler returns that can block on a stalled
46// client: the final response flush and the drain of an unread request body.
47// With no server Read/WriteTimeout those would hang forever, so finish arms
48// one last bounded read and write deadline (idle from now) before the handler
49// returns and leaves it in place. net/http clears the write deadline itself
50// after the response is finished and replaces the read deadline with the
51// keep-alive idle / header deadline before the next request, so nothing leaks
52// into the next request. The timer never touches the controller after return.
53//
54// Limits: a handler that ignores the canceled context and never touches the
55// connection again keeps its HTTP/1 connection until it returns. A handler that
56// enables HTTP/1 full-duplex mode (ResponseController.EnableFullDuplex) is not
57// covered for the unread-body drain: net/http clears the read deadline before
58// that drain, so such a handler must consume or bound its request body itself.
59// A non-positive idle disables the middleware.
60func TransferIdleTimeout(idle time.Duration) func(http.Handler) http.Handler {
61 return func(next http.Handler) http.Handler {
62 if idle <= 0 {
63 return next
64 }
65 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
66 ctx, cancel := context.WithCancel(r.Context())
67 hasBody := r.Body != nil && r.Body != http.NoBody
68 state := newIdleState(idle, http.NewResponseController(w), cancel, hasBody)
69 // finish marks the state done under the mutex before this
70 // function returns, so nothing touches the controller afterwards.
71 defer state.finish()
72
73 req := r.WithContext(ctx)
74 if hasBody {
75 req.Body = &idleBody{ReadCloser: r.Body, state: state}
76 }
77 next.ServeHTTP(&idleResponseWriter{ResponseWriter: w, state: state}, req)
78 })
79 }
80}
81
82// idleState is the request-lifetime state shared by the writer, the body and
83// the timer. The controller is only used while the state is live, under mu.
84type idleState struct {
85 idle time.Duration
86 rc *http.ResponseController
87 cancel context.CancelFunc
88
89 hasBody bool // request has a body that net/http may still have to drain
90
91 mu sync.Mutex
92 timer *time.Timer
93 last time.Time
94 done bool // handler returned or connection hijacked
95 expired bool // the timer fired and expired the connection deadlines
96 hijacked bool
97}
98
99func newIdleState(idle time.Duration, rc *http.ResponseController, cancel context.CancelFunc, hasBody bool) *idleState {
100 s := &idleState{idle: idle, rc: rc, cancel: cancel, last: time.Now(), hasBody: hasBody}
101 s.mu.Lock()
102 s.timer = time.AfterFunc(idle, s.onTimer)
103 s.mu.Unlock()
104 return s
105}
106
107// progress records activity. It never touches the controller.
108func (s *idleState) progress() {
109 s.mu.Lock()
110 if !s.done {
111 s.last = time.Now()
112 }
113 s.mu.Unlock()
114}
115
116// live reports whether the request is still being served.
117func (s *idleState) live() bool {
118 s.mu.Lock()
119 defer s.mu.Unlock()
120 return !s.done
121}
122
123// onTimer runs on timer expiry. Progress only stamps s.last, so the timer is
124// re-armed for the remainder instead of being reset on every event; this also
125// makes progress racing with expiry safe.
126func (s *idleState) onTimer() {
127 s.mu.Lock()
128 defer s.mu.Unlock()
129 if s.done {
130 return
131 }
132 if remaining := s.idle - time.Since(s.last); remaining > 0 {
133 s.timer.Reset(remaining)
134 return
135 }
136 s.expired = true
137 s.cancel()
138 // Held under mu so finish() cannot complete while these run.
139 now := time.Now()
140 _ = s.rc.SetWriteDeadline(now)
141 _ = s.rc.SetReadDeadline(now)
142}
143
144// finish ends the request lifetime. After it returns, no code path calls the
145// controller. It is safe to call more than once. Unless the connection was
146// hijacked or already expired, it first arms bounded deadlines so net/http's
147// post-handler work (final flush, unread body drain) cannot block forever.
148func (s *idleState) finish() {
149 s.mu.Lock()
150 wasDone := s.done
151 s.done = true
152 s.timer.Stop()
153 if !wasDone && !s.expired && !s.hijacked {
154 deadline := time.Now().Add(s.idle)
155 _ = s.rc.SetWriteDeadline(deadline)
156 if s.hasBody {
157 _ = s.rc.SetReadDeadline(deadline)
158 }
159 }
160 s.mu.Unlock()
161 s.cancel()
162}
163
164// flush flushes through the controller unless the request finished. The mutex
165// is not held during the flush: it may block on a slow client, and the timer
166// callback needs the mutex to expire the deadlines that unblock it. net/http
167// forbids using the writer after the handler returned or concurrently with its
168// return, so the live check is a guard against misuse, not a synchronisation
169// point for the flush itself.
170func (s *idleState) flush() error {
171 if !s.live() {
172 return errTransferFinished
173 }
174 return s.rc.Flush()
175}
176
177// hijack hijacks the connection without holding the mutex: HTTP/1 Hijack may
178// flush pending headers and block on a client that does not read, and the timer
179// needs the mutex to expire the deadlines that unblock it. Enforcement stays
180// active until the hijack succeeds; only then is the timer stopped. If the
181// timer fired meanwhile, the deadlines it set on the connection are cleared so
182// the upgraded stream is not cut.
183func (s *idleState) hijack() (net.Conn, *bufio.ReadWriter, error) {
184 if !s.live() {
185 return nil, nil, errTransferFinished
186 }
187 conn, rw, err := s.rc.Hijack()
188 if err != nil {
189 return conn, rw, err
190 }
191 s.mu.Lock()
192 s.done = true
193 s.hijacked = true
194 s.timer.Stop()
195 expired := s.expired
196 s.mu.Unlock()
197 if expired {
198 _ = conn.SetDeadline(time.Time{})
199 }
200 return conn, rw, nil
201}
202
203// idleResponseWriter reports response progress.
204type idleResponseWriter struct {
205 http.ResponseWriter
206 state *idleState
207}
208
209// Write hands p to the connection in idleWriteChunk pieces and records progress
210// after each piece.
211func (w *idleResponseWriter) Write(p []byte) (int, error) {
212 if !w.state.live() {
213 return 0, errTransferFinished
214 }
215 if len(p) == 0 {
216 return w.ResponseWriter.Write(p)
217 }
218 written := 0
219 for written < len(p) {
220 chunk := p[written:min(written+idleWriteChunk, len(p))]
221 n, err := w.ResponseWriter.Write(chunk)
222 written += n
223 if n > 0 {
224 w.state.progress()
225 }
226 if err != nil {
227 return written, err
228 }
229 }
230 return written, nil
231}
232
233// FlushError flushes and records progress. It lets http.ResponseController
234// report flush failures.
235func (w *idleResponseWriter) FlushError() error {
236 err := w.state.flush()
237 if err == nil {
238 w.state.progress()
239 }
240 return err
241}
242
243// Flush implements http.Flusher.
244func (w *idleResponseWriter) Flush() {
245 _ = w.FlushError()
246}
247
248// Hijack implements http.Hijacker and stops the inactivity timer: upgraded
249// streams (e.g. WebSocket) are exempt.
250func (w *idleResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
251 return w.state.hijack()
252}
253
254// Unwrap exposes the underlying writer to http.ResponseController.
255func (w *idleResponseWriter) Unwrap() http.ResponseWriter {
256 return w.ResponseWriter
257}
258
259// idleBody reports request body progress. It only records timestamps, so it is
260// safe to Read or Close after the handler returned (the HTTP/2 transport closes
261// the outbound body asynchronously).
262type idleBody struct {
263 io.ReadCloser
264 state *idleState
265
266 closeOnce sync.Once
267 closeErr error
268}
269
270func (b *idleBody) Read(p []byte) (int, error) {
271 n, err := b.ReadCloser.Read(p)
272 if n > 0 {
273 b.state.progress()
274 }
275 return n, err
276}
277
278// Close is idempotent.
279func (b *idleBody) Close() error {
280 b.closeOnce.Do(func() { b.closeErr = b.ReadCloser.Close() })
281 return b.closeErr
282}