Lock reads and read-modify-write in puppetca token storage

modules/puppetca_token_whitelisting/puppetca_token_whitelisting_token_storage.rb

@@ -16,11 +16,11 @@ def ensure_file
     end
 
     def read
-      YAML.safe_load File.read @tokens_file
+      lock(File::LOCK_SH) { parse }
     end
 
     def write(content)
-      lock do
+      lock(File::LOCK_EX) do
         unsafe_write content
       end
     end
@@ -29,27 +29,40 @@ def unsafe_write(content)
       File.write @tokens_file, content.to_yaml
     end
 
-    def lock(&block)
-      File.open(@tokens_file, "r+") do |f|
-        f.flock File::LOCK_EX
+    def lock(mode = File::LOCK_EX, &block)
+      File.open(@tokens_file, ((mode == File::LOCK_EX) ? 'r+' : 'r')) do |f|
+        f.flock mode
         yield
       ensure
         f.flock File::LOCK_UN
       end
     end
 
     def add(entry)
-      write read.push entry
+      modify { |tokens| tokens << entry }
     end
 
     def remove(entry)
-      write read.delete_if { |data| data == entry }
+      modify { |tokens| tokens.delete_if { |data| data == entry } }
     end
 
     def remove_if(&block)
-      lock do
-        unsafe_write read.delete_if { |token| yield(token) }
+      modify { |tokens| tokens.delete_if(&block) }
+    end
+
+    private
+
+    # Read-modify-write under a single exclusive lock so concurrent writers
+    # cannot overwrite each other's updates (last writer used to win).
+    def modify
+      lock(File::LOCK_EX) do
+        unsafe_write yield(parse)
       end
     end
+
+    # An empty file (crash mid-write) reads as "no tokens" instead of nil.
+    def parse
+      YAML.safe_load(File.read(@tokens_file)) || []
+    end
   end
 end

test/puppetca_token_whitelisting/puppetca_token_whitelisting_token_storage_test.rb

@@ -60,4 +60,24 @@ def test_remove_if_works
     end
     assert_equal ['test.bar.example.com'], @storage.read
   end
+
+  def test_should_read_empty_file_as_empty_array
+    File.write @file.path, ''
+    assert_equal [], @storage.read
+  end
+
+  def test_should_add_to_empty_file_without_crashing
+    File.write @file.path, ''
+    @storage.add 'baz.example.com'
+    assert_equal ['baz.example.com'], @storage.read
+  end
+
+  def test_should_not_lose_concurrent_adds
+    entries = (1..20).map { |i| "token-#{i}.example.com" }
+    entries.map { |entry| Thread.new { @storage.add entry } }.each(&:join)
+    stored = @storage.read
+    entries.each do |entry|
+      assert_includes stored, entry, "concurrent add of #{entry} was lost"
+    end
+  end
 end