From 3b3bddc6267036ee018fce5c2a902d1292c34873 Mon Sep 17 00:00:00 2001 From: Cody Cutrer Date: Wed, 17 Dec 2014 14:49:21 -0700 Subject: [PATCH] support a connection callback for proxies also refactors connection establishment to only go through one method --- Contributors.rdoc | 1 + lib/net/ldap.rb | 139 +++++++++++++++------------------------------- test/test_ldap.rb | 18 ++++++ 3 files changed, 64 insertions(+), 94 deletions(-) diff --git a/Contributors.rdoc b/Contributors.rdoc index bef012a9..b3b25ff6 100644 --- a/Contributors.rdoc +++ b/Contributors.rdoc @@ -20,3 +20,4 @@ Contributions since: * Erik Hetzner (egh) * nowhereman * David J. Lee (DavidJLee) +* Cody Cutrer (ccutrer) diff --git a/lib/net/ldap.rb b/lib/net/ldap.rb index b181d83d..ee758b96 100644 --- a/lib/net/ldap.rb +++ b/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| + use_connection(args[:auth]) do |conn| + @result = 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 + use_connection(args[:auth]) do |conn| + @result = 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 + use_connection(args[:auth]) do |conn| + @result = 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 + use_connection(args[:auth]) do |conn| + @result = 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 + use_connection(args[:auth]) do |conn| + @result = conn.delete(args) end @result.success? end @@ -1277,4 +1200,32 @@ def paged_searches_supported? @server_caps ||= search_root_dse @server_caps[:supportedcontrol].include?(Net::LDAP::LDAPControls::PAGED_RESULTS) end + + private + + def use_connection(auth) + if @open_connection + yield @open_connection + else + @result = 0 + begin + conn = new_connection + if (@result = conn.bind(auth || @auth)).result_code == Net::LDAP::ResultCodeSuccess + yield conn + end + ensure + conn.close if conn + end + end + end + + 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 diff --git a/test/test_ldap.rb b/test/test_ldap.rb index 9704b346..b58cbe2f 100644 --- a/test/test_ldap.rb +++ b/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