echo: allow extending CSP at request-timeecho/v0.15.0

1 files changed, 94 insertions, 15 deletions

diff --git a/echo/middleware/csp.go b/echo/middleware/csp.go
index 14047eb..644dfe1 100644
--- a/echo/middleware/csp.go
+++ b/echo/middleware/csp.go
@@ -2,6 +2,8 @@ package middleware
 
 import (
         "bytes"
+        "crypto/sha256"
+        "encoding/base64"
         "encoding/json"
         "fmt"
         "net/url"
@@ -13,7 +15,19 @@ import (
         "github.com/labstack/echo/v4/middleware"
 )
 
-const HeaderReportTo = "ReportTo"
+const (
+        HeaderReportTo       = "ReportTo"
+        cspExtendContextKey  = "__echomw_csp__csp_extend"
+        cspReplaceContextKey = "__echomw_csp__csp_replace"
+)
+
+func ReplaceCSP(c echo.Context, csp *ContentSecurityPolicyConfig) {
+        c.Set(cspReplaceContextKey, csp)
+}
+
+func ExtendCSP(c echo.Context, csp *ContentSecurityPolicyConfig) {
+        c.Set(cspExtendContextKey, csp)
+}
 
 type ContentSecurityPolicyConfig struct {
         Skipper                 middleware.Skipper
@@ -46,6 +60,53 @@ type ContentSecurityPolicyConfig struct {
         RequireTrustedTypesFor  []CSPDirective `csp:"require-trusted-types-for"` // experimental
 }
 
+func mergeField[T any](a []T, b []T) []T {
+        if b != nil {
+                v := make([]T, len(a)+len(b))
+                copy(v, a)
+                copy(v[len(a):], b)
+                return v
+        } else {
+                return a
+        }
+}
+
+// ExtendSimple returns a copy of the current policy extended with some
+// other policy. Boolean fields are not merged and will retain the value
+// of the base configuration.
+func (c *ContentSecurityPolicyConfig) ExtendSimple(o *ContentSecurityPolicyConfig) *ContentSecurityPolicyConfig {
+        return &ContentSecurityPolicyConfig{
+                Skipper:                 c.Skipper,
+                ReportOnly:              c.ReportOnly,
+                UpgradeInsecureRequests: c.UpgradeInsecureRequests,
+                BlockAllMixedContent:    c.BlockAllMixedContent,
+                DefaultSrc:              mergeField(c.DefaultSrc, o.DefaultSrc),
+                ChildSrc:                mergeField(c.ChildSrc, o.ChildSrc),
+                ConnectSrc:              mergeField(c.ConnectSrc, o.ConnectSrc),
+                FontSrc:                 mergeField(c.FontSrc, o.FontSrc),
+                FrameSrc:                mergeField(c.FrameSrc, o.FrameSrc),
+                ImageSrc:                mergeField(c.ImageSrc, o.ImageSrc),
+                ManifestSrc:             mergeField(c.ManifestSrc, o.ManifestSrc),
+                MediaSrc:                mergeField(c.MediaSrc, o.MediaSrc),
+                ObjectSrc:               mergeField(c.ObjectSrc, o.ObjectSrc),
+                ScriptSrc:               mergeField(c.ScriptSrc, o.ScriptSrc),
+                StyleSrc:                mergeField(c.StyleSrc, o.StyleSrc),
+                BaseUri:                 mergeField(c.BaseUri, o.BaseUri),
+                Sandbox:                 mergeField(c.Sandbox, o.Sandbox),
+                FormAction:              mergeField(c.FormAction, o.FormAction),
+                FrameAncestors:          mergeField(c.FrameAncestors, o.FrameAncestors),
+                ReportUri:               mergeField(c.ReportUri, o.ReportUri),
+                ReportTo:                mergeField(c.ReportTo, o.ReportTo),
+                PrefetchSrc:             mergeField(c.PrefetchSrc, o.PrefetchSrc),
+                ScriptSrcElem:           mergeField(c.ScriptSrcElem, o.ScriptSrcElem),
+                StyleSrcElem:            mergeField(c.StyleSrcElem, o.StyleSrcElem),
+                StyleSrcAttr:            mergeField(c.StyleSrcAttr, o.StyleSrcAttr),
+                WorkerSrc:               mergeField(c.WorkerSrc, o.WorkerSrc),
+                NavigateTo:              mergeField(c.NavigateTo, o.NavigateTo),
+                RequireTrustedTypesFor:  mergeField(c.RequireTrustedTypesFor, o.RequireTrustedTypesFor),
+        }
+}
+
 func (c *ContentSecurityPolicyConfig) String() string {
         st := reflect.TypeOf(*c)
         sv := reflect.ValueOf(*c)
@@ -135,6 +196,11 @@ func CSPShaString(size int, h string) CSPDirective {
         return CSPDirective(fmt.Sprintf("'sha%d-%s'", size, h))
 }
 
+func CSPSha256FromBytes(d []byte) CSPDirective {
+        s := sha256.Sum256(d)
+        return CSPShaString(256, base64.StdEncoding.EncodeToString(s[:]))
+}
+
 type CSPSandbox string
 
 const (
@@ -174,22 +240,35 @@ func ContentSecurityPolicyWithConfig(config ContentSecurityPolicyConfig) echo.Mi
                                 return next(c)
                         }
 
-                        h := c.Response().Header()
-                        if config.ReportOnly {
-                                h.Set(echo.HeaderContentSecurityPolicyReportOnly, config.String())
-                        } else {
-                                h.Set(echo.HeaderContentSecurityPolicy, config.String())
-                        }
+                        // This has to hook after the template runs but before the headers
+                        // are written because some template helper functions want to modify
+                        // the CSP state and if it renders too early that won't work.
+                        c.Response().Before(func() {
+                                liveConfig := &config
 
-                        if config.ReportTo != nil {
-                                rt := bytes.Buffer{}
-                                je := json.NewEncoder(&rt)
-                                for _, r := range config.ReportTo {
-                                        _ = je.Encode(r)
-                                        rt.WriteString(", ")
+                                if replace, ok := c.Get(cspReplaceContextKey).(*ContentSecurityPolicyConfig); ok {
+                                        liveConfig = replace
+                                } else if extend, ok := c.Get(cspExtendContextKey).(*ContentSecurityPolicyConfig); ok {
+                                        liveConfig = config.ExtendSimple(extend)
                                 }
-                                h.Set(HeaderReportTo, rt.String())
-                        }
+
+                                h := c.Response().Header()
+                                if liveConfig.ReportOnly {
+                                        h.Set(echo.HeaderContentSecurityPolicyReportOnly, liveConfig.String())
+                                } else {
+                                        h.Set(echo.HeaderContentSecurityPolicy, liveConfig.String())
+                                }
+
+                                if liveConfig.ReportTo != nil {
+                                        rt := bytes.Buffer{}
+                                        je := json.NewEncoder(&rt)
+                                        for _, r := range liveConfig.ReportTo {
+                                                _ = je.Encode(r)
+                                                rt.WriteString(", ")
+                                        }
+                                        h.Set(HeaderReportTo, rt.String())
+                                }
+                        })
 
                         return next(c)
                 }