http-loader-override-response-headers

cshum · Mar 14, 2024 · 419a20e · 419a20e
1 parent cc96a06
commit 419a20e
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git a/config/httpconfig.go b/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),

diff --git a/loader/httploader/httploader.go b/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, resp.Header.Get(key))
+					}
+				}
+			}
 		})
 		body := resp.Body
 		size, _ := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64)

diff --git a/loader/httploader/option.go b/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) {