feat(max-allowed-size): add new option max-allowed-size in bytes (#111)

* feat(max-allowed-size): add new option max-allowed-size in bytes

* fix(max-allowed-size): HEAD response handling

- consider 200~206 as valid HEAD response codes
- do not defer res.Body.Close() of HEAD request

fixtures/1024bytes

@@ -0,0 +1 @@
+1111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111

imaginary.go

@@ -33,6 +33,7 @@ var (
 	aEnableURLSource   = flag.Bool("enable-url-source", false, "Enable remote HTTP URL image source processing")
 	aEnablePlaceholder = flag.Bool("enable-placeholder", false, "Enable image response placeholder to be used in case of error")
 	aAlloweOrigins     = flag.String("allowed-origins", "", "Restrict remote image source processing to certain origins (separated by commas)")
+	aMaxAllowedSize    = flag.Int("max-allowed-size", 0, "Restrict maximum size of http image source (in bytes)")
 	aKey               = flag.String("key", "", "Define API key for authorization")
 	aMount             = flag.String("mount", "", "Mount server local directory")
 	aCertFile          = flag.String("certfile", "", "TLS certificate file path")
@@ -81,6 +82,7 @@ Options:
 	-enable-placeholder       Enable image response placeholder to be used in case of error [default: false]
   -enable-auth-forwarding   Forwards X-Forward-Authorization or Authorization header to the image source server. -enable-url-source flag must be defined. Tip: secure your server from public access to prevent attack vectors
   -allowed-origins <urls>   Restrict remote image source processing to certain origins (separated by commas)
+  -max-allowed-size <bytes> Restrict maximum size of http image source (in bytes)
   -certfile <path>          TLS certificate file path
   -keyfile <path>           TLS private key file path
   -authorization <value>    Defines a constant Authorization header value passed to all the image source servers. -enable-url-source flag must be defined. This overwrites authorization headers forwarding behavior via X-Forward-Authorization
@@ -130,6 +132,7 @@ func main() {
 		HttpWriteTimeout:  *aWriteTimeout,
 		Authorization:     *aAuthorization,
 		AlloweOrigins:     parseOrigins(*aAlloweOrigins),
+		MaxAllowedSize:    *aMaxAllowedSize,
 	}
 
 	// Create a memory release goroutine

server.go

@@ -31,6 +31,7 @@ type ServerOptions struct {
 	Placeholder       string
 	PlaceholderImage  []byte
 	AlloweOrigins     []*url.URL
+	MaxAllowedSize    int
 }
 
 func Server(o ServerOptions) error {

source.go

@@ -14,6 +14,7 @@ type SourceConfig struct {
 	MountPath       string
 	Type            ImageSourceType
 	AllowedOrigings []*url.URL
+	MaxAllowedSize  int
 }
 
 var imageSourceMap = make(map[ImageSourceType]ImageSource)
@@ -36,6 +37,7 @@ func LoadSources(o ServerOptions) {
 			AuthForwarding:  o.AuthForwarding,
 			Authorization:   o.Authorization,
 			AllowedOrigings: o.AlloweOrigins,
+			MaxAllowedSize:  o.MaxAllowedSize,
 		})
 	}
 }

source_http.go

@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"strconv"
 )
 
 const ImageSourceTypeHttp ImageSourceType = "http"
@@ -33,14 +34,26 @@ func (s *HttpImageSource) GetImage(req *http.Request) ([]byte, error) {
 }
 
 func (s *HttpImageSource) fetchImage(url *url.URL, ireq *http.Request) ([]byte, error) {
-	req := newHTTPRequest(url)
+	// Check remote image size by fetching HTTP Headers
+	if s.Config.MaxAllowedSize > 0 {
+		req := newHTTPRequest(s, ireq, "HEAD", url)
+		res, err := http.DefaultClient.Do(req)
+		if err != nil {
+			return nil, fmt.Errorf("Error fetching image http headers: %v", err)
+		}
+		res.Body.Close()
+		if res.StatusCode >= 200 && res.StatusCode <= 206 {
+			return nil, fmt.Errorf("Error fetching image http headers: (status=%d) (url=%s)", res.StatusCode, req.URL.String())
+		}
 
-	// Forward auth header to the target server, if necessary
-	if s.Config.AuthForwarding || s.Config.Authorization != "" {
-		s.setAuthorizationHeader(req, ireq)
+		contentLength, _ := strconv.Atoi(res.Header.Get("Content-Length"))
+		if contentLength > s.Config.MaxAllowedSize {
+			return nil, fmt.Errorf("Content-Length %d exceeds maximum allowed %d bytes", contentLength, s.Config.MaxAllowedSize)
+		}
 	}
 
 	// Perform the request using the default client
+	req := newHTTPRequest(s, ireq, "GET", url)
 	res, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return nil, fmt.Errorf("Error downloading image: %v", err)
@@ -76,10 +89,16 @@ func parseURL(request *http.Request) (*url.URL, error) {
 	return url.Parse(queryUrl)
 }
 
-func newHTTPRequest(url *url.URL) *http.Request {
-	req, _ := http.NewRequest("GET", url.String(), nil)
+func newHTTPRequest(s *HttpImageSource, ireq *http.Request, method string, url *url.URL) *http.Request {
+	req, _ := http.NewRequest(method, url.String(), nil)
 	req.Header.Set("User-Agent", "imaginary/"+Version)
 	req.URL = url
+
+	// Forward auth header to the target server, if necessary
+	if s.Config.AuthForwarding || s.Config.Authorization != "" {
+		s.setAuthorizationHeader(req, ireq)
+	}
+
 	return req
 }
 

source_http_test.go

@@ -9,6 +9,7 @@ import (
 )
 
 const fixtureImage = "fixtures/large.jpg"
+const fixture1024Bytes = "fixtures/1024bytes"
 
 func TestHttpImageSource(t *testing.T) {
 	var body []byte
@@ -149,3 +150,33 @@ func TestHttpImageSourceError(t *testing.T) {
 	w := httptest.NewRecorder()
 	fakeHandler(w, r)
 }
+
+func TestHttpImageSourceExceedsMaximumAllowedLength(t *testing.T) {
+	var body []byte
+	var err error
+
+	buf, _ := ioutil.ReadFile(fixture1024Bytes)
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write(buf)
+	}))
+	defer ts.Close()
+
+	source := NewHttpImageSource(&SourceConfig{
+		MaxAllowedSize: 1023,
+	})
+	fakeHandler := func(w http.ResponseWriter, r *http.Request) {
+		if !source.Matches(r) {
+			t.Fatal("Cannot match the request")
+		}
+
+		body, err = source.GetImage(r)
+		if err == nil {
+			t.Fatalf("It should not allow a request to image exceeding maximum allowed size: %s", err)
+		}
+		w.Write(body)
+	}
+
+	r, _ := http.NewRequest("GET", "http://foo/bar?url="+ts.URL, nil)
+	w := httptest.NewRecorder()
+	fakeHandler(w, r)
+}