Check if response is to large for cache

cmd/httpcache/main.go

@@ -6,6 +6,7 @@ import (
 	"github.com/donutloop/httpcache/internal/cache"
 	"github.com/donutloop/httpcache/internal/handler"
 	"github.com/donutloop/httpcache/internal/middleware"
+	"github.com/donutloop/httpcache/internal/size"
 	"github.com/donutloop/httpcache/internal/xhttp"
 	"log"
 	"net"
@@ -20,12 +21,13 @@ func main() {
 
 	fs := flag.NewFlagSet("http-proxy", flag.ExitOnError)
 	var (
-		httpAddr = fs.String("http", ":80", "serve HTTP on this address (optional)")
-		tlsAddr  = fs.String("tls", "", "serve TLS on this address (optional)")
-		cert     = fs.String("cert", "server.crt", "TLS certificate")
-		key      = fs.String("key", "server.key", "TLS key")
-		cap      = fs.Int64("cap", 100, "capacity of cache")
-		expire      = fs.Int64("expire", 5, "the items in the cache expire after or expire never")
+		httpAddr                       = fs.String("http", ":80", "serve HTTP on this address (optional)")
+		tlsAddr                        = fs.String("tls", "", "serve TLS on this address (optional)")
+		cert                           = fs.String("cert", "server.crt", "TLS certificate")
+		key                            = fs.String("key", "server.key", "TLS key")
+		cap                            = fs.Int64("cap", 100, "capacity of cache")
+		responseBodyContentLenghtLimit = fs.Int64("rbcl", 500*size.MB, "response size limit")
+		expire                         = fs.Int64("expire", 5, "the items in the cache expire after or expire never")
 	)
 	fs.Usage = usageFor(fs, "httpcache [flags]")
 	fs.Parse(os.Args[1:])
@@ -41,7 +43,7 @@ func main() {
 		}
 	}
 
-	proxy := handler.NewProxy(c, logger.Println)
+	proxy := handler.NewProxy(c, logger.Println, *responseBodyContentLenghtLimit)
 	stats := handler.NewStats(c, logger.Println)
 
 	mux := http.NewServeMux()
@@ -57,9 +59,9 @@ func main() {
 		}
 
 		xserver := xhttp.Server{
-			Server: &http.Server{Addr: *httpAddr, Handler: stack},
-			Logger: logger,
-			Listener: listener,
+			Server:          &http.Server{Addr: *httpAddr, Handler: stack},
+			Logger:          logger,
+			Listener:        listener,
 			ShutdownTimeout: 3 * time.Second,
 		}
 		if err := xserver.Start(); err != nil {
@@ -77,9 +79,9 @@ func main() {
 		}
 
 		xserver := xhttp.Server{
-			Server: &http.Server{Addr: *tlsAddr, Handler: stack},
-			Logger: logger,
-			Listener: listener,
+			Server:          &http.Server{Addr: *tlsAddr, Handler: stack},
+			Logger:          logger,
+			Listener:        listener,
 			ShutdownTimeout: 3 * time.Second,
 		}
 		if err := xserver.StartTLS(*cert, *key); err != nil {

internal/cache/cache.go

@@ -57,14 +57,13 @@ type entry struct {
 
 // NewLRUCache creates a new empty cache with the given capacity.
 func NewLRUCache(capacity int64, expiry time.Duration) *LRUCache {
-	cache :=  &LRUCache{
+	cache := &LRUCache{
 		list:     list.New(),
 		table:    make(map[string]*list.Element),
 		capacity: capacity,
-		expiry:            expiry,
+		expiry:   expiry,
 	}
 
-
 	// We have expiry start the janitor routine.
 	if expiry > 0 {
 		// Initialize a new stop GC channel.
@@ -246,4 +245,4 @@ func (c *LRUCache) StartGC() {
 			}
 		}
 	}()
-}
+}

internal/handler/proxycache.go

@@ -7,15 +7,19 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/http/httputil"
+	"strings"
 )
 
-func NewProxy(cache *cache.LRUCache, logger func(v ...interface{})) *Proxy {
+func NewProxy(cache *cache.LRUCache, logger func(v ...interface{}), contentLength int64) *Proxy {
 	return &Proxy{
 		client: &http.Client{
 			Transport: &roundtripper.LoggedTransport{
 				Transport: &roundtripper.CacheTransport{
-					Transport: http.DefaultTransport,
-					Cache:    cache,
+					Transport: &roundtripper.ResponseBodyLimitRoundTripper{
+						Transport: http.DefaultTransport,
+						Limit:     contentLength,
+					},
+					Cache: cache,
 				},
 				Logger: logger,
 			}},
@@ -33,7 +37,10 @@ func (p *Proxy) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 	req.RequestURI = ""
 	proxyResponse, err := p.client.Do(req)
 	if err != nil {
-		p.logger(err.Error())
+		if strings.Contains(err.Error(), roundtripper.ResponseIsToLarge.Error()) {
+			resp.WriteHeader(http.StatusRequestEntityTooLarge)
+			return
+		}
 		resp.WriteHeader(http.StatusInternalServerError)
 		return
 	}
@@ -73,4 +80,4 @@ func dump(request *http.Request, response *http.Response) (requestDump, response
 		return nil, nil, err
 	}
 	return dumpedRequest, dumpedResponse, nil
-}
+}

internal/handler/stats.go

@@ -10,20 +10,20 @@ import (
 
 func NewStats(c *cache.LRUCache, logger func(v ...interface{})) *Stats {
 	return &Stats{
-		c: c,
+		c:      c,
 		logger: logger,
 	}
 }
 
 type StatsResponse struct {
-	Length int64 `json:"length"`
-	Size int64   `json:"size"`
-	Capacity int64 `json:"capacity"`
-	Oldest time.Time `json:"oldest"`
+	Length   int64     `json:"length"`
+	Size     int64     `json:"size"`
+	Capacity int64     `json:"capacity"`
+	Oldest   time.Time `json:"oldest"`
 }
 
 type Stats struct {
-	c *cache.LRUCache
+	c      *cache.LRUCache
 	logger func(v ...interface{})
 }
 
@@ -47,12 +47,11 @@ func (s *Stats) Endpoint() *StatsResponse {
 	length, size, capacity, oldest := s.c.Stats()
 
 	resp := &StatsResponse{
-		Length: length,
-		Size: size,
+		Length:   length,
+		Size:     size,
 		Capacity: capacity,
-		Oldest: oldest,
+		Oldest:   oldest,
 	}
 
 	return resp
 }
-

internal/middleware/panic.go

@@ -7,15 +7,15 @@ import (
 )
 
 func NewPanic(next http.Handler, loggerFunc func(v ...interface{})) *Panic {
-	return  &Panic{
-		Next: next,
+	return &Panic{
+		Next:       next,
 		loggerFunc: loggerFunc,
 	}
 }
 
 // Panic recovers from API panics and logs encountered panics
 type Panic struct {
-	Next http.Handler
+	Next       http.Handler
 	loggerFunc func(v ...interface{})
 }
 

internal/middleware/panic_test.go

@@ -8,7 +8,7 @@ import (
 
 func TestPanic(t *testing.T) {
 	crashedHandler := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
-			panic("hello world")
+		panic("hello world")
 	})
 
 	middleware := NewPanic(crashedHandler, t.Log)

internal/roundtripper/cacher.go

@@ -3,8 +3,6 @@ package roundtripper
 import (
 	"crypto/md5"
 	"encoding/hex"
-	"errors"
-	"fmt"
 	"github.com/donutloop/httpcache/internal/cache"
 	"net/http"
 	"net/http/httputil"
@@ -24,7 +22,7 @@ func (t *CacheTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	if !ok {
 		proxyResponse, err := t.Transport.RoundTrip(req)
 		if err != nil {
-			return nil, errors.New(fmt.Sprintf("proxy couldn't forward request to destination server (%v)", err))
+			return nil, err
 		}
 		cachedResponse = &cache.CachedResponse{Resp: proxyResponse}
 		t.Cache.Set(clonedRequest, cachedResponse)

internal/roundtripper/reponse_body_limit.go

@@ -0,0 +1,26 @@
+package roundtripper
+
+import (
+	"errors"
+	"net/http"
+)
+
+var ResponseIsToLarge = errors.New("response body is to large for the cache")
+
+type ResponseBodyLimitRoundTripper struct {
+	Limit     int64
+	Transport http.RoundTripper // underlying transport (or default if nil)
+}
+
+func (t *ResponseBodyLimitRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	response, err := t.Transport.RoundTrip(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if response.ContentLength > t.Limit {
+		return nil, ResponseIsToLarge
+	}
+
+	return response, nil
+}

internal/size/size.go

@@ -0,0 +1,8 @@
+package size
+
+const (
+	_        = iota // ignore first value by assigning to blank identifier
+	KB int64 = 1 << (10 * iota)
+	MB
+	GB
+)

internal/xhttp/server.go

@@ -15,7 +15,7 @@ type Server struct {
 
 	ShutdownTimeout time.Duration
 
-	Logger       *log.Logger
+	Logger *log.Logger
 }
 
 // Start starts the server and waits for it to return.
@@ -50,4 +50,3 @@ func (s *Server) Stop() {
 
 	s.Server.Close()
 }
-