1# ***** BEGIN LICENSE BLOCK *****
2# Version: MPL 1.1/GPL 2.0/LGPL 2.1
3#
4# The contents of this file are subject to the Mozilla Public License Version
5# 1.1 (the "License"); you may not use this file except in compliance with
6# the License. You may obtain a copy of the License at
7# http://www.mozilla.org/MPL/
8#
9# Software distributed under the License is distributed on an "AS IS" basis,
10# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
11# for the specific language governing rights and limitations under the
12# License.
13#
14# The Original Code is Sync Server
15#
16# The Initial Developer of the Original Code is the Mozilla Foundation.
17# Portions created by the Initial Developer are Copyright (C) 2010
18# the Initial Developer. All Rights Reserved.
19#
20# Contributor(s):
21#   Tarek Ziade (tarek@mozilla.com)
22#
23# Alternatively, the contents of this file may be used under the terms of
24# either the GNU General Public License Version 2 or later (the "GPL"), or
25# the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
26# in which case the provisions of the GPL or the LGPL are applicable instead
27# of those above. If you wish to allow use of your version of this file only
28# under the terms of either the GPL or the LGPL, and not to allow others to
29# use your version of this file under the terms of the MPL, indicate your
30# decision by deleting the provisions above and replace them with the notice
31# and other provisions required by the GPL or the LGPL. If you do not delete
32# the provisions above, a recipient may use your version of this file under
33# the terms of any one of the MPL, the GPL or the LGPL.
34#
35# ***** END LICENSE BLOCK *****
36import threading
37import time
38import unittest
39
40import ldap
41
42import ldappool
43
44# patching StateConnector
45ldappool.StateConnector.users = {
46    'uid=tarek,ou=users,dc=mozilla':
47        {'uidNumber': ['1'],
48         'account-enabled': ['Yes'],
49         'mail': ['tarek@mozilla.com'],
50         'cn': ['tarek']},
51    'cn=admin,dc=mozilla': {'cn': ['admin'],
52                            'mail': ['admin'],
53                            'uidNumber': ['100']}}
54
55
56def _simple_bind(self, who='', cred='', *args):
57    self.connected = True
58    self.who = who
59    self.cred = cred
60
61
62ldappool.StateConnector.simple_bind_s = _simple_bind
63
64
65def _search(self, dn, *args, **kw):
66    if dn in self.users:
67        return [(dn, self.users[dn])]
68    elif dn == 'ou=users,dc=mozilla':
69        uid = kw['filterstr'].split('=')[-1][:-1]
70        for dn_, value in self.users.items():
71            if value['uidNumber'][0] != uid:
72                continue
73            return [(dn_, value)]
74
75    raise ldap.NO_SUCH_OBJECT
76
77
78ldappool.StateConnector.search_s = _search
79
80
81def _add(self, dn, user):
82    self.users[dn] = {}
83    for key, value in user:
84        if not isinstance(value, list):
85            value = [value]
86        self.users[dn][key] = value
87
88    return ldap.RES_ADD, ''
89
90
91ldappool.StateConnector.add_s = _add
92
93
94def _modify(self, dn, user):
95    if dn in self.users:
96        for type_, key, value in user:
97            if not isinstance(value, list):
98                value = [value]
99            self.users[dn][key] = value
100    return ldap.RES_MODIFY, ''
101
102
103ldappool.StateConnector.modify_s = _modify
104
105
106def _delete(self, dn):
107    if dn in self.users:
108        del self.users[dn]
109    return ldap.RES_DELETE, ''
110
111
112ldappool.StateConnector.delete_s = _delete
113
114
115class LDAPWorker(threading.Thread):
116
117    def __init__(self, pool):
118        threading.Thread.__init__(self)
119        self.pool = pool
120        self.results = []
121
122    def run(self):
123        dn = 'cn=admin,dc=mozilla'
124        for i in range(10):
125            with self.pool.connection() as conn:
126                res = conn.search_s(dn, ldap.SCOPE_BASE,
127                                    attrlist=['cn'])
128                self.results.append(res)
129
130
131class TestLDAPSQLAuth(unittest.TestCase):
132
133    def test_ctor_args(self):
134        pool = ldappool.ConnectionManager('ldap://localhost', use_tls=True)
135        self.assertEqual(pool.use_tls, True)
136
137    def test_pool(self):
138        dn = 'uid=adminuser,ou=logins,dc=mozilla'
139        passwd = 'adminuser'
140        pool = ldappool.ConnectionManager('ldap://localhost', dn, passwd)
141        workers = [LDAPWorker(pool) for i in range(10)]
142
143        for worker in workers:
144            worker.start()
145
146        for worker in workers:
147            worker.join()
148            self.assertEqual(len(worker.results), 10)
149            cn = worker.results[0][0][1]['cn']
150            self.assertEqual(cn, ['admin'])
151
152    def test_pool_full(self):
153        dn = 'uid=adminuser,ou=logins,dc=mozilla'
154        passwd = 'adminuser'
155        pool = ldappool.ConnectionManager(
156            'ldap://localhost', dn, passwd, size=1, retry_delay=1.,
157            retry_max=5, use_pool=True)
158
159        class Worker(threading.Thread):
160
161            def __init__(self, pool, duration):
162                threading.Thread.__init__(self)
163                self.pool = pool
164                self.duration = duration
165
166            def run(self):
167                with self.pool.connection() as conn:  # NOQA
168                    time.sleep(self.duration)
169
170        def tryit():
171            time.sleep(0.1)
172            with pool.connection() as conn:  # NOQA
173                pass
174
175        # an attempt on a full pool should eventually work
176        # because the connector is reused
177        for i in range(10):
178            tryit()
179
180        # we have 1 non-active connector now
181        self.assertEqual(len(pool), 1)
182
183        # an attempt with a full pool should succeed if a
184        # slot gets freed in less than one second.
185        worker1 = Worker(pool, .4)
186        worker1.start()
187
188        try:
189            tryit()
190        finally:
191            worker1.join()
192
193        # an attempt with a full pool should fail
194        # if no slot gets freed in less than one second.
195        worker1 = Worker(pool, 1.1)
196        worker1.start()
197        try:
198            self.assertRaises(ldappool.MaxConnectionReachedError, tryit)
199        finally:
200            worker1.join()
201
202        # we still have one active connector
203        self.assertEqual(len(pool), 1)
204
205    def test_pool_cleanup(self):
206        dn = 'uid=adminuser,ou=logins,dc=mozilla'
207        passwd = 'adminuser'
208        pool = ldappool.ConnectionManager('ldap://localhost', dn, passwd,
209                                          size=1, use_pool=True)
210        with pool.connection('bind1') as conn:  # NOQA
211            pass
212
213        with pool.connection('bind2') as conn:  # NOQA
214
215            pass
216
217        # the second call should have removed the first conn
218        self.assertEqual(len(pool), 1)
219
220    def test_pool_reuse(self):
221        dn = 'uid=adminuser,ou=logins,dc=mozilla'
222        passwd = 'adminuser'
223        pool = ldappool.ConnectionManager('ldap://localhost', dn, passwd,
224                                          use_pool=True)
225
226        with pool.connection() as conn:
227            self.assertTrue(conn.active)
228
229        self.assertFalse(conn.active)
230        self.assertTrue(conn.connected)
231
232        with pool.connection() as conn2:
233            pass
234
235        self.assertTrue(conn is conn2)
236
237        with pool.connection() as conn:
238            conn.connected = False
239
240        with pool.connection() as conn2:
241            pass
242
243        self.assertTrue(conn is not conn2)
244
245        # same bind and password: reuse
246        with pool.connection('bind', 'passwd') as conn:
247            self.assertTrue(conn.active)
248
249        self.assertFalse(conn.active)
250        self.assertTrue(conn.connected)
251
252        with pool.connection('bind', 'passwd') as conn2:
253            pass
254
255        self.assertTrue(conn is conn2)
256
257        # same bind different password: rebind !
258        with pool.connection('bind', 'passwd') as conn:
259            self.assertTrue(conn.active)
260
261        self.assertFalse(conn.active)
262        self.assertTrue(conn.connected)
263
264        with pool.connection('bind', 'passwd2') as conn2:
265            pass
266
267        self.assertTrue(conn is conn2)
268