From 349add70df550cd5cf1508d63f219b5d295696a6 Mon Sep 17 00:00:00 2001
From: Matthias Kurz <m.kurz@irregular.at>
Date: Tue, 12 Dec 2023 14:41:35 +0100
Subject: [PATCH] Support CSP nonce in Comet helpers

---
 .../src/main/java/play/libs/Comet.java        | 52 ++++++++++++++--
 .../src/main/scala/play/api/libs/Comet.scala  | 62 ++++++++++++++++---
 2 files changed, 101 insertions(+), 13 deletions(-)

diff --git a/core/play-java/src/main/java/play/libs/Comet.java b/core/play-java/src/main/java/play/libs/Comet.java
index 5b3cede5a7b..a271e5858b4 100644
--- a/core/play-java/src/main/java/play/libs/Comet.java
+++ b/core/play-java/src/main/java/play/libs/Comet.java
@@ -50,12 +50,24 @@ public abstract class Comet {
    * @return a flow of ByteString elements.
    */
   public static Flow<String, ByteString, NotUsed> string(String callbackName) {
+    return string(callbackName, null);
+  }
+
+  /**
+   * Produces a Flow of escaped ByteString from a series of String elements. Calls out to Comet.flow
+   * internally.
+   *
+   * @param callbackName the javascript callback method.
+   * @param nonce The CSP nonce to use for the script tag. If {@code null} no nonce will be sent.
+   * @return a flow of ByteString elements.
+   */
+  public static Flow<String, ByteString, NotUsed> string(String callbackName, String nonce) {
     return Flow.of(String.class)
         .map(
             str -> {
               return ByteString.fromString("'" + StringEscapeUtils.escapeEcmaScript(str) + "'");
             })
-        .via(flow(callbackName));
+        .via(flow(callbackName, nonce));
   }
 
   /**
@@ -66,12 +78,24 @@ public static Flow<String, ByteString, NotUsed> string(String callbackName) {
    * @return a flow of ByteString elements.
    */
   public static Flow<JsonNode, ByteString, NotUsed> json(String callbackName) {
+    return json(callbackName, null);
+  }
+
+  /**
+   * Produces a flow of ByteString using `Json.stringify` from a Flow of JsonNode. Calls out to
+   * Comet.flow internally.
+   *
+   * @param callbackName the javascript callback method.
+   * @param nonce The CSP nonce to use for the script tag. If {@code null} no nonce will be sent.
+   * @return a flow of ByteString elements.
+   */
+  public static Flow<JsonNode, ByteString, NotUsed> json(String callbackName, String nonce) {
     return Flow.of(JsonNode.class)
         .map(
             json -> {
               return ByteString.fromString(Json.stringify(json));
             })
-        .via(flow(callbackName));
+        .via(flow(callbackName, nonce));
   }
 
   /**
@@ -81,18 +105,36 @@ public static Flow<JsonNode, ByteString, NotUsed> json(String callbackName) {
    * @return a flow of ByteString elements.
    */
   public static Flow<ByteString, ByteString, NotUsed> flow(String callbackName) {
+    return flow(callbackName, null);
+  }
+
+  /**
+   * Produces a flow of ByteString with a prepended block and a script wrapper.
+   *
+   * @param callbackName the javascript callback method.
+   * @param nonce The CSP nonce to use for the script tag. If {@code null} no nonce will be sent.
+   * @return a flow of ByteString elements.
+   */
+  public static Flow<ByteString, ByteString, NotUsed> flow(String callbackName, String nonce) {
     ByteString cb = ByteString.fromString(callbackName);
     return Flow.of(ByteString.class)
         .map(
             (msg) -> {
-              return formatted(cb, msg);
+              return formatted(cb, msg, nonce);
             })
         .prepend(Source.single(initialChunk));
   }
 
-  private static ByteString formatted(ByteString callbackName, ByteString javascriptMessage) {
+  private static ByteString formatted(
+      ByteString callbackName, ByteString javascriptMessage, String nonce) {
     ByteStringBuilder b = new ByteStringBuilder();
-    b.append(ByteString.fromString("<script>"));
+    b.append(ByteString.fromString("<script"));
+    if (nonce != null && !nonce.isEmpty()) {
+      b.append(ByteString.fromString(" nonce=\""));
+      b.append(ByteString.fromString(nonce));
+      b.append(ByteString.fromString("\""));
+    }
+    b.append(ByteString.fromString(">"));
     b.append(callbackName);
     b.append(ByteString.fromString("("));
     b.append(javascriptMessage);
diff --git a/core/play/src/main/scala/play/api/libs/Comet.scala b/core/play/src/main/scala/play/api/libs/Comet.scala
index 5b4151b30a2..9ef6061c65e 100644
--- a/core/play/src/main/scala/play/api/libs/Comet.scala
+++ b/core/play/src/main/scala/play/api/libs/Comet.scala
@@ -43,10 +43,20 @@ object Comet {
    * @param callbackName the javascript callback method.
    * @return a flow of ByteString elements.
    */
-  def string(callbackName: String): Flow[String, ByteString, NotUsed] = {
+  def string(callbackName: String): Flow[String, ByteString, NotUsed] = string(callbackName, "")
+
+  /**
+   * Produces a Flow of escaped ByteString from a series of String elements.  Calls
+   * out to Comet.flow internally.
+   *
+   * @param callbackName the javascript callback method.
+   * @param nonce The CSP nonce to use for the script tag. If null no nonce will be sent.
+   * @return a flow of ByteString elements.
+   */
+  def string(callbackName: String, nonce: String): Flow[String, ByteString, NotUsed] = {
     Flow[String]
       .map(str => ByteString.fromString("'" + StringEscapeUtils.escapeEcmaScript(str) + "'"))
-      .via(flow(callbackName))
+      .via(flow(callbackName, nonce = nonce))
   }
 
   /**
@@ -56,10 +66,20 @@ object Comet {
    * @param callbackName the javascript callback method.
    * @return a flow of ByteString elements.
    */
-  def json(callbackName: String): Flow[JsValue, ByteString, NotUsed] = {
+  def json(callbackName: String): Flow[JsValue, ByteString, NotUsed] = json(callbackName, "")
+
+  /**
+   * Produces a flow of ByteString using `Json.fromJson(_).get` from a Flow of JsValue.  Calls
+   * out to Comet.flow internally.
+   *
+   * @param callbackName the javascript callback method.
+   * @param nonce The CSP nonce to use for the script tag. If null no nonce will be sent.
+   * @return a flow of ByteString elements.
+   */
+  def json(callbackName: String, nonce: String): Flow[JsValue, ByteString, NotUsed] = {
     Flow[JsValue]
       .map { msg => ByteString.fromString(Json.asciiStringify(msg)) }
-      .via(flow(callbackName))
+      .via(flow(callbackName, nonce = nonce))
   }
 
   /**
@@ -78,15 +98,41 @@ object Comet {
    */
   def flow(
       callbackName: String,
-      initialChunk: ByteString = initialByteString
+      initialChunk: ByteString
+  ): Flow[ByteString, ByteString, NotUsed] = flow(callbackName, initialChunk, "")
+
+  /**
+   * Creates a flow of ByteString.  Useful when you have objects that are not JSON or String where
+   * you may have to do your own conversion.
+   *
+   * Usage example:
+   *
+   * {{{
+   *   val htmlStream: Source[ByteString, ByteString, NotUsed] = Flow[Html].map { html =>
+   *     ByteString.fromString(html.toString())
+   *   }
+   *   ...
+   *   Ok.chunked(htmlStream via Comet.flow("parent.clockChanged"))
+   * }}}
+   */
+  def flow(
+      callbackName: String,
+      initialChunk: ByteString = initialByteString,
+      nonce: String = ""
   ): Flow[ByteString, ByteString, NotUsed] = {
     val cb: ByteString = ByteString.fromString(callbackName)
-    Flow.apply[ByteString].map(msg => formatted(cb, msg)).prepend(Source.single(initialChunk))
+    Flow.apply[ByteString].map(msg => formatted(cb, msg, nonce)).prepend(Source.single(initialChunk))
   }
 
-  private def formatted(callbackName: ByteString, javascriptMessage: ByteString): ByteString = {
+  private def formatted(callbackName: ByteString, javascriptMessage: ByteString, nonce: String = ""): ByteString = {
     val b: ByteStringBuilder = new ByteStringBuilder
-    b.append(ByteString.fromString("""<script>"""))
+    b.append(ByteString.fromString("""<script"""))
+    if (nonce != null && nonce.nonEmpty) {
+      b.append(ByteString.fromString(""" nonce=""""))
+      b.append(ByteString.fromString(nonce))
+      b.append(ByteString.fromString("""""""))
+    }
+    b.append(ByteString.fromString(""">"""))
     b.append(callbackName)
     b.append(ByteString.fromString("("))
     b.append(javascriptMessage)