Merge branch 'authorized_keys_lock' into 'master'

Authorized keys lock

4 files changed, 74 insertions, 14 deletions

diff --git a/.gitignore b/.gitignore
index 664a12b..5eb0e6e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@ config.yml
 tmp/*
 *.log
 /*.log.*
+authorized_keys.lock
diff --git a/CHANGELOG b/CHANGELOG
index d033a6c..778ca4d 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,3 +1,6 @@
+v1.9.4
+  - Use lock file when modify authorized_keys
+
 v1.9.3
   - Ignore force push detection for new branch or branch remove push
 
diff --git a/lib/gitlab_keys.rb b/lib/gitlab_keys.rb
index cade09e..675f5e1 100644
--- a/lib/gitlab_keys.rb
+++ b/lib/gitlab_keys.rb
@@ -36,13 +36,15 @@ class GitlabKeys
   end
 
   def batch_add_keys
-    open(auth_file, 'a') do |file|
-      stdin.each_line do |input|
-        tokens = input.strip.split("\t")
-        abort("#{$0}: invalid input #{input.inspect}") unless tokens.count == 2
-        key_id, public_key = tokens
-        $logger.info "Adding key #{key_id} => #{public_key.inspect}"
-        file.puts(key_line(key_id, public_key))
+    lock do
+      open(auth_file, 'a') do |file|
+        stdin.each_line do |input|
+          tokens = input.strip.split("\t")
+          abort("#{$0}: invalid input #{input.inspect}") unless tokens.count == 2
+          key_id, public_key = tokens
+          $logger.info "Adding key #{key_id} => #{public_key.inspect}"
+          file.puts(key_line(key_id, public_key))
+        end
       end
     end
     true
@@ -57,15 +59,17 @@ class GitlabKeys
   end
 
   def rm_key
-    $logger.info "Removing key #{@key_id}"
-    Tempfile.open('authorized_keys') do |temp|
-      open(auth_file, 'r+') do |current|
-        current.each do |line|
-          temp.puts(line) unless line.include?("/bin/gitlab-shell #{@key_id}\"")
+    lock do
+      $logger.info "Removing key #{@key_id}"
+      Tempfile.open('authorized_keys') do |temp|
+        open(auth_file, 'r+') do |current|
+          current.each do |line|
+            temp.puts(line) unless line.include?("/bin/gitlab-shell #{@key_id}\"")
+          end
         end
+        temp.close
+        FileUtils.cp(temp.path, auth_file)
       end
-      temp.close
-      FileUtils.cp(temp.path, auth_file)
     end
     true
   end
@@ -74,4 +78,20 @@ class GitlabKeys
     open(auth_file, 'w') { |file| file.puts '# Managed by gitlab-shell' }
     true
   end
+
+
+  def lock(timeout = 10)
+    File.open(lock_file, "w+") do |f|
+      begin
+        f.flock File::LOCK_EX
+        Timeout::timeout(timeout) { yield }
+      ensure
+        f.flock File::LOCK_UN
+      end
+    end
+  end
+
+  def lock_file
+    @lock_file ||= File.join(ROOT_PATH, "authorized_keys.lock")
+  end
 end
diff --git a/spec/gitlab_keys_spec.rb b/spec/gitlab_keys_spec.rb
index 5bf4c04..2ef7241 100644
--- a/spec/gitlab_keys_spec.rb
+++ b/spec/gitlab_keys_spec.rb
@@ -145,6 +145,42 @@ describe GitlabKeys do
     end
   end
 
+  describe :lock do
+    it "should raise exception if operation lasts more then timeout" do
+      key = GitlabKeys.new
+      expect do
+        key.send :lock, 1 do
+          sleep 2
+        end
+      end.to raise_error
+    end
+
+    it "should actually lock file" do
+      $global = ""
+      key = GitlabKeys.new
+
+      thr1 = Thread.new do
+        key.send :lock do
+          # Put bigger sleep here to test if main thread will
+          # wait for lock file released before executing code
+          sleep 1
+          $global << "foo"
+        end
+      end
+
+      # make sure main thread start lock command after
+      # thread above
+      sleep 0.5
+
+      key.send :lock do
+        $global << "bar"
+      end
+
+      thr1.join
+      $global.should == "foobar"
+    end
+  end
+
   def build_gitlab_keys(*args)
     argv(*args)
     GitlabKeys.new