feat: add search with channels inspired by go-ldap#319

t2y · t2y · commit f7929a807f50 · 2023-06-01T13:37:13.000+09:00
diff --git a/examples_test.go b/examples_test.go
@@ -1,6 +1,7 @@
 package ldap
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
@@ -50,6 +51,34 @@ func ExampleConn_Search() {
 	}
 }
 
+// This example demonstrates how to search with channel
+func ExampleConn_SearchWithChannel() {
+	l, err := DialURL(fmt.Sprintf("%s:%d", "ldap.example.com", 389))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer l.Close()
+
+	searchRequest := NewSearchRequest(
+		"dc=example,dc=com", // The base dn to search
+		ScopeWholeSubtree, NeverDerefAliases, 0, 0, false,
+		"(&(objectClass=organizationalPerson))", // The filter to apply
+		[]string{"dn", "cn"},                    // A list attributes to retrieve
+		nil,
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ch := l.SearchWithChannel(ctx, searchRequest)
+	for res := range ch {
+		if res.Error != nil {
+			log.Fatalf("Error searching: %s", res.Error)
+		}
+		fmt.Printf("%s has DN %s\n", res.Entry.GetAttributeValue("cn"), res.Entry.DN)
+	}
+}
+
 // This example demonstrates how to start a TLS connection
 func ExampleConn_StartTLS() {
 	l, err := DialURL("ldap://ldap.example.com:389")
diff --git a/ldap_test.go b/ldap_test.go
@@ -1,6 +1,7 @@
 package ldap
 
 import (
+	"context"
 	"crypto/tls"
 	"testing"
 
@@ -344,3 +345,63 @@ func TestEscapeDN(t *testing.T) {
 		})
 	}
 }
+
+func TestSearchWithChannel(t *testing.T) {
+	l, err := DialURL(ldapServer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	searchRequest := NewSearchRequest(
+		baseDN,
+		ScopeWholeSubtree, DerefAlways, 0, 0, false,
+		filter[2],
+		attributes,
+		nil)
+
+	srs := make([]*Entry, 0)
+	ctx := context.Background()
+	for sr := range l.SearchWithChannel(ctx, searchRequest) {
+		if sr.Error != nil {
+			t.Fatal(err)
+		}
+		srs = append(srs, sr.Entry)
+	}
+
+	t.Logf("TestSearchWithChannel: %s -> num of entries = %d", searchRequest.Filter, len(srs))
+}
+
+func TestSearchWithChannelAndCancel(t *testing.T) {
+	l, err := DialURL(ldapServer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	searchRequest := NewSearchRequest(
+		baseDN,
+		ScopeWholeSubtree, DerefAlways, 0, 0, false,
+		filter[2],
+		attributes,
+		nil)
+
+	cancelNum := 10
+	srs := make([]*Entry, 0)
+	ctx, cancel := context.WithCancel(context.Background())
+	for sr := range l.SearchWithChannel(ctx, searchRequest) {
+		if sr.Error != nil {
+			t.Fatal(err)
+		}
+		srs = append(srs, sr.Entry)
+		if len(srs) == cancelNum {
+			cancel()
+		}
+	}
+	if len(srs) > cancelNum+2 {
+		// The cancel process is asynchronous,
+		// so a few entries after it canceled might be received
+		t.Errorf("Got entries %d, expected less than %d", len(srs), cancelNum+2)
+	}
+	t.Logf("TestSearchWithChannel: %s -> num of entries = %d", searchRequest.Filter, len(srs))
+}
diff --git a/search.go b/search.go
@@ -1,6 +1,7 @@
 package ldap
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"reflect"
@@ -375,6 +376,28 @@ func (s *SearchResult) appendTo(r *SearchResult) {
 	r.Controls = append(r.Controls, s.Controls...)
 }
 
+// SearchSingleResult holds the server's single response to a search request
+type SearchSingleResult struct {
+	// Entry is the returned entry
+	Entry *Entry
+	// Referral is the returned referral
+	Referral string
+	// Controls are the returned controls
+	Controls []Control
+	// Error is set when the search request was failed
+	Error error
+}
+
+// Print outputs a human-readable description
+func (s *SearchSingleResult) Print() {
+	s.Entry.Print()
+}
+
+// PrettyPrint outputs a human-readable description with indenting
+func (s *SearchSingleResult) PrettyPrint(indent int) {
+	s.Entry.PrettyPrint(indent)
+}
+
 // SearchRequest represents a search request to send to the server
 type SearchRequest struct {
 	BaseDN       string
@@ -559,6 +582,111 @@ func (l *Conn) Search(searchRequest *SearchRequest) (*SearchResult, error) {
 	}
 }
 
+// SearchWithChannel performs a search request and returns all search results
+// via the returned channel as soon as they are received. This means you get
+// all results until an error happens (or the search successfully finished),
+// e.g. for size / time limited requests all are recieved via the channel
+// until the limit is reached.
+func (l *Conn) SearchWithChannel(ctx context.Context, searchRequest *SearchRequest) chan *SearchSingleResult {
+	ch := make(chan *SearchSingleResult)
+	go func() {
+		defer close(ch)
+		if l.IsClosing() {
+			return
+		}
+
+		packet := ber.Encode(ber.ClassUniversal, ber.TypeConstructed, ber.TagSequence, nil, "LDAP Request")
+		packet.AppendChild(ber.NewInteger(ber.ClassUniversal, ber.TypePrimitive, ber.TagInteger, l.nextMessageID(), "MessageID"))
+		// encode search request
+		err := searchRequest.appendTo(packet)
+		if err != nil {
+			ch <- &SearchSingleResult{Error: err}
+			return
+		}
+		l.Debug.PrintPacket(packet)
+
+		msgCtx, err := l.sendMessage(packet)
+		if err != nil {
+			ch <- &SearchSingleResult{Error: err}
+			return
+		}
+		defer l.finishMessage(msgCtx)
+
+		foundSearchSingleResultDone := false
+		for !foundSearchSingleResultDone {
+			select {
+			case <-ctx.Done():
+				l.Debug.Printf("%d: %s", msgCtx.id, ctx.Err().Error())
+				return
+			default:
+				l.Debug.Printf("%d: waiting for response", msgCtx.id)
+				packetResponse, ok := <-msgCtx.responses
+				if !ok {
+					err := NewError(ErrorNetwork, errors.New("ldap: response channel closed"))
+					ch <- &SearchSingleResult{Error: err}
+					return
+				}
+				packet, err = packetResponse.ReadPacket()
+				l.Debug.Printf("%d: got response %p", msgCtx.id, packet)
+				if err != nil {
+					ch <- &SearchSingleResult{Error: err}
+					return
+				}
+
+				if l.Debug {
+					if err := addLDAPDescriptions(packet); err != nil {
+						ch <- &SearchSingleResult{Error: err}
+						return
+					}
+					ber.PrintPacket(packet)
+				}
+
+				switch packet.Children[1].Tag {
+				case ApplicationSearchResultEntry:
+					entry := new(Entry)
+					entry.DN = packet.Children[1].Children[0].Value.(string)
+					for _, child := range packet.Children[1].Children[1].Children {
+						attr := new(EntryAttribute)
+						attr.Name = child.Children[0].Value.(string)
+						for _, value := range child.Children[1].Children {
+							attr.Values = append(attr.Values, value.Value.(string))
+							attr.ByteValues = append(attr.ByteValues, value.ByteValue)
+						}
+						entry.Attributes = append(entry.Attributes, attr)
+					}
+					ch <- &SearchSingleResult{Entry: entry}
+
+				case ApplicationSearchResultDone:
+					if err := GetLDAPError(packet); err != nil {
+						ch <- &SearchSingleResult{Error: err}
+						return
+					}
+					if len(packet.Children) == 3 {
+						result := &SearchSingleResult{}
+						for _, child := range packet.Children[2].Children {
+							decodedChild, err := DecodeControl(child)
+							if err != nil {
+								werr := fmt.Errorf("failed to decode child control: %w", err)
+								ch <- &SearchSingleResult{Error: werr}
+								return
+							}
+							result.Controls = append(result.Controls, decodedChild)
+						}
+						ch <- result
+					}
+					foundSearchSingleResultDone = true
+
+				case ApplicationSearchResultReference:
+					ref := packet.Children[1].Children[0].Value.(string)
+					ch <- &SearchSingleResult{Referral: ref}
+				}
+			}
+		}
+		l.Debug.Printf("%d: returning", msgCtx.id)
+	}()
+	return ch
+}
+
 // unpackAttributes will extract all given LDAP attributes and it's values
 // from the ber.Packet
 func unpackAttributes(children []*ber.Packet) []*EntryAttribute {
diff --git a/v3/examples_test.go b/v3/examples_test.go
@@ -1,6 +1,7 @@
 package ldap
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
@@ -50,6 +51,34 @@ func ExampleConn_Search() {
 	}
 }
 
+// This example demonstrates how to search with channel
+func ExampleConn_SearchWithChannel() {
+	l, err := DialURL(fmt.Sprintf("%s:%d", "ldap.example.com", 389))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer l.Close()
+
+	searchRequest := NewSearchRequest(
+		"dc=example,dc=com", // The base dn to search
+		ScopeWholeSubtree, NeverDerefAliases, 0, 0, false,
+		"(&(objectClass=organizationalPerson))", // The filter to apply
+		[]string{"dn", "cn"},                    // A list attributes to retrieve
+		nil,
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	ch := l.SearchWithChannel(ctx, searchRequest)
+	for res := range ch {
+		if res.Error != nil {
+			log.Fatalf("Error searching: %s", res.Error)
+		}
+		fmt.Printf("%s has DN %s\n", res.Entry.GetAttributeValue("cn"), res.Entry.DN)
+	}
+}
+
 // This example demonstrates how to start a TLS connection
 func ExampleConn_StartTLS() {
 	l, err := DialURL("ldap://ldap.example.com:389")
diff --git a/v3/ldap_test.go b/v3/ldap_test.go
@@ -1,6 +1,7 @@
 package ldap
 
 import (
+	"context"
 	"crypto/tls"
 	"testing"
 
@@ -344,3 +345,63 @@ func TestEscapeDN(t *testing.T) {
 		})
 	}
 }
+
+func TestSearchWithChannel(t *testing.T) {
+	l, err := DialURL(ldapServer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	searchRequest := NewSearchRequest(
+		baseDN,
+		ScopeWholeSubtree, DerefAlways, 0, 0, false,
+		filter[2],
+		attributes,
+		nil)
+
+	srs := make([]*Entry, 0)
+	ctx := context.Background()
+	for sr := range l.SearchWithChannel(ctx, searchRequest) {
+		if sr.Error != nil {
+			t.Fatal(err)
+		}
+		srs = append(srs, sr.Entry)
+	}
+
+	t.Logf("TestSearchWithChannel: %s -> num of entries = %d", searchRequest.Filter, len(srs))
+}
+
+func TestSearchWithChannelAndCancel(t *testing.T) {
+	l, err := DialURL(ldapServer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	searchRequest := NewSearchRequest(
+		baseDN,
+		ScopeWholeSubtree, DerefAlways, 0, 0, false,
+		filter[2],
+		attributes,
+		nil)
+
+	cancelNum := 10
+	srs := make([]*Entry, 0)
+	ctx, cancel := context.WithCancel(context.Background())
+	for sr := range l.SearchWithChannel(ctx, searchRequest) {
+		if sr.Error != nil {
+			t.Fatal(err)
+		}
+		srs = append(srs, sr.Entry)
+		if len(srs) == cancelNum {
+			cancel()
+		}
+	}
+	if len(srs) > cancelNum+2 {
+		// The cancel process is asynchronous,
+		// so a few entries after it canceled might be received
+		t.Errorf("Got entries %d, expected less than %d", len(srs), cancelNum+2)
+	}
+	t.Logf("TestSearchWithChannel: %s -> num of entries = %d", searchRequest.Filter, len(srs))
+}
diff --git a/v3/search.go b/v3/search.go
