diff options
| -rw-r--r-- | .gitignore | 1 | ||||
| -rw-r--r-- | lib/gitlab_keys.rb | 48 | ||||
| -rw-r--r-- | spec/gitlab_keys_spec.rb | 36 |
3 files changed, 71 insertions, 14 deletions
@@ -2,3 +2,4 @@ config.yml tmp/* *.log /*.log.* +authorized_keys.lock 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 |