Fix incorrect escaping of large payloads

Gemfile

@@ -4,6 +4,8 @@ source "https://rubygems.org"
 
 gemspec
 
+gem "rake", "~> 13"
+
 group :development, :test do
   gem "activerecord"
   gem "debug"

README.md

@@ -58,4 +58,13 @@ Deletion of stale payloads (2 minutes or older) are triggered every 100 large pa
 
 - Clone repo
 - `bundle install` to install dependencies
-- `bundle exec ruby test/postgresql_test.rb` to run tests
+- Ensure you have docker-engine and run
+  ```
+  bundle exec rake
+  ```
+
+- Alternatively
+  ```
+  bundle exec rake test
+  ```
+- See `Rakefile` for environment variables that you can override

Rakefile

@@ -0,0 +1,69 @@
+require "pg"
+
+def postgres_ready?(url, timeout:, retry_interval: 0.1)
+  start_time = Time.now
+
+  loop do
+    conn = PG.connect(url)
+    conn.exec("SELECT 1")
+    conn.close
+    return true
+  rescue PG::Error => e
+    if Time.now - start_time > timeout
+      puts "Failed to connect to PostgreSQL: #{e.message}"
+      return false
+    else
+      sleep retry_interval
+    end
+  end
+end
+
+ENV["DB_USERNAME"] ||= "actioncable"
+ENV["DB_PASSWORD"] ||= "actioncable"
+ENV["DB_HOST"] ||= "localhost"
+ENV["DB_PORT"] ||= "5432"
+
+ENV["DB_URL"] ||= begin
+  username = ENV.fetch("DB_USERNAME", "")
+  password = ENV.fetch("DB_PASSWORD", "")
+  host = ENV.fetch("DB_HOST", "localhost")
+  port = ENV.fetch("DB_PORT", "5432")
+
+  url = "postgresql://"
+  unless username.empty?
+    url += username
+    url += ":#{password}" unless password.empty?
+    url += "@"
+  end
+  url += host
+  url += ":#{port}" unless port.empty?
+  url
+end
+
+namespace :docker do
+  task :up do
+    `docker compose up -d`
+
+    timeout = 10
+    puts "* Waiting for PostgreSQL database to accepting connections"
+    if postgres_ready?(ENV.fetch("DB_URL"), timeout: timeout)
+      puts "* PostgreSQL database is up and accepting connections"
+    else
+      puts "* PostgreSQL database is not ready after #{timeout} seconds"
+      `docker compose down -v`
+      exit 1
+    end
+  end
+
+  task :down do
+    `docker compose down -v`
+  end
+end
+
+task :test do
+  system("bundle", "exec", "ruby", "test/postgresql_test.rb")
+end
+
+task "test:docker" => ["docker:down", "docker:up", "test", "docker:down"]
+
+task default: "test:docker"

docker-compose.yml

@@ -0,0 +1,10 @@
+version: "3"
+services:
+  db:
+    image: postgres:13.2
+    environment:
+      - POSTGRES_USER=${DB_USERNAME}
+      - POSTGRES_PASSWORD=${DB_PASSWORD}
+    ports:
+      - "127.0.0.1:${DB_PORT}:5432"
+

lib/action_cable/subscription_adapter/enhanced_postgresql.rb

@@ -38,7 +38,6 @@ def broadcast(channel, payload)
 
         with_broadcast_connection do |pg_conn|
           channel = pg_conn.escape_identifier(channel_identifier(channel))
-          payload = pg_conn.escape_string(payload)
 
           if payload.bytesize > MAX_NOTIFY_SIZE
             payload_id = insert_large_payload(pg_conn, payload)
@@ -53,7 +52,7 @@ def broadcast(channel, payload)
             payload = "#{LARGE_PAYLOAD_PREFIX}#{encrypted_payload_id}"
           end
 
-          pg_conn.exec("NOTIFY #{channel}, '#{payload}'")
+          pg_conn.exec_params("NOTIFY #{channel}, #{pg_conn.escape_literal(payload)}")
         end
       end
 

test/postgresql_test.rb

@@ -13,7 +13,15 @@ class PostgresqlAdapterTest < ActionCable::TestCase
   include ChannelPrefixTest
 
   def setup
-    database_config = { "adapter" => "postgresql", "database" => "actioncable_enhanced_postgresql_test" }
+    database_config = {
+      "adapter" => "postgresql",
+      "database" => "actioncable_enhanced_postgresql_test",
+      "host": ENV.fetch("DB_HOST"),
+      "username": ENV.fetch("DB_USERNAME"),
+      "password": ENV.fetch("DB_PASSWORD")
+    }
+
+    database_config.merge(port: ENV["DB_PORT"]) unless ENV["DB_PORT"].empty?
 
     # Create the database unless it already exists
     begin
@@ -113,6 +121,38 @@ def test_large_payload_broadcast
     end
   end
 
+  # Postgres has a NOTIFY payload limit of 8000 bytes which requires special handling to avoid
+  # "PG::InvalidParameterValue: ERROR: payload string too long" errors.
+  def test_large_payload_escapes_correctly
+    large_payloads_table = ActionCable::SubscriptionAdapter::EnhancedPostgresql::LARGE_PAYLOADS_TABLE
+    ActiveRecord::Base.connection_pool.with_connection do |connection|
+      connection.execute("DROP TABLE IF EXISTS #{large_payloads_table}")
+    end
+
+    server = ActionCable::Server::Base.new
+    server.config.cable = cable_config.with_indifferent_access
+    server.config.logger = Logger.new(StringIO.new).tap { |l| l.level = Logger::UNKNOWN }
+    adapter = server.config.pubsub_adapter.new(server)
+
+    ascii_string = (32..126).map(&:chr).join.encode("UTF-8")
+
+    expected_length = (ActionCable::SubscriptionAdapter::EnhancedPostgresql::MAX_NOTIFY_SIZE + 1)
+
+    large_payload = (ascii_string * (1.0 * expected_length/ascii_string.length).ceil)[...expected_length]
+
+    subscribe_as_queue("channel", adapter) do |queue|
+      adapter.broadcast("channel", large_payload)
+
+      # The large payload is stored in the database at this point
+      assert_equal 1, ActiveRecord::Base.connection.query("SELECT COUNT(*) FROM #{large_payloads_table}").first.first
+
+      got = queue.pop
+
+      assert_equal large_payload.length, got.length, "Expected lengths to match"
+      assert_equal large_payload, got, "Expected values to match"
+    end
+  end
+
   def test_automatic_payload_deletion
     inserts_per_delete = ActionCable::SubscriptionAdapter::EnhancedPostgresql::INSERTS_PER_DELETE
     large_payloads_table = ActionCable::SubscriptionAdapter::EnhancedPostgresql::LARGE_PAYLOADS_TABLE
@@ -157,11 +197,11 @@ def test_explicit_url_configuration
     rescue ActiveRecord::DatabaseAlreadyExists
     end
 
-    pg_conn = PG::Connection.open(dbname: explicit_database)
+    pg_conn = PG::Connection.new("#{ENV.fetch("DB_URL")}/#{explicit_database}")
     pg_conn.exec("DROP TABLE IF EXISTS #{large_payloads_table}")
 
     server = ActionCable::Server::Base.new
-    server.config.cable = cable_config.merge(url: "postgres://localhost/#{explicit_database}").with_indifferent_access
+    server.config.cable = cable_config.merge(url: "#{ENV.fetch("DB_URL")}/#{explicit_database}").with_indifferent_access
     server.config.logger = Logger.new(StringIO.new).tap { |l| l.level = Logger::UNKNOWN }
     adapter = server.config.pubsub_adapter.new(server)
 

test/test_helper.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+$LOAD_PATH.unshift File.expand_path("../lib", __dir__)
+
 require "debug"
 
 require "active_support/testing/strict_warnings"