support a connection callback for proxies

Contributors.rdoc

@@ -20,3 +20,4 @@ Contributions since:
 * Erik Hetzner (egh)
 * nowhereman
 * David J. Lee (DavidJLee)
+* Cody Cutrer (ccutrer)

lib/net/ldap.rb

@@ -441,6 +441,10 @@ def self.result2string(code) #:nodoc:
   # described below. The following arguments are supported:
   # * :host => the LDAP server's IP-address (default 127.0.0.1)
   # * :port => the LDAP server's TCP port (default 389)
+  # * :connect_cb => a Proc that will be called when a new connection is
+  #   needed. This should return an actual Ruby IO object. Useful for
+  #   manually handling connecting, like if you want to go through a proxy
+  #   server. It will receive :host: and :port: as arguments.
   # * :auth => a Hash containing authorization parameters. Currently
   #   supported values include: {:method => :anonymous} and {:method =>
   #   :simple, :username => your_user_name, :password => your_password }
@@ -469,6 +473,7 @@ def self.result2string(code) #:nodoc:
   def initialize(args = {})
     @host = args[:host] || DefaultHost
     @port = args[:port] || DefaultPort
+    @connect_cb = args[:connect_cb]
     @verbose = false # Make this configurable with a switch on the class.
     @auth = args[:auth] || DefaultAuth
     @base = args[:base] || DefaultTreebase
@@ -670,12 +675,7 @@ def open
 
     instrument "open.net_ldap" do |payload|
       begin
-        @open_connection =
-          Net::LDAP::Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
+        @open_connection = new_connection
         payload[:connection] = @open_connection
         payload[:bind]       = @open_connection.bind(@auth)
         yield self
@@ -745,27 +745,11 @@ def search(args = {})
     result_set = return_result_set ? [] : nil
 
     instrument "search.net_ldap", args do |payload|
-      if @open_connection
-        @result = @open_connection.search(args) { |entry|
+      @result = use_connection(args) do |conn|
+        conn.search(args) { |entry|
           result_set << entry if result_set
           yield entry if block_given?
         }
-      else
-        begin
-          conn = Net::LDAP::Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
-          if (@result = conn.bind(args[:auth] || @auth)).result_code == Net::LDAP::ResultCodeSuccess
-            @result = conn.search(args) { |entry|
-              result_set << entry if result_set
-              yield entry if block_given?
-            }
-          end
-        ensure
-          conn.close if conn
-        end
       end
 
       if return_result_set
@@ -844,11 +828,7 @@ def bind(auth = @auth)
         payload[:bind]       = @result = @open_connection.bind(auth)
       else
         begin
-          conn = Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
+          conn = new_connection
           payload[:connection] = conn
           payload[:bind]       = @result = conn.bind(auth)
         ensure
@@ -946,22 +926,8 @@ def bind_as(args = {})
   #  end
   def add(args)
     instrument "add.net_ldap", args do |payload|
-      if @open_connection
-        @result = @open_connection.add(args)
-      else
-        @result = 0
-        begin
-          conn = Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
-          if (@result = conn.bind(args[:auth] || @auth)).result_code == Net::LDAP::ResultCodeSuccess
-            @result = conn.add(args)
-          end
-        ensure
-          conn.close if conn
-        end
+      @result = use_connection(args) do |conn|
+        conn.add(args)
       end
       @result.success?
     end
@@ -1050,24 +1016,9 @@ def add(args)
   # does _not_ imply transactional atomicity, which LDAP does not provide.
   def modify(args)
     instrument "modify.net_ldap", args do |payload|
-      if @open_connection
-        @result = @open_connection.modify(args)
-      else
-        @result = 0
-        begin
-          conn = Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
-          if (@result = conn.bind(args[:auth] || @auth)).result_code == Net::LDAP::ResultCodeSuccess
-            @result = conn.modify(args)
-          end
-        ensure
-          conn.close if conn
-        end
+      @result = use_connection(args) do |conn|
+        conn.modify(args)
       end
-
       @result.success?
     end
   end
@@ -1127,22 +1078,8 @@ def delete_attribute(dn, attribute)
   # _Documentation_ _stub_
   def rename(args)
     instrument "rename.net_ldap", args do |payload|
-      if @open_connection
-        @result = @open_connection.rename(args)
-      else
-        @result = 0
-        begin
-          conn = Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
-          if (@result = conn.bind(args[:auth] || @auth)).result_code == Net::LDAP::ResultCodeSuccess
-            @result = conn.rename(args)
-          end
-        ensure
-          conn.close if conn
-        end
+      @result = use_connection(args) do |conn|
+        conn.rename(args)
       end
       @result.success?
     end
@@ -1160,22 +1097,8 @@ def rename(args)
   #  ldap.delete :dn => dn
   def delete(args)
     instrument "delete.net_ldap", args do |payload|
-      if @open_connection
-        @result = @open_connection.delete(args)
-      else
-        @result = 0
-        begin
-          conn = Connection.new \
-            :host                    => @host,
-            :port                    => @port,
-            :encryption              => @encryption,
-            :instrumentation_service => @instrumentation_service
-          if (@result = conn.bind(args[:auth] || @auth)).result_code == Net::LDAP::ResultCodeSuccess
-            @result = conn.delete(args)
-          end
-        ensure
-          conn.close
-        end
+      @result = use_connection(args) do |conn|
+        conn.delete(args)
       end
       @result.success?
     end
@@ -1277,4 +1200,38 @@ def paged_searches_supported?
     @server_caps ||= search_root_dse
     @server_caps[:supportedcontrol].include?(Net::LDAP::LDAPControls::PAGED_RESULTS)
   end
+
+  private
+
+  # Yields an open connection if there is one, otherwise establishes a new
+  # connection, binds, and yields it. If binding fails, it will return the
+  # result from that, and :use_connection: will not yield at all. If not
+  # the return value is whatever is returned from the block.
+  def use_connection(args)
+    if @open_connection
+      yield @open_connection
+    else
+      begin
+        conn = new_connection
+        if (result = conn.bind(args[:auth] || @auth)).result_code == Net::LDAP::ResultCodeSuccess
+          yield conn
+        else
+          return result
+        end
+      ensure
+        conn.close if conn
+      end
+    end
+  end
+
+  # Establish a new connection to the LDAP server
+  def new_connection
+    socket = @connect_cb.call(@host, @port) if @connect_cb
+    Net::LDAP::Connection.new \
+      :socket                  => socket,
+      :host                    => @host,
+      :port                    => @port,
+      :encryption              => @encryption,
+      :instrumentation_service => @instrumentation_service
+  end
 end # class LDAP

test/test_ldap.rb

@@ -57,4 +57,22 @@ def test_instrument_search_with_size
     assert_equal "(uid=user1)", payload[:filter]
     assert_equal result.size, payload[:size]
   end
+
+  def test_connect_cb
+    flexmock(Net::LDAP::Connection).should_receive(:new).with(
+        :socket => 42,
+        :host => "test.mocked.com",
+        :port => 636,
+        :encryption => nil,
+        :instrumentation_service => @service).and_return(@connection)
+    flexmock(@connection).should_receive(:bind).and_return(flexmock(:bind_result, :result_code => Net::LDAP::ResultCodeSuccess))
+
+    @subject = Net::LDAP.new \
+      :connect_cb => lambda { |host, port| 42 },
+      :host => "test.mocked.com", :port => 636,
+      :force_no_page => true, # so server capabilities are not queried
+      :instrumentation_service => @service
+
+    @subject.open {}
+  end
 end