enforce type correctness for memcache

authserv/ratelimit.py

@@ -18,6 +18,12 @@ _whitelist_rx = [
     for x in WHITELIST]
 
 
+def _tostr(s):
+    if isinstance(s, unicode):
+        return s.encode('utf-8')
+    return s
+
+
 def whitelisted(ip):
     for rx in _whitelist_rx:
         if rx.search(ip):
@@ -79,7 +85,7 @@ class RateLimit(object):
         self.period = period
 
     def check(self, mc, key):
-        key = self.prefix + key
+        key = _tostr(self.prefix + key)
         try:
             result = mc.incr(key)
         except:
@@ -109,7 +115,7 @@ def ratelimit_http_request(key_fn, count=0, period=0):
     def decoratorfn(fn):
         @functools.wraps(fn)
         def _ratelimit(*args, **kwargs):
-            key = key_fn(fn.__name__, args)
+            key = _tostr(key_fn(fn.__name__, args))
             if key and not rl.check(app.memcache, key):
                 app.logger.debug('ratelimited: %s', key)
                 abort(503)
@@ -128,12 +134,12 @@ class BlackList(object):
         self.ttl = ttl
 
     def check(self, mc, key):
-        key = self.prefix + key
+        key = _tostr(self.prefix + key)
         result = mc.get(key)
         return result is None
 
     def incr(self, mc, key):
-        key = self.prefix + key
+        key = _tostr(self.prefix + key)
         if not self.rl.check(mc, key):
             mc.set(key, 'true', time=self.ttl)
 
@@ -155,7 +161,7 @@ def blacklist_on_auth_failure(key_fn, count=0, period=0, ttl=0, check_wl=False):
     def decoratorfn(fn):
         @functools.wraps(fn)
         def _blacklist(*args, **kwargs):
-            key = key_fn(fn.__name__, args)
+            key = _tostr(key_fn(fn.__name__, args))
             if not key:
                 return fn(*args, **kwargs)