From 4237472cb11fa548c203e6fa21ff372728f60f3b Mon Sep 17 00:00:00 2001
From: ned <ned@appliedtrust.com>
Date: Tue, 2 Jun 2015 12:56:59 -0600
Subject: Added search SizeLimit support

---
 server_search.go      | 10 +++++-----
 server_search_test.go | 36 +++++++++++++++++++++++++++++++++---
 server_test.go        |  2 +-
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/server_search.go b/server_search.go
index 4dda9a3..74d06b9 100644
--- a/server_search.go
+++ b/server_search.go
@@ -46,11 +46,6 @@ func HandleSearchRequest(req *ber.Packet, controls *[]Control, messageID uint64,
 
 	for i, entry := range searchResp.Entries {
 		if server.EnforceLDAP {
-			// size limit
-			if searchReq.SizeLimit > 0 && i >= searchReq.SizeLimit {
-				break
-			}
-
 			// filter
 			keep, resultCode := ServerApplyFilter(filterPacket, entry)
 			if resultCode != LDAPResultSuccess {
@@ -77,6 +72,11 @@ func HandleSearchRequest(req *ber.Packet, controls *[]Control, messageID uint64,
 				}
 			}
 
+			// size limit
+			if searchReq.SizeLimit > 0 && i >= searchReq.SizeLimit {
+				break
+			}
+
 			// attributes
 			if len(searchReq.Attributes) > 1 || (len(searchReq.Attributes) == 1 && len(searchReq.Attributes[0]) > 0) {
 				entry, err = filterAttributes(entry, searchReq.Attributes)
diff --git a/server_search_test.go b/server_search_test.go
index c3f42b0..ed05f68 100644
--- a/server_search_test.go
+++ b/server_search_test.go
@@ -66,13 +66,23 @@ func TestSearchSizelimit(t *testing.T) {
 
 	go func() {
 		cmd := exec.Command("ldapsearch", "-H", ldapURL, "-x",
-			"-b", serverBaseDN, "-D", "cn=testy,"+serverBaseDN, "-w", "iLike2test", "-z", "9") // effectively no limit for this test
+			"-b", serverBaseDN, "-D", "cn=testy,"+serverBaseDN, "-w", "iLike2test") // no limit for this test
 		out, _ := cmd.CombinedOutput()
 		if !strings.Contains(string(out), "result: 0 Success") {
 			t.Errorf("ldapsearch failed: %v", string(out))
 		}
 		if !strings.Contains(string(out), "numEntries: 3") {
-			t.Errorf("ldapsearch sizelimit failed - not enough entries: %v", string(out))
+			t.Errorf("ldapsearch sizelimit unlimited failed - not enough entries: %v", string(out))
+		}
+
+		cmd = exec.Command("ldapsearch", "-H", ldapURL, "-x",
+			"-b", serverBaseDN, "-D", "cn=testy,"+serverBaseDN, "-w", "iLike2test", "-z", "9") // effectively no limit for this test
+		out, _ = cmd.CombinedOutput()
+		if !strings.Contains(string(out), "result: 0 Success") {
+			t.Errorf("ldapsearch failed: %v", string(out))
+		}
+		if !strings.Contains(string(out), "numEntries: 3") {
+			t.Errorf("ldapsearch sizelimit 9 failed - not enough entries: %v", string(out))
 		}
 
 		cmd = exec.Command("ldapsearch", "-H", ldapURL, "-x",
@@ -82,7 +92,27 @@ func TestSearchSizelimit(t *testing.T) {
 			t.Errorf("ldapsearch failed: %v", string(out))
 		}
 		if !strings.Contains(string(out), "numEntries: 2") {
-			t.Errorf("ldapsearch sizelimit failed - too many entries: %v", string(out))
+			t.Errorf("ldapsearch sizelimit 2 failed - too many entries: %v", string(out))
+		}
+
+		cmd = exec.Command("ldapsearch", "-H", ldapURL, "-x",
+			"-b", serverBaseDN, "-D", "cn=testy,"+serverBaseDN, "-w", "iLike2test", "-z", "1")
+		out, _ = cmd.CombinedOutput()
+		if !strings.Contains(string(out), "result: 0 Success") {
+			t.Errorf("ldapsearch failed: %v", string(out))
+		}
+		if !strings.Contains(string(out), "numEntries: 1") {
+			t.Errorf("ldapsearch sizelimit 1 failed - too many entries: %v", string(out))
+		}
+
+		cmd = exec.Command("ldapsearch", "-H", ldapURL, "-x",
+			"-b", serverBaseDN, "-D", "cn=testy,"+serverBaseDN, "-w", "iLike2test", "-z", "0")
+		out, _ = cmd.CombinedOutput()
+		if !strings.Contains(string(out), "result: 0 Success") {
+			t.Errorf("ldapsearch failed: %v", string(out))
+		}
+		if !strings.Contains(string(out), "numEntries: 3") {
+			t.Errorf("ldapsearch sizelimit 0 failed - wrong number of entries: %v", string(out))
 		}
 		done <- true
 	}()
diff --git a/server_test.go b/server_test.go
index 7e813ec..88c47bf 100644
--- a/server_test.go
+++ b/server_test.go
@@ -284,7 +284,7 @@ func TestSearchStats(t *testing.T) {
 
 	stats := s.GetStats()
 	log.Println(stats)
-	if stats.Conns != 2 || stats.Binds != 1 {
+	if stats.Conns != 1 || stats.Binds != 1 {
 		t.Errorf("Stats data missing or incorrect: %v", w.buffer.String())
 	}
 	quit <- true
-- 
cgit v1.2.3