From 537dc5e1bb750da7e15a3e45e81124c3ef27b954 Mon Sep 17 00:00:00 2001
From: Dave Sims <davesims@github.com>
Date: Tue, 12 Jul 2016 16:00:57 -0500
Subject: [PATCH] Restarting referrals from
 https://github.com/github/github-ldap/pull/87

---
 lib/github/ldap.rb                            | 48 +++++++++++++++++++
 lib/github/ldap/member_search/recursive.rb    |  5 +-
 .../ldap/membership_validators/recursive.rb   |  5 +-
 3 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/lib/github/ldap.rb b/lib/github/ldap.rb
index 0545247..b37f3a4 100644
--- a/lib/github/ldap.rb
+++ b/lib/github/ldap.rb
@@ -1,5 +1,6 @@
 require 'net/ldap'
 require 'forwardable'
+require 'uri'
 
 require 'github/ldap/filter'
 require 'github/ldap/domain'
@@ -100,6 +101,9 @@ def initialize(options = {})
 
       # enables instrumenting queries
       @instrumentation_service = options[:instrumentation_service]
+
+      # referral connection handle
+      @referrals = {}
     end
 
     # Public - Whether membership checks should recurse into nested groups when
@@ -187,7 +191,51 @@ def search(options, &block)
               result.concat Array(rs) unless rs == false
             end
           end
+        if options[:return_referrals]
+          ref_result = []
+          result.delete_if do |entry|
+            if entry.respond_to?('search_referrals')
+              rs = search_referrals(entry[:search_referrals], options, &block)
+              ref_result.concat Array(rs) unless rs == false
+              true
+            else
+              false
+            end
+          end
+          result.concat ref_result
+        end
+        return [] if result == false
+        Array(result)
+      end
+    end
 
+    # Internal: Searches the referral LDAP servers
+    #
+    # referal_hosts: list of referral hosts uris
+    # options: is a hash with the same options that Net::LDAP::Connection#search supports.
+    # block: is an optional block to pass to the search.
+    #
+    # Returns an Array of Net::LDAP::Entry.
+    def search_referrals(referral_uris, options, &block)
+      instrument "search_referrals.github_ldap", options.dup do |payload|
+        options.delete(:return_referrals)
+        result = []
+        referral_uris.each do |referral_url|
+          unless @referrals.has_key? referral_url
+            uri = URI(referral_url)
+            @referrals[referral_url] = Net::LDAP.new({
+              host: uri.host,
+              port: uri.port,
+              base: uri.path.sub(/^\//, ''),
+              auth: @connection.instance_variable_get(:@auth),
+              encryption: @connection.instance_variable_get(:@encryption),
+              instrumentation_service: @connection.instance_variable_get(:@instrumentation_service)
+            })
+          end
+            options.delete(:base)
+            rs = @referrals[referral_url].search(options, &block)
+            result.concat Array(rs) unless rs == false
+        end
         return [] if result == false
         Array(result)
       end
diff --git a/lib/github/ldap/member_search/recursive.rb b/lib/github/ldap/member_search/recursive.rb
index a36aa4d..7c59b8f 100644
--- a/lib/github/ldap/member_search/recursive.rb
+++ b/lib/github/ldap/member_search/recursive.rb
@@ -123,7 +123,8 @@ def find_groups_by_dn(dn)
             base: dn,
             scope: Net::LDAP::SearchScope_BaseObject,
             attributes: attrs,
-            filter: ALL_GROUPS_FILTER
+            filter: ALL_GROUPS_FILTER,
+            return_referrals: true
         end
         private :find_groups_by_dn
 
@@ -133,7 +134,7 @@ def find_groups_by_dn(dn)
         def entries_by_uid(members)
           filter = members.map { |uid| Net::LDAP::Filter.eq(ldap.uid, uid) }.reduce(:|)
           domains.each_with_object([]) do |domain, entries|
-            entries.concat domain.search(filter: filter, attributes: attrs)
+            entries.concat domain.search(filter: filter, attributes: attrs, return_referrals: true)
           end.compact
         end
         private :entries_by_uid
diff --git a/lib/github/ldap/membership_validators/recursive.rb b/lib/github/ldap/membership_validators/recursive.rb
index 3b78545..d6398bc 100644
--- a/lib/github/ldap/membership_validators/recursive.rb
+++ b/lib/github/ldap/membership_validators/recursive.rb
@@ -51,7 +51,7 @@ def perform(entry, depth_override = nil)
 
           domains.each do |domain|
             # find groups entry is an immediate member of
-            membership = domain.search(filter: member_filter(entry), attributes: ATTRS)
+            membership = domain.search(filter: member_filter(entry), attributes: ATTRS, return_referrals: true)
 
             # success if any of these groups match the restricted auth groups
             return true if membership.any? { |entry| group_dns.include?(entry.dn) }
@@ -62,7 +62,7 @@ def perform(entry, depth_override = nil)
             # recurse to at most `depth`
             (depth_override || depth).times do |n|
               # find groups whose members include membership groups
-              membership = domain.search(filter: membership_filter(membership), attributes: ATTRS)
+              membership = domain.search(filter: membership_filter(membership), attributes: ATTRS, return_referrals: true)
 
               # success if any of these groups match the restricted auth groups
               return true if membership.any? { |entry| group_dns.include?(entry.dn) }
@@ -115,3 +115,4 @@ def group_dns
     end
   end
 end
+