fix: Use locking when creating new keypairs

Gemfile.lock

@@ -6,6 +6,7 @@ PATH
       activerecord (>= 6.0)
       jwt (~> 2.1)
       lockbox (~> 0.4)
+      with_advisory_lock (~> 4.6)
 
 GEM
   remote: https://rubygems.org/
@@ -134,6 +135,8 @@ GEM
     tzinfo (2.0.4)
       concurrent-ruby (~> 1.0)
     unicode-display_width (2.1.0)
+    with_advisory_lock (4.6.0)
+      activerecord (>= 4.2)
     zeitwerk (2.5.1)
 
 PLATFORMS
@@ -155,4 +158,4 @@ DEPENDENCIES
   timecop
 
 BUNDLED WITH
-   2.2.32
+   2.3.8

keypairs.gemspec

@@ -28,6 +28,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency 'activerecord', '>= 6.0'                     # Depend on activerecord as ORM
   spec.add_dependency 'jwt', '~> 2.1'                              # Working with JSON Web Tokens
   spec.add_dependency 'lockbox', '~> 0.4'                          # Encrypt and decrypt attributes
+  spec.add_dependency 'with_advisory_lock', '~> 4.6'               # Advisory database-level locks
 
   spec.add_development_dependency 'appraisal'                      # Test against multiple gem versions
   spec.add_development_dependency 'brakeman'                       # Static analysis security vulnerability scanner

lib/keypair.rb

@@ -2,6 +2,7 @@
 
 require 'lockbox'
 require 'jwt'
+require 'with_advisory_lock'
 
 # This class contains functionality needed for signing messages
 # and publishing JWK[s].
@@ -34,7 +35,7 @@
 # @attr [Time] not_before The time before which no payloads may be signed using the keypair.
 # @attr [Time] not_after The time after which no payloads may be signed using the keypair.
 # @attr [Time] expires_at The time after which the keypair may not be used for signature validation.
-class Keypair < ActiveRecord::Base
+class Keypair < ActiveRecord::Base # rubocop:disable Metrics/ClassLength
   ALGORITHM = 'RS256'
   ROTATION_INTERVAL = 1.month
 
@@ -55,12 +56,16 @@ class Keypair < ActiveRecord::Base
   #   Non-expired keypairs are considered valid and can be used to validate signatures and export public jwks.
   scope :valid, -> { where(arel_table[:expires_at].gt(Time.zone.now)) }
 
+  # @param create [Boolean] if true, will create a new current keypair if one doesn't exist
   # @return [Keypair] the keypair used to sign messages and autorotates if it has expired.
-  def self.current
-    order(not_before: :asc)
-      .where(arel_table[:not_before].lteq(Time.zone.now))
-      .where(arel_table[:not_after].gteq(Time.zone.now))
-      .last || create!
+  def self.current(create: true)
+    current = order(not_before: :asc)
+              .where(arel_table[:not_before].lteq(Time.zone.now))
+              .where(arel_table[:not_after].gteq(Time.zone.now))
+              .last
+    return current unless current.nil?
+
+    create_with_lock! if create
   end
 
   # The JWK Set of our valid keypairs.
@@ -87,11 +92,8 @@ def self.current
   # @see https://www.imsglobal.org/spec/security/v1p0/#h_key-set-url
   def self.keyset
     valid_keys = valid.order(not_before: :asc).to_a
-    # If we don't have any keys or if we don't have a future key (i.e. the last key is the current key)
-    while valid_keys.last.nil? || valid_keys.last.not_before <= Time.zone.now
-      # There is an automatic fallback to Time.zone.now if not_before is not set
-      valid_keys << create!(not_before: valid_keys.last&.not_after)
-    end
+
+    valid_keys = rotate_keys_with_lock! if valid_keys.last.nil? || valid_keys.last.not_before <= Time.zone.now
 
     {
       keys: valid_keys.map(&:public_jwk_export)
@@ -229,4 +231,33 @@ def expires_at_after_not_after
 
     errors.add(:expires_at, 'must be after not after')
   end
+
+  class << self
+    private
+
+    LOCK_NAME = 'keypairs:rotate'
+
+    def rotate_keys_with_lock!
+      with_advisory_lock(LOCK_NAME) do
+        valid_keys = valid.order(not_before: :asc).to_a
+
+        # If we don't have any keys or if we don't have a future key (i.e. the last key is the current key)
+        while valid_keys.last.nil? || valid_keys.last.not_before <= Time.zone.now
+          # There is an automatic fallback to Time.zone.now if not_before is not set
+          valid_keys << create!(not_before: valid_keys.last&.not_after)
+        end
+
+        valid_keys
+      end
+    end
+
+    def create_with_lock!
+      with_advisory_lock(LOCK_NAME) do
+        current = current(create: false)
+        next current if current
+
+        create!
+      end
+    end
+  end
 end

spec/models/keypair_spec.rb

@@ -138,6 +138,17 @@
         it 'returns an instance' do
           expect(described_class.current).to be_a described_class
         end
+
+        it 'only creates 1 keypair with multiple threads' do
+          expect do
+            threads = Array.new(3).map do |_i|
+              Thread.new do
+                expect(described_class.current).to be_a Keypair
+              end
+            end
+            threads.each(&:join)
+          end.to change(described_class, :count).by(1)
+        end
       end
 
       context 'with valid keypairs' do
@@ -193,6 +204,17 @@
         it 'contains the public_jwk_export of only the valid keypairs' do
           expect(subject[:keys]).to eq(expected)
         end
+
+        it 'only creates 1 keypair with multiple threads' do
+          expect do
+            threads = Array.new(3).map do |_i|
+              Thread.new do
+                expect(described_class.keyset[:keys].length).to eq 3
+              end
+            end
+            threads.each(&:join)
+          end.to change(described_class, :count).by(1)
+        end
       end
 
       context 'with past, present and future keypairs' do

spec/support/with_advisory_lock.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+RSpec.configure do |config|
+  config.before(:suite) do
+    ENV['FLOCK_DIR'] = Dir.mktmpdir
+  end
+
+  config.after(:suite) do
+    FileUtils.remove_entry_secure ENV['FLOCK_DIR']
+  end
+end