/
search.go
139 lines (118 loc) · 4.31 KB
/
search.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package ldapsession
import (
"errors"
"fmt"
"github.com/go-ldap/ldap/v3"
"github.com/sirupsen/logrus"
)
func (w *LDAPSession) MakeSimpleSearchRequest(filter string, attrs []string) *ldap.SearchRequest {
return ldap.NewSearchRequest(
w.BaseDN,
ldap.ScopeWholeSubtree,
ldap.NeverDerefAliases,
0, 0, false,
filter,
attrs,
nil)
}
// GetPagedSearchResults is a synchronous operation that will populate and return an ldap.SearchResult object
func (w *LDAPSession) GetPagedSearchResults(request *ldap.SearchRequest) (result *ldap.SearchResult, err error) {
w.Log.WithFields(logrus.Fields{"filter": request.Filter, "attributes": request.Attributes}).Infof("sending LDAP search request")
return w.LConn.SearchWithPaging(request, 1000)
}
func (w *LDAPSession) GetSearchResults(request *ldap.SearchRequest) (result *ldap.SearchResult, err error) {
w.Log.WithFields(logrus.Fields{"filter": request.Filter, "attributes": request.Attributes}).Infof("sending LDAP search request")
return w.LConn.Search(request)
}
func (w *LDAPSession) ManualWriteSearchResultsToChan(results *ldap.SearchResult) {
w.Log.Debugf("received search results, writing %d entries to channel", len(results.Entries))
defer w.CloseChannels()
for _, entry := range results.Entries {
w.Channels.Entries <- entry
}
for _, referral := range results.Referrals {
w.Channels.Referrals <- referral
}
for _, control := range results.Controls {
w.Channels.Controls <- control
}
}
// ExecuteSearchRequest performs a paged search and writes results to the LDAPsession's defined results channel.
// it only returns an err
func (w *LDAPSession) ExecuteSearchRequest(searchRequest *ldap.SearchRequest) error {
w.Log.WithFields(logrus.Fields{"filter": searchRequest.Filter, "attributes": searchRequest.Attributes}).Infof("sending LDAP search request")
if w.Channels == nil {
return fmt.Errorf("no channels defined. Call SetChannels first, or use GetPagedSearchResults instead")
}
defer func() {
w.Log.Debugf("search finished. closing channels...")
w.CloseChannels()
}()
// basically a re-implementation of the standard function: https://github.com/go-ldap/ldap/blob/master/v3/search.go#L253
// but writes entries to a channel as it gets them instead of waiting for all pages to complete
var pagingControl *ldap.ControlPaging
control := ldap.FindControl(searchRequest.Controls, ldap.ControlTypePaging)
if control == nil {
pagingControl = ldap.NewControlPaging(w.PageSize)
searchRequest.Controls = append(searchRequest.Controls, pagingControl)
} else {
castControl, ok := control.(*ldap.ControlPaging)
if !ok {
return fmt.Errorf("expected paging control to be of type *ControlPaging, got %v", control)
}
if castControl.PagingSize != w.PageSize {
return fmt.Errorf("paging size given in search request (%d) conflicts with size given in search call (%d)", castControl.PagingSize, w.PageSize)
}
pagingControl = castControl
}
pageNumber := 0
PagedSearch:
for {
select {
case <-w.ctx.Done():
w.Log.Warn("cancel received. aborting remaining pages")
return nil
default:
w.Log.Debugf("making paged request...\n")
result, err := w.LConn.Search(searchRequest)
w.Log.Debugf("Looking for Paging Control...\n")
pageNumber++
if err != nil {
return err
}
if result == nil {
return ldap.NewError(ldap.ErrorNetwork, errors.New("ldap: packet not received"))
}
for _, entry := range result.Entries {
w.Channels.Entries <- entry
}
w.Log.Infof("Received page %d with %d LDAP entries...", pageNumber, len(result.Entries))
for _, referral := range result.Referrals {
w.Channels.Referrals <- referral
}
for _, control := range result.Controls {
w.Channels.Controls <- control
}
w.Log.Debugf("Looking for Paging Control...")
pagingResult := ldap.FindControl(result.Controls, ldap.ControlTypePaging)
if pagingResult == nil {
pagingControl = nil
w.Log.Debugf("Could not find paging control. Breaking...")
break PagedSearch
}
cookie := pagingResult.(*ldap.ControlPaging).Cookie
if len(cookie) == 0 {
pagingControl = nil
w.Log.Debugf("Could not find cookie. Breaking...")
break PagedSearch
}
pagingControl.SetCookie(cookie)
}
}
if pagingControl != nil {
w.Log.Debugf("Abandoning Paging...")
pagingControl.PagingSize = 0
w.LConn.Search(searchRequest)
}
return nil
}