Allow tidewave.ai in CSP if toolbar is enabled (#254)

Author: SteffenDE

lib/tidewave.ex

@@ -87,25 +87,34 @@ defmodule Tidewave do
   defp maybe_rewrite_csp(conn) do
     case Plug.Conn.get_resp_header(conn, "content-security-policy") do
       [csp | _] ->
-        csp = rewrite_csp(csp)
+        csp = rewrite_csp(conn, csp)
         Plug.Conn.put_resp_header(conn, "content-security-policy", csp)
 
       _ ->
         conn
     end
   end
 
-  defp rewrite_csp(csp) do
+  defp rewrite_csp(conn, csp) do
     policy_directives = String.split(csp, ";", trim: true)
 
+    toolbar_host =
+      case conn.private.tidewave_config do
+        %{toolbar: true} ->
+          Application.get_env(:tidewave, :client_url, "https://tidewave.ai") <> " "
+
+        _ ->
+          ""
+      end
+
     for policy_directive <- policy_directives,
         policy_directive = String.trim(policy_directive),
         not String.starts_with?(policy_directive, "frame-ancestors") do
       case String.split(policy_directive, " ", parts: 2) do
         ["script-src", directives] ->
           case :binary.match(directives, "'unsafe-eval'") do
-            :nomatch -> "script-src 'unsafe-eval' #{directives}"
-            _ -> "script-src #{directives}"
+            :nomatch -> "script-src #{toolbar_host}'unsafe-eval' #{directives}"
+            _ -> "script-src #{toolbar_host}#{directives}"
           end
 
         [policy, directives] ->

test/tidewave_test.exs

@@ -114,21 +114,36 @@ defmodule TidewaveTest do
     assert Plug.Conn.get_resp_header(conn, "x-frame-options") == []
   end
 
-  test "updates CSP header if set" do
+  test "updates CSP header if set (no toolbar)" do
     conn =
       conn(:get, "/foo")
       |> Plug.Conn.put_resp_header(
         "content-security-policy",
         "default-src 'self' http://example.com; connect-src 'none'; script-src 'self'; frame-ancestors 'none'"
       )
-      |> Tidewave.call(Tidewave.init([]))
+      |> Tidewave.call(Tidewave.init(toolbar: false))
       |> Plug.Conn.send_resp(200, "foo")
 
     assert Plug.Conn.get_resp_header(conn, "content-security-policy") == [
              "default-src 'self' http://example.com; connect-src 'none'; script-src 'unsafe-eval' 'self'"
            ]
   end
 
+  test "updates CSP header if set (toolbar)" do
+    conn =
+      conn(:get, "/foo")
+      |> Plug.Conn.put_resp_header(
+        "content-security-policy",
+        "default-src 'self' http://example.com; connect-src 'none'; script-src 'self'; frame-ancestors 'none'"
+      )
+      |> Tidewave.call(Tidewave.init(toolbar: true))
+      |> Plug.Conn.send_resp(200, "foo")
+
+    assert Plug.Conn.get_resp_header(conn, "content-security-policy") == [
+             "default-src 'self' http://example.com; connect-src 'none'; script-src https://tidewave.ai 'unsafe-eval' 'self'"
+           ]
+  end
+
   test "updates CSP headers with flags and trailing space" do
     conn =
       conn(:get, "/foo")
@@ -140,7 +155,7 @@ defmodule TidewaveTest do
       |> Plug.Conn.send_resp(200, "foo")
 
     assert Plug.Conn.get_resp_header(conn, "content-security-policy") ==
-             ["upgrade-insecure-requests; script-src 'unsafe-eval' 'self'; "]
+             ["upgrade-insecure-requests; script-src https://tidewave.ai 'unsafe-eval' 'self'; "]
   end
 
   describe "/mcp" do