From ec86275b6060db33a03b57d3b0926c9ffff62c0a Mon Sep 17 00:00:00 2001
From: Tyler Brock <tyler.brock@gmail.com>
Date: Thu, 15 Mar 2012 13:50:02 -0400
Subject: [PATCH] RUBY-424 Authenticating with only secondary fails

Authentication command now prefers to use primary node but will fall
back on secondary if no primary is available
---
 lib/mongo/connection.rb              | 19 +++++++++++++++++++
 lib/mongo/db.rb                      |  7 +++++--
 lib/mongo/repl_set_connection.rb     | 21 +++++++++++++++++++--
 test/auxillary/repl_set_auth_test.rb | 24 +++++++++++++++++++-----
 4 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/lib/mongo/connection.rb b/lib/mongo/connection.rb
index 7081ac9..3fe8089 100644
--- a/lib/mongo/connection.rb
+++ b/lib/mongo/connection.rb
@@ -480,6 +480,12 @@ module Mongo
       @max_bson_size
     end
 
+    # Prefer primary pool but fall back to secondary
+    def checkout_best
+      connect unless connected?
+      @primary_pool.checkout
+    end
+
     # Checkout a socket for reading (i.e., a secondary node).
     # Note: this is overridden in ReplSetConnection.
     def checkout_reader
@@ -513,6 +519,19 @@ module Mongo
       end
     end
 
+    # Excecutes block with the best available socket
+    def best_available_socket
+      socket = nil
+      begin
+        socket = checkout_best
+        yield socket
+      ensure
+        if socket
+          socket.pool.checkin(socket)
+        end
+      end
+    end
+
     protected
 
     def valid_opts
diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index 72f0e8d..1711879 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -111,10 +111,13 @@ module Mongo
         if !save_auth
           raise MongoArgumentError, "If using connection pooling, :save_auth must be set to true."
         end
-        @connection.authenticate_pools
       end
 
-      issue_authentication(username, password, save_auth)
+      @connection.best_available_socket do |socket|
+        issue_authentication(username, password, save_auth, :socket => socket)
+      end
+
+      @connection.authenticate_pools
     end
 
     def issue_authentication(username, password, save_auth=true, opts={})
diff --git a/lib/mongo/repl_set_connection.rb b/lib/mongo/repl_set_connection.rb
index 406df36..4aa10b8 100644
--- a/lib/mongo/repl_set_connection.rb
+++ b/lib/mongo/repl_set_connection.rb
@@ -287,14 +287,18 @@ module Mongo
     end
 
     def authenticate_pools
-      primary_pool.authenticate_existing
+      if primary_pool
+        primary_pool.authenticate_existing
+      end
       secondary_pools.each do |pool|
         pool.authenticate_existing
       end
     end
 
     def logout_pools(db)
-      primary_pool.logout_existing(db)
+      if primary_pool
+        primary_pool.logout_existing(db)
+      end
       secondary_pools.each do |pool|
         pool.logout_existing(db)
       end
@@ -323,6 +327,19 @@ module Mongo
         raise ConnectionFailure.new("Could not checkout a socket.")
       end
     end
+
+    # Checkout best available socket by trying primary
+    # pool first and then falling back to secondary.
+    def checkout_best
+      checkout do
+        socket = get_socket_from_pool(:primary)
+        if !socket
+          connect
+          socket = get_socket_from_pool(:secondary)
+        end
+        socket
+      end
+    end
     
     # Checkout a socket for reading (i.e., a secondary node).
     # Note that @read_pool might point to the primary pool
diff --git a/test/auxillary/repl_set_auth_test.rb b/test/auxillary/repl_set_auth_test.rb
index afa3984..8f6e630 100644
--- a/test/auxillary/repl_set_auth_test.rb
+++ b/test/auxillary/repl_set_auth_test.rb
@@ -1,21 +1,22 @@
 $:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
 require './test/test_helper'
 require './test/tools/auth_repl_set_manager'
+require './test/replica_sets/rs_test_helper'
 
 class AuthTest < Test::Unit::TestCase
   include Mongo
 
   def setup
-    @manager = AuthReplSetManager.new(:start_port => 40000)
-    @manager.start_set
+    @rs = AuthReplSetManager.new(:start_port => 40000)
+    @rs.start_set
   end
 
   def teardown
-    @manager.cleanup_set
+    #@rs.cleanup_set
   end
 
   def test_repl_set_auth
-    @conn = ReplSetConnection.new(build_seeds(3), :name => @manager.name)
+    @conn = ReplSetConnection.new(build_seeds(3), :name => @rs.name)
 
     # Add an admin user
     @conn['admin'].add_user("me", "secret")
@@ -51,7 +52,20 @@ class AuthTest < Test::Unit::TestCase
     end
 
     # But not when authenticated
-    @slave1['admin'].authenticate("me", "secret")
+    assert @slave1['admin'].authenticate("me", "secret")
     assert @slave1['foo']['stuff'].find_one
+
+    # Same should apply when using :secondary_only
+    @second_only = ReplSetConnection.new(build_seeds(3), 
+      :require_primary => false, :read => :secondary_only)
+
+    # Find should fail
+    assert_raise_error Mongo::OperationFailure, "unauthorized" do
+      @second_only['foo']['stuff'].find_one
+    end
+
+    # But not when authenticated
+    assert @second_only['admin'].authenticate("me", "secret")
+    assert @second_only['foo']['stuff'].find_one
   end
 end