Check that advisory lock prefix fits inside four bytes

brandur · brandur · commit f0496ade151d · 2024-07-06T20:40:28.000-07:00
Add an extra safety check that an advisory lock prefix fits inside of
four bytes, which is all that's reserved for one.
diff --git 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 = self.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
@@ -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
@@ -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 = {}
