2016-06-15 19:07:33 +02:00
package middlewares
import (
"bufio"
"bytes"
2017-05-03 10:20:33 +02:00
"context"
2017-02-02 17:09:47 +01:00
"io/ioutil"
2016-06-15 19:07:33 +02:00
"net"
"net/http"
2016-12-30 09:21:13 +01:00
"github.com/containous/traefik/log"
"github.com/vulcand/oxy/utils"
2016-06-15 19:07:33 +02:00
)
2017-04-18 08:22:06 +02:00
// Compile time validation responseRecorder implements http interfaces correctly.
2016-08-03 14:50:52 +02:00
var (
2017-04-18 08:22:06 +02:00
_ Stateful = & retryResponseRecorder { }
2016-08-03 14:50:52 +02:00
)
2016-06-15 19:07:33 +02:00
// Retry is a middleware that retries requests
type Retry struct {
attempts int
next http . Handler
2017-04-18 08:22:06 +02:00
listener RetryListener
2016-06-15 19:07:33 +02:00
}
// NewRetry returns a new Retry instance
2017-04-18 08:22:06 +02:00
func NewRetry ( attempts int , next http . Handler , listener RetryListener ) * Retry {
2016-06-15 19:07:33 +02:00
return & Retry {
attempts : attempts ,
next : next ,
2017-04-18 08:22:06 +02:00
listener : listener ,
2016-06-15 19:07:33 +02:00
}
}
func ( retry * Retry ) ServeHTTP ( rw http . ResponseWriter , r * http . Request ) {
2017-02-02 17:09:47 +01:00
// if we might make multiple attempts, swap the body for an ioutil.NopCloser
// cf https://github.com/containous/traefik/issues/1008
if retry . attempts > 1 {
body := r . Body
defer body . Close ( )
r . Body = ioutil . NopCloser ( body )
}
2016-06-15 19:07:33 +02:00
attempts := 1
for {
2017-05-03 10:20:33 +02:00
netErrorOccurred := false
// We pass in a pointer to netErrorOccurred so that we can set it to true on network errors
// when proxying the HTTP requests to the backends. This happens in the custom RecordingErrorHandler.
newCtx := context . WithValue ( r . Context ( ) , defaultNetErrCtxKey , & netErrorOccurred )
2017-04-18 08:22:06 +02:00
recorder := newRetryResponseRecorder ( )
2016-06-15 19:07:33 +02:00
recorder . responseWriter = rw
2017-05-03 10:20:33 +02:00
retry . next . ServeHTTP ( recorder , r . WithContext ( newCtx ) )
2017-09-20 18:40:03 +02:00
// It's a stream request and the body gets already sent to the client.
// Therefore we should not send the response a second time.
if recorder . streamingResponseStarted {
break
}
2017-05-03 10:20:33 +02:00
if ! netErrorOccurred || attempts >= retry . attempts {
2016-07-04 19:32:19 +02:00
utils . CopyHeaders ( rw . Header ( ) , recorder . Header ( ) )
2016-06-15 19:07:33 +02:00
rw . WriteHeader ( recorder . Code )
rw . Write ( recorder . Body . Bytes ( ) )
break
}
attempts ++
log . Debugf ( "New attempt %d for request: %v" , attempts , r . URL )
2017-08-28 12:50:02 +02:00
retry . listener . Retried ( r , attempts )
2016-06-15 19:07:33 +02:00
}
}
2017-05-03 10:20:33 +02:00
// netErrorCtxKey is a custom type that is used as key for the context.
type netErrorCtxKey string
// defaultNetErrCtxKey is the actual key which value is used to record network errors.
var defaultNetErrCtxKey netErrorCtxKey = "NetErrCtxKey"
// NetErrorRecorder is an interface to record net errors.
type NetErrorRecorder interface {
// Record can be used to signal the retry middleware that an network error happened
// and therefore the request should be retried.
Record ( ctx context . Context )
}
// DefaultNetErrorRecorder is the default NetErrorRecorder implementation.
type DefaultNetErrorRecorder struct { }
// Record is recording network errors by setting the context value for the defaultNetErrCtxKey to true.
func ( DefaultNetErrorRecorder ) Record ( ctx context . Context ) {
val := ctx . Value ( defaultNetErrCtxKey )
if netErrorOccurred , isBoolPointer := val . ( * bool ) ; isBoolPointer {
* netErrorOccurred = true
}
2016-06-15 19:07:33 +02:00
}
2017-04-18 08:22:06 +02:00
// RetryListener is used to inform about retry attempts.
type RetryListener interface {
// Retried will be called when a retry happens, with the request attempt passed to it.
// For the first retry this will be attempt 2.
2017-08-28 12:50:02 +02:00
Retried ( req * http . Request , attempt int )
}
// RetryListeners is a convenience type to construct a list of RetryListener and notify
// each of them about a retry attempt.
type RetryListeners [ ] RetryListener
// Retried exists to implement the RetryListener interface. It calls Retried on each of its slice entries.
func ( l RetryListeners ) Retried ( req * http . Request , attempt int ) {
for _ , retryListener := range l {
retryListener . Retried ( req , attempt )
}
2017-04-18 08:22:06 +02:00
}
// retryResponseRecorder is an implementation of http.ResponseWriter that
// records its mutations for later inspection.
type retryResponseRecorder struct {
2016-06-15 19:07:33 +02:00
Code int // the HTTP response code from WriteHeader
HeaderMap http . Header // the HTTP response headers
Body * bytes . Buffer // if non-nil, the bytes.Buffer to append written data to
2017-09-20 18:40:03 +02:00
responseWriter http . ResponseWriter
err error
streamingResponseStarted bool
2016-06-15 19:07:33 +02:00
}
2017-04-18 08:22:06 +02:00
// newRetryResponseRecorder returns an initialized retryResponseRecorder.
func newRetryResponseRecorder ( ) * retryResponseRecorder {
return & retryResponseRecorder {
2016-06-15 19:07:33 +02:00
HeaderMap : make ( http . Header ) ,
Body : new ( bytes . Buffer ) ,
Code : 200 ,
}
}
// Header returns the response headers.
2017-04-18 08:22:06 +02:00
func ( rw * retryResponseRecorder ) Header ( ) http . Header {
2016-06-15 19:07:33 +02:00
m := rw . HeaderMap
if m == nil {
m = make ( http . Header )
rw . HeaderMap = m
}
return m
}
// Write always succeeds and writes to rw.Body, if not nil.
2017-04-18 08:22:06 +02:00
func ( rw * retryResponseRecorder ) Write ( buf [ ] byte ) ( int , error ) {
2016-08-03 14:50:52 +02:00
if rw . err != nil {
return 0 , rw . err
2016-06-15 19:07:33 +02:00
}
2016-08-03 14:50:52 +02:00
return rw . Body . Write ( buf )
2016-06-15 19:07:33 +02:00
}
// WriteHeader sets rw.Code.
2017-04-18 08:22:06 +02:00
func ( rw * retryResponseRecorder ) WriteHeader ( code int ) {
2016-07-04 19:32:19 +02:00
rw . Code = code
2016-06-15 19:07:33 +02:00
}
// Hijack hijacks the connection
2017-04-18 08:22:06 +02:00
func ( rw * retryResponseRecorder ) Hijack ( ) ( net . Conn , * bufio . ReadWriter , error ) {
2016-06-15 19:07:33 +02:00
return rw . responseWriter . ( http . Hijacker ) . Hijack ( )
}
2016-08-03 14:50:52 +02:00
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone
// away.
2017-04-18 08:22:06 +02:00
func ( rw * retryResponseRecorder ) CloseNotify ( ) <- chan bool {
2016-08-03 14:50:52 +02:00
return rw . responseWriter . ( http . CloseNotifier ) . CloseNotify ( )
}
// Flush sends any buffered data to the client.
2017-04-18 08:22:06 +02:00
func ( rw * retryResponseRecorder ) Flush ( ) {
2017-09-20 18:40:03 +02:00
if ! rw . streamingResponseStarted {
utils . CopyHeaders ( rw . responseWriter . Header ( ) , rw . Header ( ) )
rw . responseWriter . WriteHeader ( rw . Code )
rw . streamingResponseStarted = true
}
2016-08-03 14:50:52 +02:00
_ , err := rw . responseWriter . Write ( rw . Body . Bytes ( ) )
if err != nil {
2017-04-18 08:22:06 +02:00
log . Errorf ( "Error writing response in retryResponseRecorder: %s" , err )
2016-08-03 14:50:52 +02:00
rw . err = err
}
rw . Body . Reset ( )
flusher , ok := rw . responseWriter . ( http . Flusher )
if ok {
flusher . Flush ( )
}
}