feat: override response header option for HTTP Loader (#417)

* override header if specified

* http-loader-override-response-headers

* test cases

* cleanup

Author: cshum

blob.go

@@ -52,7 +52,8 @@ type Blob struct {
 	contentType   string
 	memory        *memory
 
-	Stat *Stat
+	Header http.Header
+	Stat   *Stat
 }
 
 // Stat Blob stat attributes

config/config_test.go

@@ -60,6 +60,7 @@ func TestBasic(t *testing.T) {
 		"-imagor-cache-header-ttl", "169h",
 		"-imagor-cache-header-swr", "167h",
 		"-http-loader-insecure-skip-verify-transport",
+		"-http-loader-override-response-headers", "cache-control,content-type",
 		"-http-loader-base-url", "https://www.example.com/foo.org",
 	})
 	app := srv.App.(*imagor.Imagor)
@@ -85,6 +86,7 @@ func TestBasic(t *testing.T) {
 	httpLoader := app.Loaders[0].(*httploader.HTTPLoader)
 	assert.True(t, httpLoader.Transport.(*http.Transport).TLSClientConfig.InsecureSkipVerify)
 	assert.Equal(t, "https://www.example.com/foo.org", httpLoader.BaseURL.String())
+	assert.Equal(t, []string{"cache-control", "content-type"}, httpLoader.OverrideResponseHeaders)
 }
 
 func TestVersion(t *testing.T) {

config/httpconfig.go

@@ -14,6 +14,8 @@ func withHTTPLoader(fs *flag.FlagSet, cb func() (*zap.Logger, bool)) imagor.Opti
 	var (
 		httpLoaderForwardHeaders = fs.String("http-loader-forward-headers", "",
 			"Forward request header to HTTP Loader request by csv e.g. User-Agent,Accept")
+		httpLoaderOverrideResponseHeaders = fs.String("http-loader-override-response-headers", "",
+			"Override HTTP Loader response header to image response by csv e.g. Cache-Control,Expires")
 		httpLoaderForwardClientHeaders = fs.Bool("http-loader-forward-client-headers", false,
 			"Forward browser client request headers to HTTP Loader request")
 		httpLoaderForwardAllHeaders = fs.Bool("http-loader-forward-all-headers", false,
@@ -58,6 +60,7 @@ func withHTTPLoader(fs *flag.FlagSet, cb func() (*zap.Logger, bool)) imagor.Opti
 						*httpLoaderForwardClientHeaders || *httpLoaderForwardAllHeaders),
 					httploader.WithAccept(*httpLoaderAccept),
 					httploader.WithForwardHeaders(*httpLoaderForwardHeaders),
+					httploader.WithOverrideResponseHeaders(*httpLoaderOverrideResponseHeaders),
 					httploader.WithAllowedSources(*httpLoaderAllowedSources),
 					httploader.WithAllowedSourceRegexps(*httpLoaderAllowedSourceRegexp),
 					httploader.WithMaxAllowedSize(*httpLoaderMaxAllowedSize),

imagor.go

@@ -203,6 +203,11 @@ func (app *Imagor) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if r.Header.Get("Imagor-Raw") != "" {
 		w.Header().Set("Content-Security-Policy", "script-src 'none'")
 	}
+	if h := blob.Header; h != nil {
+		for key := range h {
+			w.Header().Set(key, h.Get(key))
+		}
+	}
 	if checkStatNotModified(w, r, blob.Stat) {
 		w.WriteHeader(http.StatusNotModified)
 		return

imagor_test.go

@@ -331,6 +331,31 @@ func TestWithRaw(t *testing.T) {
 	assert.Equal(t, "bar", w.Header().Get("Content-Type"))
 }
 
+func TestWithOverrideHeader(t *testing.T) {
+	app := New(
+		WithDebug(true),
+		WithUnsafe(true),
+		WithLogger(zap.NewExample()),
+		WithLoaders(loaderFunc(func(r *http.Request, image string) (*Blob, error) {
+			blob := NewBlobFromBytes([]byte("foo"))
+			blob.SetContentType("bar")
+			blob.Header = make(http.Header)
+			blob.Header.Set("Content-Type", "tada")
+			blob.Header.Set("Foo", "bar")
+			blob.Header.Set("asdf", "fghj")
+			return blob, nil
+		})),
+	)
+	w := httptest.NewRecorder()
+	app.ServeHTTP(w, httptest.NewRequest(
+		http.MethodGet, "https://example.com/unsafe/filters:fill(red):raw()/gopher.png", nil))
+	assert.Equal(t, 200, w.Code)
+	assert.Equal(t, "foo", w.Body.String())
+	assert.Equal(t, "script-src 'none'", w.Header().Get("Content-Security-Policy"))
+	assert.Equal(t, "tada", w.Header().Get("Content-Type"))
+	assert.Equal(t, "fghj", w.Header().Get("ASDF"))
+}
+
 func TestNewBlobFromPathNotFound(t *testing.T) {
 	loader := loaderFunc(func(r *http.Request, image string) (*Blob, error) {
 		return NewBlobFromFile("./non-exists-path"), nil

loader/httploader/httploader.go

@@ -63,6 +63,9 @@ type HTTPLoader struct {
 	// OverrideHeaders override image request headers
 	OverrideHeaders map[string]string
 
+	// OverrideResponseHeaders override image response header from HTTP Loader response
+	OverrideResponseHeaders []string
+
 	// AllowedSources list of sources allowed to load from
 	AllowedSources []AllowedSource
 
@@ -204,6 +207,14 @@ func (h *HTTPLoader) Get(r *http.Request, image string) (*imagor.Blob, error) {
 		}
 		once.Do(func() {
 			blob.SetContentType(resp.Header.Get("Content-Type"))
+			if len(h.OverrideResponseHeaders) > 0 {
+				blob.Header = make(http.Header)
+				for _, key := range h.OverrideResponseHeaders {
+					if val := resp.Header.Get(key); val != "" {
+						blob.Header.Set(key, val)
+					}
+				}
+			}
 		})
 		body := resp.Body
 		size, _ := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64)

loader/httploader/httploader_test.go

@@ -23,7 +23,7 @@ func (t testTransport) RoundTrip(r *http.Request) (w *http.Response, err error)
 		w = &http.Response{
 			StatusCode: http.StatusOK,
 			Body:       io.NopCloser(strings.NewReader(res)),
-			Header:     map[string][]string{},
+			Header:     make(http.Header),
 		}
 		w.Header.Set("Content-Type", "image/jpeg")
 		return
@@ -48,6 +48,7 @@ type test struct {
 	name   string
 	target string
 	result string
+	header map[string]string
 	err    string
 }
 
@@ -80,6 +81,11 @@ func doTests(t *testing.T, loader imagor.Loader, tests []test) {
 				}
 				assert.Equal(t, tt.err, msg)
 			}
+			if tt.header != nil {
+				for key, val := range tt.header {
+					assert.Equal(t, val, b.Header.Get(key))
+				}
+			}
 		})
 	}
 }
@@ -492,6 +498,31 @@ func TestWithForwardHeadersOverrideUserAgent(t *testing.T) {
 	})
 }
 
+func TestWithOverrideResponseHeader(t *testing.T) {
+	doTests(t, New(
+		WithTransport(roundTripFunc(func(r *http.Request) (w *http.Response, err error) {
+			res := &http.Response{
+				StatusCode: http.StatusOK,
+				Header:     map[string][]string{},
+				Body:       io.NopCloser(strings.NewReader("ok")),
+			}
+			res.Header.Set("Content-Type", "image/jpeg")
+			res.Header.Set("Foo", "Bar")
+			return res, nil
+		})),
+		WithOverrideResponseHeaders("foo"),
+	), []test{
+		{
+			name:   "user agent",
+			target: "https://foo.bar/baz",
+			result: "ok",
+			header: map[string]string{
+				"Foo": "Bar",
+			},
+		},
+	})
+}
+
 func TestWithForwardClientHeaders(t *testing.T) {
 	doTests(t, New(
 		WithTransport(roundTripFunc(func(r *http.Request) (w *http.Response, err error) {

loader/httploader/option.go

@@ -59,6 +59,21 @@ func WithForwardHeaders(headers ...string) Option {
 	}
 }
 
+// WithOverrideResponseHeaders with override selected response headers option
+func WithOverrideResponseHeaders(headers ...string) Option {
+	return func(h *HTTPLoader) {
+		for _, raw := range headers {
+			splits := strings.Split(raw, ",")
+			for _, header := range splits {
+				header = strings.TrimSpace(header)
+				if len(header) > 0 {
+					h.OverrideResponseHeaders = append(h.OverrideResponseHeaders, header)
+				}
+			}
+		}
+	}
+}
+
 // WithForwardClientHeaders with forward browser request headers option
 func WithForwardClientHeaders(enabled bool) Option {
 	return func(h *HTTPLoader) {