From c0e8a525bdc03220bee24fc3985f0d83fdc52da1 Mon Sep 17 00:00:00 2001
From: Kyle Banker <kylebanker@gmail.com>
Date: Thu, 25 Feb 2010 14:58:32 -0500
Subject: [PATCH] reauthenticate on reconnect

---
 Rakefile                              |  5 ++
 lib/mongo/connection.rb               | 92 +++++++++++++++++++--------
 lib/mongo/db.rb                       | 24 +++++--
 lib/mongo/gridfs/grid_io.rb           |  2 +-
 test/auxillary/authentication_test.rb | 68 ++++++++++++++++++++
 test/auxillary/autoreconnect_test.rb  |  3 +-
 test/connection_test.rb               | 36 ++++++++++-
 test/unit/connection_test.rb          |  7 +-
 8 files changed, 198 insertions(+), 39 deletions(-)
 create mode 100644 test/auxillary/authentication_test.rb

diff --git a/Rakefile b/Rakefile
index 0e3464a..c3ae028 100644
--- a/Rakefile
+++ b/Rakefile
@@ -81,6 +81,11 @@ namespace :test do
     t.verbose    = true
   end
 
+  Rake::TestTask.new(:authentication) do |t|
+    t.test_files = FileList['test/auxillary/authentication_test.rb']
+    t.verbose    = true
+  end
+
   task :drop_databases do |t|
     puts "Dropping test database..."
     require File.join(File.dirname(__FILE__), 'lib', 'mongo')
diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb
index 346fdcc..4b8a369 100644
--- a/lib/mongo/connection.rb
+++ b/lib/mongo/connection.rb
@@ -90,8 +90,10 @@ module Mongo
     #
     # @core connections
     def initialize(pair_or_host=nil, port=nil, options={})
+      @auths        = []
+
       if block_given?
-        @nodes, @auths = yield self
+        @nodes = yield self
       else
         @nodes = format_pair(pair_or_host, port)
       end
@@ -126,10 +128,7 @@ module Mongo
       @options  = options
 
       should_connect = options[:connect].nil? ? true : options[:connect]
-      if should_connect
-        connect_to_master
-        authenticate_databases if @auths
-      end
+      connect_to_master if should_connect
     end
 
     # Initialize a paired connection to MongoDB.
@@ -154,7 +153,7 @@ module Mongo
       # Block returns an array, the first element being an array of nodes and the second an array
       # of authorizations for the database.
       new(nil, nil, opts) do |con|
-        [[con.pair_val_to_connection(nodes[0]), con.pair_val_to_connection(nodes[1])], []]
+        [con.pair_val_to_connection(nodes[0]), con.pair_val_to_connection(nodes[1])]
       end
     end
 
@@ -172,12 +171,68 @@ module Mongo
       end
     end
 
+    # Apply each of the saved database authentications.
+    #
+    # @return [Boolean] returns true if authentications exist and succeeed, false
+    #   if none exists.
+    #
+    # @raise [AuthenticationError] raises an exception if any one
+    #   authentication fails.
+    def apply_saved_authentication
+      return false if @auths.empty?
+      @auths.each do |auth|
+        self[auth['db_name']].authenticate(auth['username'], auth['password'], false)
+      end
+      true
+    end
+
+    # Save an authentication to this connection. When connecting,
+    # the connection will attempt to re-authenticate on every db
+    # specificed in the list of auths.
+    #
+    # @param [String] db_name
+    # @param [String] username
+    # @param [String] password
+    #
+    # @return [Hash] a hash representing the authentication just added.
+    def add_auth(db_name, username, password)
+      remove_auth(db_name)
+      auth = {}
+      auth['db_name']  = db_name
+      auth['username'] = username
+      auth['password'] = password
+      @auths << auth
+      auth
+    end
+
+    # Remove a saved authentication for this connection.
+    #
+    # @param [String] db_name
+    #
+    # @return [Boolean]
+    def remove_auth(db_name)
+      return unless @auths
+      if @auths.reject! { |a| a['db_name'] == db_name }
+        true
+      else
+        false
+      end
+    end
+
+    # Remove all authenication information stored in this connection.
+    #
+    # @return [true] this operation return true because it always succeeds.
+    def clear_auths
+      @auths = []
+      true
+    end
+
     # Return a hash with all database names
     # and their respective sizes on disk.
     #
     # @return [Hash]
     def database_info
-      doc = self['admin'].command(:listDatabases => 1)
+      doc = self['admin'].command({:listDatabases => 1}, false, true)
       returning({}) do |info|
         doc['databases'].each { |db| info[db['name']] = db['sizeOnDisk'].to_i }
       end
@@ -232,7 +287,7 @@ module Mongo
       oh[:fromhost] = from_host
       oh[:fromdb]   = from
       oh[:todb]     = to
-      self["admin"].command(oh)
+      self["admin"].command(oh, false, true)
     end
 
     # Increment and return the next available request id.
@@ -250,7 +305,7 @@ module Mongo
     #
     # @return [Hash]
     def server_info
-      db("admin").command({:buildinfo => 1}, {:admin => true, :check_response => true})
+      self["admin"].command({:buildinfo => 1}, false, true)
     end
 
     # Get the build version of the current server.
@@ -365,6 +420,7 @@ module Mongo
           result = self['admin'].command({:ismaster => 1}, false, false, socket)
           if result['ok'] == 1 && ((is_master = result['ismaster'] == 1) || @slave_ok)
             @host, @port = host, port
+            apply_saved_authentication
           end
 
           # Note: slave_ok can be true only when connecting to a single node.
@@ -470,13 +526,13 @@ module Mongo
           raise MongoArgumentError, "MongoDB URI must include all three of username, password, " +
             "and db if any one of these is specified."
         else
-          auths << [uname, pwd, db]
+          add_auth(db, uname, pwd)
         end
 
         nodes << [host, port]
       end
 
-      [nodes, auths]
+      nodes
     end
 
     private
@@ -654,19 +710,5 @@ module Mongo
       end
       message
     end
-
-    # Authenticate for any auth info provided on instantiating the connection.
-    # Only called when a MongoDB URI has been used to instantiate the connection, and
-    # when that connection specifies databases and authentication credentials.
-    #
-    # @raise [MongoDBError]
-    def authenticate_databases
-      @auths.each do |auth|
-        user    = auth[0]
-        pwd     = auth[1]
-        db_name = auth[2]
-        self.db(db_name).authenticate(user, pwd)
-      end
-    end
   end
 end
diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index 99fd198..6846baa 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -82,7 +82,7 @@ module Mongo
     # @raise [AuthenticationError]
     #
     # @core authenticate authenticate-instance_method
-    def authenticate(username, password)
+    def authenticate(username, password, save_authorization=true)
       doc = command(:getnonce => 1)
       raise "error retrieving nonce: #{doc}" unless ok?(doc)
       nonce = doc['nonce']
@@ -92,8 +92,14 @@ module Mongo
       auth['user'] = username
       auth['nonce'] = nonce
       auth['key'] = Digest::MD5.hexdigest("#{nonce}#{username}#{hash_password(username, password)}")
-      ok?(command(auth)) ||
-        raise(MongoDBError::AuthenticationError, "Failed to authenticate user '#{username}' on db '#{self.name}'")
+      if ok?(command(auth))
+        if save_authorization
+          @connection.add_auth(@name, username, password)
+        end
+        true
+      else
+        raise(Mongo::AuthenticationError, "Failed to authenticate user '#{username}' on db '#{self.name}'")
+      end
     end
 
     # Adds a user to this database for use with authentication. If the user already
@@ -125,15 +131,21 @@ module Mongo
       end
     end
 
-    # Deauthorizes use for this database for this connection.
+    # Deauthorizes use for this database for this connection. Also removes
+    # any saved authorization in the connection class associated with this
+    # database.
     #
     # @raise [MongoDBError] if logging out fails.
     #
     # @return [Boolean]
     def logout
       doc = command(:logout => 1)
-      return true if ok?(doc)
-      raise MongoDBError, "error logging out: #{doc.inspect}"
+      if ok?(doc)
+        @connection.remove_auth(@name)
+        true
+      else
+        raise MongoDBError, "error logging out: #{doc.inspect}"
+      end
     end
 
     # Get an array of collection names in this database.
diff --git a/lib/mongo/gridfs/grid_io.rb b/lib/mongo/gridfs/grid_io.rb
index cc60e6c..e79886e 100644
--- a/lib/mongo/gridfs/grid_io.rb
+++ b/lib/mongo/gridfs/grid_io.rb
@@ -320,7 +320,7 @@ module Mongo
       if @safe
         @client_md5 = @local_md5.hexdigest
         if @local_md5 != @server_md5
-          raise @local_md5 != @server_md5, "File on server failed MD5 check"
+          raise GridError, "File on server failed MD5 check"
         end
       else
         @server_md5
diff --git a/test/auxillary/authentication_test.rb b/test/auxillary/authentication_test.rb
new file mode 100644
index 0000000..b78a9be
--- /dev/null
+++ b/test/auxillary/authentication_test.rb
@@ -0,0 +1,68 @@
+$:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
+require 'mongo'
+require 'test/unit'
+require 'test/test_helper'
+
+# NOTE: This test requires bouncing the server.
+# It also requires that a user exists on the admin database.
+class AuthenticationTest < Test::Unit::TestCase
+  include Mongo
+
+  def setup
+    @conn = Mongo::Connection.new
+    @db1 = @conn.db('mongo-ruby-test-auth1')
+    @db2 = @conn.db('mongo-ruby-test-auth2')
+    @admin = @conn.db('admin')
+  end
+
+  def teardown
+    @db1.authenticate('user1', 'secret')
+    @db2.authenticate('user2', 'secret')
+    @conn.drop_database('mongo-ruby-test-auth1')
+    @conn.drop_database('mongo-ruby-test-auth2')
+  end
+
+  def test_authenticate
+    @admin.authenticate('bob', 'secret')
+    @db1.add_user('user1', 'secret')
+    @db2.add_user('user2', 'secret')
+    @admin.logout
+
+    assert_raise Mongo::OperationFailure do
+      @db1['stuff'].insert({:a => 2}, :safe => true)
+    end
+
+    assert_raise Mongo::OperationFailure do
+      @db2['stuff'].insert({:a => 2}, :safe => true)
+    end
+
+    @db1.authenticate('user1', 'secret')
+    @db2.authenticate('user2', 'secret')
+
+    assert @db1['stuff'].insert({:a => 2}, :safe => true)
+    assert @db2['stuff'].insert({:a => 2}, :safe => true)
+
+    puts "Please bounce the server."
+    gets
+
+    # Here we reconnect.
+    begin
+      @db1['stuff'].find.to_a
+      rescue Mongo::ConnectionFailure
+    end
+
+    assert @db1['stuff'].insert({:a => 2}, :safe => true)
+    assert @db2['stuff'].insert({:a => 2}, :safe => true)
+
+    @db1.logout
+    assert_raise Mongo::OperationFailure do
+      @db1['stuff'].insert({:a => 2}, :safe => true)
+    end
+
+    @db2.logout
+    assert_raise Mongo::OperationFailure do
+      assert @db2['stuff'].insert({:a => 2}, :safe => true)
+    end
+  end
+
+end
diff --git a/test/auxillary/autoreconnect_test.rb b/test/auxillary/autoreconnect_test.rb
index 729c883..6b721ad 100644
--- a/test/auxillary/autoreconnect_test.rb
+++ b/test/auxillary/autoreconnect_test.rb
@@ -29,7 +29,7 @@ class AutoreconnectTest < Test::Unit::TestCase
 
     begin
       @coll.find.to_a
-      rescue Mongo::ConnectionFailure 
+      rescue Mongo::ConnectionFailure
     end
 
     results = []
@@ -38,5 +38,4 @@ class AutoreconnectTest < Test::Unit::TestCase
         assert results.any? {|r| r['a'] == a}, "Could not find record for a => #{a}"
       end
   end
-
 end
diff --git a/test/connection_test.rb b/test/connection_test.rb
index fab70f8..59e2b76 100644
--- a/test/connection_test.rb
+++ b/test/connection_test.rb
@@ -69,8 +69,6 @@ class TestConnection < Test::Unit::TestCase
     assert_kind_of Array, names
     assert names.length >= 1
     assert names.include?('ruby-mongo-info-test')
-
-    @mongo.drop_database('ruby-mongo-info-test')
   end
 
   def test_logging
@@ -80,7 +78,7 @@ class TestConnection < Test::Unit::TestCase
     db = Connection.new(@host, @port, :logger => logger).db('ruby-mongo-test')
     assert output.string.include?("admin.$cmd.find")
   end
-  
+
   def test_connection_logger
     output = StringIO.new
     logger = Logger.new(output)
@@ -124,6 +122,38 @@ class TestConnection < Test::Unit::TestCase
     assert_equal ['foo', 123], nodes[1]
   end
 
+  context "Saved authentications" do
+    setup do
+      @conn = Mongo::Connection.new
+      @auth = {'db_name' => 'test', 'username' => 'bob', 'password' => 'secret'}
+      @conn.add_auth(@auth['db_name'], @auth['username'], @auth['password'])
+    end
+
+    should "save the authentication" do
+      assert_equal @auth, @conn.auths[0]
+    end
+
+    should "replace the auth if given a new auth for the same db" do
+      auth = {'db_name' => 'test', 'username' => 'mickey', 'password' => 'm0u53'}
+      @conn.add_auth(auth['db_name'], auth['username'], auth['password'])
+      assert_equal 1, @conn.auths.length
+      assert_equal auth, @conn.auths[0]
+    end
+
+    should "remove auths by database" do
+      @conn.remove_auth('non-existent database')
+      assert_equal 1, @conn.auths.length
+
+      @conn.remove_auth('test')
+      assert_equal 0, @conn.auths.length
+    end
+
+    should "remove all auths" do
+      @conn.clear_auths
+      assert_equal 0, @conn.auths.length
+    end
+  end
+
   context "Connection exceptions" do
     setup do
       @conn = Mongo::Connection.new('localhost', 27017, :pool_size => 10, :timeout => 10)
diff --git a/test/unit/connection_test.rb b/test/unit/connection_test.rb
index ec0a5f7..870cfec 100644
--- a/test/unit/connection_test.rb
+++ b/test/unit/connection_test.rb
@@ -75,8 +75,10 @@ class ConnectionTest < Test::Unit::TestCase
         @conn = Connection.from_uri("mongodb://kyle:s3cr3t@localhost:27017/app,mickey:m0u5e@mydb.com:27018/dsny", :connect => false)
         assert_equal ['localhost', 27017], @conn.nodes[0]
         assert_equal ['mydb.com', 27018], @conn.nodes[1]
-        assert_equal ['kyle', 's3cr3t', 'app'], @conn.auths[0]
-        assert_equal ['mickey', 'm0u5e', 'dsny'], @conn.auths[1]
+        auth_hash = {'username' => 'kyle', 'password' => 's3cr3t', 'db_name' => 'app'}
+        assert_equal auth_hash, @conn.auths[0]
+        auth_hash = {'username' => 'mickey', 'password' => 'm0u5e', 'db_name' => 'dsny'}
+        assert_equal auth_hash, @conn.auths[1]
       end
 
       should "attempt to connect" do
@@ -86,6 +88,7 @@ class ConnectionTest < Test::Unit::TestCase
         admin_db = new_mock_db
         admin_db.expects(:command).returns({'ok' => 1, 'ismaster' => 1})
         @conn.expects(:[]).with('admin').returns(admin_db)
+        @conn.expects(:apply_saved_authentication)
         @conn.connect_to_master
       end