From 9f2f22bff5de4acc97f18326da924edb24cd4a7e Mon Sep 17 00:00:00 2001
From: Nicolas Capponi <nicolas.capponi@forgerock.com>
Date: Mon, 30 Sep 2013 09:42:41 +0000
Subject: [PATCH] Fix OPENDJ-916 - MemoryBackend in SDK should support size limit Review CR-2386

---
 opendj3/opendj-ldap-sdk/src/test/java/org/forgerock/opendj/ldap/MemoryBackendTestCase.java |  146 +++++++++++++++++++++++++-----------
 opendj3/opendj-ldap-sdk/src/main/java/org/forgerock/opendj/ldap/MemoryBackend.java         |   75 +++++++++++-------
 2 files changed, 146 insertions(+), 75 deletions(-)

diff --git a/opendj3/opendj-ldap-sdk/src/main/java/org/forgerock/opendj/ldap/MemoryBackend.java b/opendj3/opendj-ldap-sdk/src/main/java/org/forgerock/opendj/ldap/MemoryBackend.java
index f2a89ef..812fdfc 100644
--- a/opendj3/opendj-ldap-sdk/src/main/java/org/forgerock/opendj/ldap/MemoryBackend.java
+++ b/opendj3/opendj-ldap-sdk/src/main/java/org/forgerock/opendj/ldap/MemoryBackend.java
@@ -34,7 +34,7 @@
 
 import java.io.IOException;
 import java.util.Collection;
-import java.util.NavigableMap;
+import java.util.Map;
 import java.util.concurrent.ConcurrentSkipListMap;
 
 import org.forgerock.i18n.LocalizedIllegalArgumentException;
@@ -385,7 +385,6 @@
             final SearchResultHandler resultHandler) {
         try {
             final DN dn = request.getName();
-            final Entry baseEntry = getRequiredEntry(request, dn);
             final SearchScope scope = request.getScope();
             final Filter filter = request.getFilter();
             final Matcher matcher = filter.matcher(schema);
@@ -393,36 +392,14 @@
                     new AttributeFilter(request.getAttributes(), schema).typesOnly(request
                             .isTypesOnly());
             if (scope.equals(SearchScope.BASE_OBJECT)) {
+                final Entry baseEntry = getRequiredEntry(request, dn);
                 if (matcher.matches(baseEntry).toBoolean()) {
                     sendEntry(attributeFilter, resultHandler, baseEntry);
                 }
-            } else if (scope.equals(SearchScope.SINGLE_LEVEL)) {
-                final NavigableMap<DN, Entry> subtree =
-                        entries.subMap(dn, dn.child(RDN.maxValue()));
-                for (final Entry entry : subtree.values()) {
-                    // Check for cancellation.
-                    requestContext.checkIfCancelled(false);
-                    final DN childDN = entry.getName();
-                    if (childDN.isChildOf(dn)) {
-                        if (matcher.matches(entry).toBoolean()
-                                && !sendEntry(attributeFilter, resultHandler, entry)) {
-                            // Caller has asked to stop sending results.
-                            break;
-                        }
-                    }
-                }
-            } else if (scope.equals(SearchScope.WHOLE_SUBTREE)) {
-                final NavigableMap<DN, Entry> subtree =
-                        entries.subMap(dn, dn.child(RDN.maxValue()));
-                for (final Entry entry : subtree.values()) {
-                    // Check for cancellation.
-                    requestContext.checkIfCancelled(false);
-                    if (matcher.matches(entry).toBoolean()
-                            && !sendEntry(attributeFilter, resultHandler, entry)) {
-                        // Caller has asked to stop sending results.
-                        break;
-                    }
-                }
+            } else if (scope.equals(SearchScope.SINGLE_LEVEL) || scope.equals(SearchScope.WHOLE_SUBTREE)) {
+                searchWithSubordinates(
+                        requestContext, resultHandler, dn, matcher,
+                        attributeFilter, request.getSizeLimit(), scope);
             } else {
                 throw newErrorResult(ResultCode.PROTOCOL_ERROR,
                         "Search request contains an unsupported search scope");
@@ -490,6 +467,46 @@
         return entries.size();
     }
 
+    /**
+     * Perform a search for scope that includes subordinates, i.e., either
+     * <code>SearchScope.SINGLE_LEVEL</code> or <code>SearchScope.WHOLE_SUBTREE</code>.
+     *
+     * @param requestContext context of this request
+     * @param resultHandler handler which should be used to send back the search results to the client.
+     * @param dn distinguished name of the base entry used for this request
+     * @param matcher to filter entries that matches this request
+     * @param attributeFilter to select attributes to return in search results
+     * @param sizeLimit maximum number of entries to return. A value of zero indicates no restriction
+     *          on number of entries.
+     * @throws CancelledResultException
+     *           If a cancellation request has been received and processing of
+     *           the request should be aborted if possible.
+     * @throws ErrorResultException
+     *           If the request is unsuccessful.
+     */
+    private void searchWithSubordinates(final RequestContext requestContext, final SearchResultHandler resultHandler,
+            final DN dn, final Matcher matcher, final AttributeFilter attributeFilter, final int sizeLimit,
+            SearchScope scope) throws CancelledResultException, ErrorResultException {
+
+        final Map<DN, Entry> subtree = entries.subMap(dn, dn.child(RDN.maxValue()));
+        int numberOfResults = 0;
+        for (final Entry entry : subtree.values()) {
+            requestContext.checkIfCancelled(false);
+            if (scope.equals(SearchScope.WHOLE_SUBTREE) || entry.getName().isChildOf(dn)) {
+                if (matcher.matches(entry).toBoolean()) {
+                    if (sizeLimit > 0 && numberOfResults >= sizeLimit) {
+                        throw newErrorResult(newResult(ResultCode.SIZE_LIMIT_EXCEEDED));
+                    }
+                    numberOfResults++;
+                    boolean acceptMoreResults = sendEntry(attributeFilter, resultHandler, entry);
+                    if (!acceptMoreResults) {
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
     private <R extends Result> R addResultControls(final Request request, final Entry before,
             final Entry after, final R result) throws ErrorResultException {
         try {
diff --git a/opendj3/opendj-ldap-sdk/src/test/java/org/forgerock/opendj/ldap/MemoryBackendTestCase.java b/opendj3/opendj-ldap-sdk/src/test/java/org/forgerock/opendj/ldap/MemoryBackendTestCase.java
index a682970..c30e392 100644
--- a/opendj3/opendj-ldap-sdk/src/test/java/org/forgerock/opendj/ldap/MemoryBackendTestCase.java
+++ b/opendj3/opendj-ldap-sdk/src/test/java/org/forgerock/opendj/ldap/MemoryBackendTestCase.java
@@ -34,7 +34,9 @@
 import static org.forgerock.opendj.ldif.LDIFEntryReader.valueOfLDIFEntry;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 
 import org.forgerock.opendj.ldap.controls.AssertionRequestControl;
 import org.forgerock.opendj.ldap.controls.PermissiveModifyRequestControl;
@@ -43,6 +45,7 @@
 import org.forgerock.opendj.ldap.controls.PreReadRequestControl;
 import org.forgerock.opendj.ldap.controls.PreReadResponseControl;
 import org.forgerock.opendj.ldap.requests.Requests;
+import org.forgerock.opendj.ldap.responses.SearchResultEntry;
 import org.forgerock.opendj.ldif.ConnectionEntryReader;
 import org.forgerock.opendj.ldif.LDIFEntryReader;
 import org.testng.annotations.Test;
@@ -53,6 +56,8 @@
 @SuppressWarnings("javadoc")
 public class MemoryBackendTestCase extends SdkTestCase {
 
+    int numberOfEntriesInBackend;
+
     @Test
     public void testAdd() throws Exception {
         final Connection connection = getConnection();
@@ -410,11 +415,50 @@
     }
 
     @Test
+    public void testSearchOneLevelWithSizeLimit() throws Exception {
+        final Connection connection = getConnection();
+        final ConnectionEntryReader reader =
+                connection.search(Requests.newSearchRequest("dc=com", SearchScope.SINGLE_LEVEL, "(objectClass=*)").
+                        setSizeLimit(1));
+        assertThat(reader.readEntry()).isEqualTo(
+                valueOfLDIFEntry("dn: dc=example,dc=com", "objectClass: domain",
+                        "objectClass: top", "dc: example"));
+        try {
+            reader.hasNext();
+            failWasExpected(ErrorResultIOException.class);
+        } catch (ErrorResultIOException e) {
+            assertThat(e.getCause().getResult().getResultCode()).isEqualTo(ResultCode.SIZE_LIMIT_EXCEEDED);
+        }
+    }
+
+    @Test
     public void testSearchSubtree() throws Exception {
         final Connection connection = getConnection();
-        assertThat(
-                connection.searchSingleEntry("dc=example,dc=com", SearchScope.WHOLE_SUBTREE,
-                        "(uid=test1)")).isEqualTo(getUser1Entry());
+        assertThat(connection.searchSingleEntry("dc=example,dc=com", SearchScope.WHOLE_SUBTREE, "(uid=test1)")).
+                isEqualTo(getUser1Entry());
+    }
+
+    @Test
+    public void testSearchSubtreeReturnsAllEntries() throws Exception {
+        final Connection connection = getConnection();
+        Collection<SearchResultEntry> entries = new ArrayList<SearchResultEntry>();
+        connection.search(Requests.newSearchRequest("dc=com", SearchScope.WHOLE_SUBTREE, "(objectclass=*)"), entries);
+        assertThat(entries).hasSize(numberOfEntriesInBackend);
+    }
+
+    @Test
+    public void testSearchSubtreeWithSizeLimit() throws Exception {
+        final Connection connection = getConnection();
+        Collection<SearchResultEntry> entries = new ArrayList<SearchResultEntry>();
+        try {
+            connection.search(
+                    Requests.newSearchRequest("dc=example,dc=com", SearchScope.WHOLE_SUBTREE, "(objectClass=*)").
+                    setSizeLimit(2), entries);
+            failWasExpected(ErrorResultException.class);
+        } catch (ErrorResultException e) {
+            assertThat(e.getResult().getResultCode()).isEqualTo(ResultCode.SIZE_LIMIT_EXCEEDED);
+            assertThat(entries).hasSize(2);
+        }
     }
 
     @Test(expectedExceptions = EntryNotFoundException.class)
@@ -468,53 +512,63 @@
 
     private Connection getConnection() throws IOException {
         // @formatter:off
-        final MemoryBackend backend =
-                new MemoryBackend(new LDIFEntryReader(
-                        "dn: dc=com",
-                        "objectClass: domain",
-                        "objectClass: top",
-                        "dc: com",
-                        "",
-                        "dn: dc=example,dc=com",
-                        "objectClass: domain",
-                        "objectClass: top",
-                        "dc: example",
-                        "entryDN: dc=example,dc=com",
-                        "entryUUID: fc252fd9-b982-3ed6-b42a-c76d2546312c",
-                        "",
-                        "dn: ou=People,dc=example,dc=com",
-                        "objectClass: organizationalunit",
-                        "objectClass: top",
-                        "ou: People",
-                        "",
-                        "dn: uid=test1,ou=People,dc=example,dc=com",
-                        "objectClass: top",
-                        "objectClass: person",
-                        "uid: test1",
-                        "userpassword: password",
-                        "cn: test user 1",
-                        "sn: user 1",
-                        "entryDN: uid=test1,ou=people,dc=example,dc=com",
-                        "entryUUID: fc252fd9-b982-3ed6-b42a-c76d2546312c",
-                        "",
-                        "dn: uid=test2,ou=People,dc=example,dc=com",
-                        "objectClass: top",
-                        "objectClass: person",
-                        "uid: test2",
-                        "userpassword: password",
-                        "cn: test user 2",
-                        "sn: user 2",
-                        "",
-                        "dn: dc=xxx,dc=com",
-                        "objectClass: domain",
-                        "objectClass: top",
-                        "dc: xxx"
-                ));
+        String[] ldifEntries = new String[] {
+            "dn: dc=com",
+            "objectClass: domain",
+            "objectClass: top",
+            "dc: com",
+            "",
+            "dn: dc=example,dc=com",
+            "objectClass: domain",
+            "objectClass: top",
+            "dc: example",
+            "entryDN: dc=example,dc=com",
+            "entryUUID: fc252fd9-b982-3ed6-b42a-c76d2546312c",
+            "",
+            "dn: ou=People,dc=example,dc=com",
+            "objectClass: organizationalunit",
+            "objectClass: top",
+            "ou: People",
+            "",
+            "dn: uid=test1,ou=People,dc=example,dc=com",
+            "objectClass: top",
+            "objectClass: person",
+            "uid: test1",
+            "userpassword: password",
+            "cn: test user 1",
+            "sn: user 1",
+            "entryDN: uid=test1,ou=people,dc=example,dc=com",
+            "entryUUID: fc252fd9-b982-3ed6-b42a-c76d2546312c",
+            "",
+            "dn: uid=test2,ou=People,dc=example,dc=com",
+            "objectClass: top",
+            "objectClass: person",
+            "uid: test2",
+            "userpassword: password",
+            "cn: test user 2",
+            "sn: user 2",
+            "",
+            "dn: dc=xxx,dc=com",
+            "objectClass: domain",
+            "objectClass: top",
+            "dc: xxx"
+        };
         // @formatter:on
-
+        numberOfEntriesInBackend = getNumberOfEntries(ldifEntries);
+        final MemoryBackend backend = new MemoryBackend(new LDIFEntryReader(ldifEntries));
         return newInternalConnection(backend);
     }
 
+    private int getNumberOfEntries(String[] ldifEntries) {
+        int entries = 0;
+        for (int i = 0; i < ldifEntries.length; i++) {
+            if (ldifEntries[i].startsWith("dn: ")) {
+                entries++;
+            }
+        }
+        return entries;
+    }
+
     private Entry getUser1Entry() {
         return valueOfLDIFEntry("dn: uid=test1,ou=People,dc=example,dc=com", "objectClass: top",
                 "objectClass: person", "uid: test1", "userpassword: password", "cn: test user 1",

--
Gitblit v1.10.0