From 93166a5b200f62b7bb474670a1a1931f6b031f8f Mon Sep 17 00:00:00 2001
From: Will Norris <will@willnorris.com>
Date: Tue, 29 Nov 2016 13:50:19 -0800
Subject: [PATCH] add support for per-request timeout

Adds a -timeout flag for specifying the timeout.  Currently, this
returns a 503 response on timeout, though it should really be a 504,
since imageproxy is acting as a gateway.

Ref: #75
---
 cmd/imageproxy/main.go |  2 ++
 imageproxy.go          | 16 +++++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/cmd/imageproxy/main.go b/cmd/imageproxy/main.go
index 613ec95..396ca12 100644
--- a/cmd/imageproxy/main.go
+++ b/cmd/imageproxy/main.go
@@ -49,6 +49,7 @@ var cacheDir = flag.String("cacheDir", "", "(Deprecated; use 'cache' instead) di
 var cacheSize = flag.Uint64("cacheSize", 0, "Deprecated: this flag does nothing")
 var signatureKey = flag.String("signatureKey", "", "HMAC key used in calculating request signatures")
 var scaleUp = flag.Bool("scaleUp", false, "allow images to scale beyond their original dimensions")
+var timeout = flag.Duration("timeout", 0, "time limit for requests served by this proxy")
 var version = flag.Bool("version", false, "print version information")
 
 func main() {
@@ -91,6 +92,7 @@ func main() {
 		}
 	}
 
+	p.Timeout = *timeout
 	p.ScaleUp = *scaleUp
 
 	server := &http.Server{
diff --git a/imageproxy.go b/imageproxy.go
index c8682e9..22a7c56 100644
--- a/imageproxy.go
+++ b/imageproxy.go
@@ -58,6 +58,11 @@ type Proxy struct {
 
 	// Allow images to scale beyond their original dimensions.
 	ScaleUp bool
+
+	// Timeout specifies a time limit for requests served by this Proxy.
+	// If a call runs for longer than its time limit, a 504 Gateway Timeout
+	// response is returned.  A Timeout of zero means no timeout.
+	Timeout time.Duration
 }
 
 // NewProxy constructs a new proxy.  The provided http RoundTripper will be
@@ -87,7 +92,7 @@ func NewProxy(transport http.RoundTripper, cache Cache) *Proxy {
 	return &proxy
 }
 
-// ServeHTTP handles image requests.
+// ServeHTTP handles incoming requests.
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if r.URL.Path == "/favicon.ico" {
 		return // ignore favicon requests
@@ -98,6 +103,15 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	var h http.Handler = http.HandlerFunc(p.serveImage)
+	if p.Timeout > 0 {
+		h = http.TimeoutHandler(h, p.Timeout, "")
+	}
+	h.ServeHTTP(w, r)
+}
+
+// serveImage handles incoming requests for proxied images.
+func (p *Proxy) serveImage(w http.ResponseWriter, r *http.Request) {
 	req, err := NewRequest(r, p.DefaultBaseURL)
 	if err != nil {
 		msg := fmt.Sprintf("invalid request URL: %v", err)