1 files changed, 21 insertions, 13 deletions
diff --git a/actionpack/lib/action_dispatch/http/content_security_policy.rb b/actionpack/lib/action_dispatch/http/content_security_policy.rb
index a3407c9698..35041fd072 100644
--- a/actionpack/lib/action_dispatch/http/content_security_policy.rb
+++ b/actionpack/lib/action_dispatch/http/content_security_policy.rb
@@ -21,13 +21,8 @@ module ActionDispatch #:nodoc:
         return response if policy_present?(headers)
 
         if policy = request.content_security_policy
-          if policy.directives["script-src"]
-            if nonce = request.content_security_policy_nonce
-              policy.directives["script-src"] << "'nonce-#{nonce}'"
-            end
-          end
-
-          headers[header_name(request)] = policy.build(request.controller_instance)
+          nonce = request.content_security_policy_nonce
+          headers[header_name(request)] = policy.build(request.controller_instance, nonce)
         end
 
         response
@@ -113,7 +108,9 @@ module ActionDispatch #:nodoc:
       blob:           "blob:",
       filesystem:     "filesystem:",
       report_sample:  "'report-sample'",
-      strict_dynamic: "'strict-dynamic'"
+      strict_dynamic: "'strict-dynamic'",
+      ws:             "ws:",
+      wss:            "wss:"
     }.freeze
 
     DIRECTIVES = {
@@ -129,12 +126,15 @@ module ActionDispatch #:nodoc:
       manifest_src:    "manifest-src",
       media_src:       "media-src",
       object_src:      "object-src",
+      prefetch_src:    "prefetch-src",
       script_src:      "script-src",
       style_src:       "style-src",
       worker_src:      "worker-src"
     }.freeze
 
-    private_constant :MAPPINGS, :DIRECTIVES
+    NONCE_DIRECTIVES = %w[script-src].freeze
+
+    private_constant :MAPPINGS, :DIRECTIVES, :NONCE_DIRECTIVES
 
     attr_reader :directives
 
@@ -203,8 +203,8 @@ module ActionDispatch #:nodoc:
       end
     end
 
-    def build(context = nil)
-      build_directives(context).compact.join("; ")
+    def build(context = nil, nonce = nil)
+      build_directives(context, nonce).compact.join("; ")
     end
 
     private
@@ -227,10 +227,14 @@ module ActionDispatch #:nodoc:
         end
       end
 
-      def build_directives(context)
+      def build_directives(context, nonce)
         @directives.map do |directive, sources|
           if sources.is_a?(Array)
-            "#{directive} #{build_directive(sources, context).join(' ')}"
+            if nonce && nonce_directive?(directive)
+              "#{directive} #{build_directive(sources, context).join(' ')} 'nonce-#{nonce}'"
+            else
+              "#{directive} #{build_directive(sources, context).join(' ')}"
+            end
           elsif sources
             directive
           else
@@ -259,5 +263,9 @@ module ActionDispatch #:nodoc:
           raise RuntimeError, "Unexpected content security policy source: #{source.inspect}"
         end
       end
+
+      def nonce_directive?(directive)
+        NONCE_DIRECTIVES.include?(directive)
+      end
   end
 end