Use a ThreadMappedPool for memcache connections

Now requires pylibmc.

authserv/app_common.py

@@ -4,7 +4,8 @@ from authserv import auth
 from authserv import protocol
 from authserv.ratelimit import *
 
-_blacklisted = (protocol.ERR_AUTHENTICATION_FAILURE, 'blacklisted', None)
+_user_blacklisted = (protocol.ERR_AUTHENTICATION_FAILURE, 'user_blacklisted', None)
+_ip_blacklisted = (protocol.ERR_AUTHENTICATION_FAILURE, 'ip_blacklisted', None)
 
 
 def check_ratelimit(request, username, source_ip):
@@ -50,17 +51,25 @@ def _validate_username(username):
 
 def do_auth(username, service, shard, password, otp_token, source_ip,
             password_only=False):
+    with current_app.memcache.reserve() as mc:
+        return _do_auth(mc, username, service, shard, password, otp_token,
+                        source_ip, password_only)
+
+
+def _do_auth(mc, username, service, shard, password, otp_token, source_ip,
+            password_only):
     # Username must be an ASCII string.
     bl = AuthBlackList(current_app.config.get('BLACKLIST_COUNT', 5),
                        current_app.config.get('BLACKLIST_PERIOD', 600),
-                       current_app.config.get('BLACKLIST_TIME', 6*3600))
+                       current_app.config.get('BLACKLIST_TIME', 6*3600),
+                       mc)
     if current_app.config.get('ENABLE_BLACKLIST'):
         if bl.is_blacklisted('u', username):
-            return _blacklisted
+            return _user_blacklisted
         if (source_ip
             and not whitelisted(source_ip, current_app.config.get('SOURCE_IP_WHITELIST'))
             and bl.is_blacklisted('ip', source_ip)):
-            return _blacklisted
+            return _ip_blacklisted
 
     retval = protocol.ERR_AUTHENTICATION_FAILURE
     errmsg = 'user does not exist'

authserv/ratelimit.py

@@ -76,13 +76,14 @@ def ratelimit_http_request(req, value, tag='', count=0, period=0):
 
     rl = RateLimit(count, period)
     key = key_sep.join([req.method, req.path, value])
-    return rl.check(current_app.memcache, key)
+    with current_app.memcache.reserve() as mc:
+        return rl.check(mc, key)
 
 
 class BlackList(object):
     """Block requests once a rate limit is triggered."""
 
-    prefix = 'authserv/b:'
+    prefix = 'authserv/b/'
 
     def __init__(self, count, period, ttl):
         self.rl = RateLimit(count, period)
@@ -99,17 +100,20 @@ class BlackList(object):
             mc.set(key, 'true', time=self.ttl)
 
 
-class AuthBlackList(BlackList):
+class AuthBlackList(object):
+
+    def __init__(self, count, period, ttl, mc):
+        self.blacklist = BlackList(count, period, ttl)
+        self.mc = mc
 
     def is_blacklisted(self, tag, value):
         if not value:
             return False
         key = key_sep.join([tag, value])
-        return not self.check(current_app.memcache, key)
+        return not self.blacklist.check(self.mc, key)
 
     def auth_failure(self, tag, value):
         if not value:
             return
         key = key_sep.join([tag, value])
-        self.incr(current_app.memcache, key)
-
+        self.blacklist.incr(self.mc, key)

authserv/server.py

@@ -7,6 +7,7 @@ import logging
 import logging.handlers
 import optparse
 import os
+import pylibmc
 import signal
 import sys
 from authserv import app_main
@@ -27,14 +28,9 @@ def create_app(app, userdb=None, mc=None):
     app.userdb = userdb
 
     if not mc:
-        try:
-            import pylibmc
-            mc = pylibmc.Client(
-                app.config['MEMCACHE_ADDR'], binary=True)
-        except ImportError:
-            import memcache
-            mc = memcache.Client(
-                app.config['MEMCACHE_ADDR'], debug=0)
+        client = pylibmc.Client(
+            app.config['MEMCACHE_ADDR'], binary=True)
+        mc = pylibmc.ThreadMappedPool(client)
     app.memcache = mc
 
     return app

authserv/test/__init__.py

+import contextlib
 import crypt
 import logging
 import os
@@ -39,6 +40,16 @@ class FakeMemcache(object):
         return value
 
 
+class FakeMemcachePool(object):
+
+    def __init__(self, t=None):
+        self.mc = FakeMemcache(t)
+
+    @contextlib.contextmanager
+    def reserve(self):
+        yield self.mc
+
+
 class FakeUser(model.User):
 
     def __init__(self, username, password=None, asps=None, otp_key=None, shard=None):

authserv/test/test_app_main.py

@@ -21,7 +21,7 @@ class ServerTest(unittest.TestCase):
             }
         app = server.create_app(app_main.app,
                                 userdb=FakeUserDb(self.users),
-                                mc=FakeMemcache(_time))
+                                mc=FakeMemcachePool(_time))
         app.config.update({
                 'TESTING': True,
                 'DEBUG': True,

authserv/test/test_app_nginx.py

@@ -17,7 +17,7 @@ class ServerTest(unittest.TestCase):
             }
         app = server.create_app(app_nginx.app,
                                 userdb=FakeUserDb(self.users),
-                                mc=FakeMemcache(_time))
+                                mc=FakeMemcachePool(_time))
         app.config.update({
                 'TESTING': True,
                 'DEBUG': True,

authserv/test/test_integration.py

@@ -63,7 +63,7 @@ class SSLServerTest(unittest.TestCase):
         def _runserver():
             app = server.create_app(app_main.app,
                                     userdb=FakeUserDb(cls.users),
-                                    mc=FakeMemcache(time.time))
+                                    mc=FakeMemcachePool(time.time))
             app.config.update({
                 'TESTING': True,
                 'DEBUG': True,

authserv/test/test_ratelimit.py

@@ -8,25 +8,27 @@ class RateLimitTest(unittest.TestCase):
         self.tick = 0
         def _time():
             return self.tick
-        self.mc = FakeMemcache(_time)
+        self.mc = FakeMemcachePool(_time)
         self.rl = RateLimit(100, 10)
 
     def test_ratelimit_pass(self):
         n = 200
         ok = 0
-        for i in xrange(n):
-            self.tick = i
-            if self.rl.check(self.mc, 'key'):
-                ok += 1
+        with self.mc.reserve() as mc:
+            for i in xrange(n):
+                self.tick = i
+                if self.rl.check(mc, 'key'):
+                    ok += 1
         self.assertEquals(n, ok)
 
     def test_ratelimit_fail(self):
         n = 200
         ok = 0
-        for i in xrange(n):
-            self.tick = i / 20
-            if self.rl.check(self.mc, 'key'):
-                ok += 1
+        with self.mc.reserve() as mc:
+            for i in xrange(n):
+                self.tick = i / 20
+                if self.rl.check(mc, 'key'):
+                    ok += 1
         self.assertEquals(100, ok)
 
 

setup.py

@@ -4,12 +4,12 @@ from setuptools import setup, find_packages
 
 setup(
     name="authserv",
-    version="0.1.1",
+    version="0.1.2",
     description="Authentication server",
     author="Autistici/Inventati",
     author_email="info@autistici.org",
     url="https://git.autistici.org/ai/authserv",
-    install_requires=["gevent", "python-ldap", "Flask", "python-memcached"],
+    install_requires=["gevent", "python-ldap", "Flask", "pylibmc"],
     setup_requires=[],
     zip_safe=False,
     packages=find_packages(),