clean up code to use timeout instead of two keys

3 files changed, 68 insertions, 55 deletions

diff --git a/nova/api/ec2/__init__.py b/nova/api/ec2/__init__.py
index 19eb666cd..381b0e871 100644
--- a/nova/api/ec2/__init__.py
+++ b/nova/api/ec2/__init__.py
@@ -22,7 +22,6 @@ Starting point for routing EC2 requests.
 
 import logging
 import routes
-import time
 import webob
 import webob.dec
 import webob.exc
@@ -44,6 +43,8 @@ flags.DEFINE_integer('lockout_attempts', 5,
                      'Number of failed auths before lockout.')
 flags.DEFINE_integer('lockout_minutes', 15,
                      'Number of minutes to lockout if triggered.')
+flags.DEFINE_integer('lockout_window', 15,
+                     'Number of minutes for lockout window.')
 flags.DEFINE_list('lockout_memcached_servers', None,
                   'Memcached servers or None for in process cache.')
 
@@ -53,7 +54,6 @@ _log.setLevel(logging.DEBUG)
 
 
 class API(wsgi.Middleware):
-
     """Routing for all EC2 API requests."""
 
     def __init__(self):
@@ -63,57 +63,53 @@ class API(wsgi.Middleware):
 
 
 class Lockout(wsgi.Middleware):
-    """Only allow x failed auths in a y minute period.
+    """Lockout for x minutes on y failed auths in a z minute period.
 
-    x = lockout_attempts flag
-    y = lockout_timeout flag
+    x = lockout_timeout flag
+    y = lockout_window flag
+    z = lockout_attempts flag
 
     Uses memcached if lockout_memcached_servers flag is set, otherwise it
     uses a very simple in-proccess cache. Due to the simplicity of
-    the implementation, the timeout window is reset with every failed
-    request, so it actually blocks if there are x failed logins with no
-    more than y minutes between any two failures.
+    the implementation, the timeout window is started with the first
+    failed request, so it will block if there are x failed logins within
+    that period.
 
     There is a possible race condition where simultaneous requests could
     sneak in before the lockout hits, but this is extremely rare and would
     only result in a couple of extra failed attempts."""
 
-    def __init__(self, application, time_fn=time.time):
-        """The middleware can use a custom time function for testing."""
-        self.time_fn = time_fn
+    def __init__(self, application, time_fn=None):
+        """middleware can pass a custom time function to fake for testing."""
         if FLAGS.lockout_memcached_servers:
             import memcache
+            self.mc = memcache.Client(FLAGS.lockout_memcached_servers,
+                                      debug=0)
         else:
-            from nova import fakememcache as memcache
-        self.mc = memcache.Client(FLAGS.lockout_memcached_servers, debug=0)
+            from nova import fakememcache
+            self.mc = fakememcache.Client(time_fn=time_fn)
         super(Lockout, self).__init__(application)
 
     @webob.dec.wsgify
     def __call__(self, req):
         access_key = req.params['AWSAccessKeyId']
-        failures_key = "%s-failures" % access_key
-        last_key = "%s-last" % access_key
-        now = self.time_fn()
-        timeout = now - FLAGS.lockout_minutes * 60
-        # NOTE(vish): To use incr, failures has to be a string.
+        failures_key = "authfailures-%s" % access_key
         failures = int(self.mc.get(failures_key) or 0)
-        last = self.mc.get(last_key)
-        if (failures and failures >= FLAGS.lockout_attempts
-            and last > timeout):
-            self.mc.set(last_key, now)
+        if failures >= FLAGS.lockout_attempts:
             detail = "Too many failed authentications."
             raise webob.exc.HTTPForbidden(detail=detail)
         res = req.get_response(self.application)
         if res.status_int == 403:
-            if last > timeout:
-                failures = int(self.mc.incr(failures_key))
-                if failures >= FLAGS.lockout_attempts:
-                    _log.warn('Access key %s has had %d failed authentications'
-                              ' and will be locked out for %d minutes.' %
-                              (access_key, failures, FLAGS.lockout_minutes))
-            else:
-                self.mc.set(failures_key, '1')
-            self.mc.set(last_key, now)
+            failures = self.mc.incr(failures_key)
+            if failures is None:
+                # NOTE(vish): To use incr, failures has to be a string.
+                self.mc.set(failures_key, '1', time=FLAGS.lockout_window * 60)
+            elif failures >= FLAGS.lockout_attempts:
+                _log.warn('Access key %s has had %d failed authentications'
+                          ' and will be locked out for %d minutes.' %
+                          (access_key, failures, FLAGS.lockout_minutes))
+                self.mc.set(failures_key, str(failures),
+                            time=FLAGS.lockout_minutes * 60)
         return res
 
 
diff --git a/nova/fakememcache.py b/nova/fakememcache.py
index 0b2e3b6c1..0b4037ef6 100644
--- a/nova/fakememcache.py
+++ b/nova/fakememcache.py
@@ -18,33 +18,43 @@
 
 """Super simple fake memcache client."""
 
+import time
+
 
 class Client(object):
     """Replicates a tiny subset of memcached client interface."""
-    __cache = {}
 
-    def __init__(self, *args, **kwargs):
-        """Ignores all constructor params."""
-        pass
+    def __init__(self, time_fn=time.time, *args, **kwargs):
+        """Time fn is to allow testing through a custom function"""
+        self.time_fn = time_fn
+        self.cache = {}
 
     def get(self, key):
         """Retrieves the value for a key or None."""
-        return self.__cache.get(key, None)
+        (timeout, value) = self.cache.get(key, (0, None))
+        if timeout == 0 or self.time_fn() < timeout:
+            return value
+        return None
 
-    def set(self, key, value):
+    def set(self, key, value, time=0, min_compress_len=0):
         """Sets the value for a key."""
-        self.__cache[key] = value
+        timeout = 0
+        if time != 0:
+            timeout = self.time_fn() + time
+        self.cache[key] = (timeout, value)
         return True
 
-    def add(self, key, value):
+    def add(self, key, value, time=0, min_compress_len=0):
         """Sets the value for a key if it doesn't exist."""
-        if key in self.__cache:
+        if not self.get(key) is None:
             return False
-        return self.set(key, value)
+        return self.set(key, value, time, min_compress_len)
 
     def incr(self, key, delta=1):
         """Increments the value for a key."""
-        if not key in self.__cache:
-            return 0
-        self.__cache[key] = str(int(self.__cache[key]) + 1)
-        return self.__cache[key]
+        value = self.get(key)
+        if value is None:
+            return None
+        new_value = int(value) + delta
+        self.cache[key] = (self.cache[key][0], str(new_value))
+        return new_value
diff --git a/nova/tests/middleware_unittest.py b/nova/tests/middleware_unittest.py
index bbbd4a5a7..61a790c1f 100644
--- a/nova/tests/middleware_unittest.py
+++ b/nova/tests/middleware_unittest.py
@@ -48,9 +48,9 @@ class LockoutTestCase(test.TrialTestCase):
         """Helper method to force timeouts."""
         return self.local_time
 
-    def _trigger_lockout(self, access_key):
-        """Send x failed requests where x = lockout_attempts."""
-        for i in xrange(FLAGS.lockout_attempts):
+    def _send_bad_attempts(self, access_key, num_attempts=1):
+        """Fail x."""
+        for i in xrange(num_attempts):
             req = webob.Request.blank('/?AWSAccessKeyId=%s&die=1' % access_key)
             self.assertEqual(req.get_response(self.lockout).status_int, 403)
 
@@ -59,24 +59,31 @@ class LockoutTestCase(test.TrialTestCase):
         req = webob.Request.blank('/?AWSAccessKeyId=%s' % access_key)
         return (req.get_response(self.lockout).status_int == 403)
 
-    def _timeout(self):
+    def _advance_time(self, time):
         """Increment time to 1 second past the lockout."""
-        self.local_time = 1 + self.local_time + FLAGS.lockout_minutes * 60
+        self.local_time = self.local_time + time
 
     def test_lockout(self):
-        self._trigger_lockout('test')
+        self._send_bad_attempts('test', FLAGS.lockout_attempts)
         self.assertTrue(self._is_locked_out('test'))
 
     def test_timeout(self):
-        self._trigger_lockout('test')
+        self._send_bad_attempts('test', FLAGS.lockout_attempts)
         self.assertTrue(self._is_locked_out('test'))
-        self._timeout()
+        self._advance_time(FLAGS.lockout_minutes * 60)
         self.assertFalse(self._is_locked_out('test'))
 
     def test_multiple_keys(self):
-        self._trigger_lockout('test1')
+        self._send_bad_attempts('test1', FLAGS.lockout_attempts)
         self.assertTrue(self._is_locked_out('test1'))
         self.assertFalse(self._is_locked_out('test2'))
-        self._timeout()
+        self._advance_time(FLAGS.lockout_minutes * 60)
         self.assertFalse(self._is_locked_out('test1'))
         self.assertFalse(self._is_locked_out('test2'))
+
+    def test_window_timeout(self):
+        self._send_bad_attempts('test', FLAGS.lockout_attempts - 1)
+        self.assertFalse(self._is_locked_out('test'))
+        self._advance_time(FLAGS.lockout_window * 60)
+        self._send_bad_attempts('test', FLAGS.lockout_attempts - 1)
+        self.assertFalse(self._is_locked_out('test'))