Merge pull request #116 from jch/connection-write-request

Add private Net::LDAP::Connection#write_request

Author: jch

lib/net/ldap/connection.rb

@@ -10,7 +10,7 @@ def initialize(server)
     @instrumentation_service = server[:instrumentation_service]
 
     begin
-      @conn = TCPSocket.new(server[:host], server[:port])
+      @conn = server[:socket] || TCPSocket.new(server[:host], server[:port])
     rescue SocketError
       raise Net::LDAP::LdapError, "No such address or other socket error."
     rescue Errno::ECONNREFUSED
@@ -87,10 +87,8 @@ def setup_encryption(args)
       # additional branches requiring server validation and peer certs, etc.
       # go here.
     when :start_tls
-      msgid = next_msgid.to_ber
       request = [Net::LDAP::StartTlsOid.to_ber_contextspecific(0)].to_ber_appsequence(Net::LDAP::PDU::ExtendedRequest)
-      request_pkt = [msgid, request].to_ber_sequence
-      write request_pkt
+      write(request)
       be = read
       raise Net::LDAP::LdapError, "no start_tls result" if be.nil?
       pdu = Net::LDAP::PDU.new(be)
@@ -130,14 +128,17 @@ def read(syntax = Net::LDAP::AsnSyntax)
   end
   private :read
 
-  # Internal: Writes the given packet to the configured connection.
+  # Internal: Write a BER formatted packet with the next message id to the
+  # configured connection.
   #
-  # - packet: the BER data packet to write on the socket.
+  # - request: required BER formatted request
+  # - controls: optional BER formatted controls
   #
   # Returns the return value from writing to the connection, which in some
   # cases is the Integer number of bytes written to the socket.
-  def write(packet)
+  def write(request, controls = nil)
     instrument "write.net_ldap_connection" do |payload|
+      packet = [next_msgid.to_ber, request, controls].compact.to_ber_sequence
       payload[:content_length] = @conn.write(packet)
     end
   end
@@ -176,11 +177,9 @@ def bind_simple(auth)
 
     raise Net::LDAP::LdapError, "Invalid binding information" unless (user && psw)
 
-    msgid = next_msgid.to_ber
     request = [LdapVersion.to_ber, user.to_ber,
       psw.to_ber_contextspecific(0)].to_ber_appsequence(0)
-    request_pkt = [msgid, request].to_ber_sequence
-    write request_pkt
+    write(request)
 
     (be = read and pdu = Net::LDAP::PDU.new(be)) or raise Net::LDAP::LdapError, "no bind result"
 
@@ -215,11 +214,9 @@ def bind_sasl(auth)
 
     n = 0
     loop {
-      msgid = next_msgid.to_ber
       sasl = [mech.to_ber, cred.to_ber].to_ber_contextspecific(3)
       request = [LdapVersion.to_ber, "".to_ber, sasl].to_ber_appsequence(0)
-      request_pkt = [msgid, request].to_ber_sequence
-      write request_pkt
+      write(request)
 
       (be = read and pdu = Net::LDAP::PDU.new(be)) or raise Net::LDAP::LdapError, "no bind result"
       return pdu unless pdu.result_code == 14 # saslBindInProgress
@@ -393,8 +390,7 @@ def search(args = {})
         controls << sort_control if sort_control
         controls = controls.empty? ? nil : controls.to_ber_contextspecific(0)
 
-        pkt = [next_msgid.to_ber, request, controls].compact.to_ber_sequence
-        write pkt
+        write(request, controls)
 
         result_pdu = nil
         controls = []
@@ -502,8 +498,7 @@ def modify(args)
     ops = self.class.modify_ops args[:operations]
     request = [ modify_dn.to_ber,
       ops.to_ber_sequence ].to_ber_appsequence(6)
-    pkt = [ next_msgid.to_ber, request ].to_ber_sequence
-    write pkt
+    write(request)
 
     (be = read) && (pdu = Net::LDAP::PDU.new(be)) && (pdu.app_tag == Net::LDAP::PDU::ModifyResponse) or raise Net::LDAP::LdapError, "response missing or invalid"
 
@@ -525,8 +520,7 @@ def add(args)
     }
 
     request = [add_dn.to_ber, add_attrs.to_ber_sequence].to_ber_appsequence(8)
-    pkt = [next_msgid.to_ber, request].to_ber_sequence
-    write pkt
+    write(request)
 
     (be = read) &&
       (pdu = Net::LDAP::PDU.new(be)) &&
@@ -548,8 +542,7 @@ def rename(args)
     request = [old_dn.to_ber, new_rdn.to_ber, delete_attrs.to_ber]
     request << new_superior.to_ber_contextspecific(0) unless new_superior == nil
 
-    pkt = [next_msgid.to_ber, request.to_ber_appsequence(12)].to_ber_sequence
-    write pkt
+    write(request.to_ber_appsequence(12))
 
     (be = read) &&
     (pdu = Net::LDAP::PDU.new( be )) && (pdu.app_tag == Net::LDAP::PDU::ModifyRDNResponse) or
@@ -565,8 +558,7 @@ def delete(args)
     dn = args[:dn] or raise "Unable to delete empty DN"
     controls = args.include?(:control_codes) ? args[:control_codes].to_ber_control : nil #use nil so we can compact later
     request = dn.to_s.to_ber_application_string(10)
-    pkt = [next_msgid.to_ber, request, controls].compact.to_ber_sequence
-    write pkt
+    write(request, controls)
 
     (be = read) && (pdu = Net::LDAP::PDU.new(be)) && (pdu.app_tag == Net::LDAP::PDU::DeleteResponse) or raise Net::LDAP::LdapError, "response missing or invalid"
 

test/test_ldap_connection.rb

@@ -42,6 +42,29 @@ def test_modify_ops_replace
     expected = [ "0#\n\x01\x020\x1E\x04\x04mail1\x16\x04\x14testuser@example.com" ]
     assert_equal(expected, result)
   end
+
+  def test_write
+    mock = flexmock("socket")
+    mock.should_receive(:write).with([1.to_ber, "request"].to_ber_sequence).and_return(true)
+    conn = Net::LDAP::Connection.new(:socket => mock)
+    conn.send(:write, "request")
+  end
+
+  def test_write_with_controls
+    mock = flexmock("socket")
+    mock.should_receive(:write).with([1.to_ber, "request", "controls"].to_ber_sequence).and_return(true)
+    conn = Net::LDAP::Connection.new(:socket => mock)
+    conn.send(:write, "request", "controls")
+  end
+
+  def test_write_increments_msgid
+    mock = flexmock("socket")
+    mock.should_receive(:write).with([1.to_ber, "request1"].to_ber_sequence).and_return(true)
+    mock.should_receive(:write).with([2.to_ber, "request2"].to_ber_sequence).and_return(true)
+    conn = Net::LDAP::Connection.new(:socket => mock)
+    conn.send(:write, "request1")
+    conn.send(:write, "request2")
+  end
 end