From 845560127c1ad56c3f0cbb48325672ac0c1c1bd6 Mon Sep 17 00:00:00 2001
From: Brandur <brandur@brandur.org>
Date: Sat, 6 Jul 2024 20:37:18 -0700
Subject: [PATCH] Check that advisory lock prefix fits inside four bytes

Add an extra safety check that an advisory lock prefix fits inside of
four bytes, which is all that's reserved for one.
---
 CHANGELOG.md        |  2 ++
 lib/client.rb       | 12 +++++++++++-
 sig/client.rbs      |  8 +++++---
 spec/client_spec.rb | 16 ++++++++++++++++
 4 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index edaf68a..d72f83b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- Advisory lock prefixes are now checked to make sure they fit inside of four bytes. [PR #24](https://github.com/riverqueue/riverqueue-ruby/pull/24).
+
 ## [0.5.0] - 2024-07-05
 
 ### Changed
diff --git a/lib/client.rb b/lib/client.rb
index 42dd5da..ac2d222 100644
--- a/lib/client.rb
+++ b/lib/client.rb
@@ -26,7 +26,7 @@ module River
   class Client
     def initialize(driver, advisory_lock_prefix: nil)
       @driver = driver
-      @advisory_lock_prefix = advisory_lock_prefix
+      @advisory_lock_prefix = check_advisory_lock_prefix_bounds(advisory_lock_prefix)
       @time_now_utc = -> { Time.now.utc } # for test time stubbing
     end
 
@@ -137,6 +137,16 @@ def insert_many(args)
       @driver.job_insert_many(all_params)
     end
 
+    private def check_advisory_lock_prefix_bounds(advisory_lock_prefix)
+      return nil if advisory_lock_prefix.nil?
+
+      # 2**32-1 is 0xffffffff (the largest number that's four bytes)
+      if advisory_lock_prefix < 0 || advisory_lock_prefix > 2**32 - 1
+        raise ArgumentError, "advisory lock prefix must fit inside four bytes"
+      end
+      advisory_lock_prefix
+    end
+
     # Default states that are used during a unique insert. Can be overridden by
     # setting UniqueOpts#by_state.
     DEFAULT_UNIQUE_STATES = [
diff --git a/sig/client.rbs b/sig/client.rbs
index 99a7eae..aaa01b7 100644
--- a/sig/client.rbs
+++ b/sig/client.rbs
@@ -8,13 +8,15 @@ module River
     @driver: _Driver
     @time_now_utc: ^() -> Time
 
-    DEFAULT_UNIQUE_STATES: Array[jobStateAll]
-    EMPTY_INSERT_OPTS: InsertOpts
-
     def initialize: (_Driver driver, ?advisory_lock_prefix: Integer?) -> void
     def insert: (jobArgs, ?insert_opts: InsertOpts) -> InsertResult
     def insert_many: (Array[jobArgs | InsertManyParams]) -> Integer
 
+    private def check_advisory_lock_prefix_bounds: (Integer?) -> Integer?
+
+    DEFAULT_UNIQUE_STATES: Array[jobStateAll]
+    EMPTY_INSERT_OPTS: InsertOpts
+
     private def check_unique_job: (Driver::JobInsertParams, UniqueOpts?) { () -> InsertResult } -> InsertResult
     private def make_insert_params: (jobArgs, InsertOpts, ?is_insert_many: bool) -> [Driver::JobInsertParams, UniqueOpts?]
     private def truncate_time: (Time, Integer) -> Time
diff --git a/spec/client_spec.rb b/spec/client_spec.rb
index 7b4a7b3..9129525 100644
--- a/spec/client_spec.rb
+++ b/spec/client_spec.rb
@@ -173,6 +173,22 @@ class SimpleArgsWithInsertOpts < SimpleArgs
       )
     end
 
+    it "errors if advisory lock prefix is larger than four bytes" do
+      River::Client.new(mock_driver, advisory_lock_prefix: 123)
+
+      expect do
+        River::Client.new(mock_driver, advisory_lock_prefix: -1)
+      end.to raise_error(ArgumentError, "advisory lock prefix must fit inside four bytes")
+
+      # 2^32-1 is 0xffffffff (1s for 32 bits) which fits
+      River::Client.new(mock_driver, advisory_lock_prefix: 2**32 - 1)
+
+      # 2^32 is 0x100000000, which does not
+      expect do
+        River::Client.new(mock_driver, advisory_lock_prefix: 2**32)
+      end.to raise_error(ArgumentError, "advisory lock prefix must fit inside four bytes")
+    end
+
     it "errors if args don't respond to #kind" do
       args_klass = Class.new do
         def to_json = {}