[tor-commits] [stem/master] Don't use @lru_cache for exit policy rules

[atagar at torproject.org]

commit 863ddacb961ffe255f02eeea115515bd7f3ffd41
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Mar 17 09:19:32 2014 -0700

    Don't use @lru_cache for exit policy rules
    
    Lession learned today: @lru_cache is handy for efficiency but don't rely on the
    cache to function.
    
    Our ExitPolicy class expected its _get_rules() method to be called exactly once
    due to having cached results. However, pickling then unpicking an ExitPolicy
    class caused our @lru_cache to be cleared, resulting in a TypeError.
    
    Replacing this method's @lru_cache with a standard one, and adding a test for
    pickleability. Thanks to Nicholas Hopper for the catch!
---
 stem/exit_policy.py             |   88 ++++++++++++++++++++-------------------
 test/unit/exit_policy/policy.py |   16 +++++++
 2 files changed, 62 insertions(+), 42 deletions(-)

diff --git a/stem/exit_policy.py b/stem/exit_policy.py
index 169eb13..78db793 100644
--- a/stem/exit_policy.py
+++ b/stem/exit_policy.py
@@ -167,6 +167,8 @@ class ExitPolicy(object):
     else:
       self._input_rules = rules
 
+    self._rules = None
+
     # Result when no rules apply. According to the spec policies default to 'is
     # allowed', but our microdescriptor policy subclass might want to change
     # this.
@@ -295,54 +297,56 @@ class ExitPolicy(object):
 
     return (label_prefix + ", ".join(display_ranges)).strip()
 
-  @lru_cache()
   def _get_rules(self):
-    rules = []
-    is_all_accept, is_all_reject = True, True
-
-    if isinstance(self._input_rules, bytes):
-      decompressed_rules = zlib.decompress(self._input_rules).split(b',')
-    else:
-      decompressed_rules = self._input_rules
-
-    for rule in decompressed_rules:
-      if isinstance(rule, bytes):
-        rule = stem.util.str_tools._to_unicode(rule)
-
-      if isinstance(rule, unicode):
-        rule = ExitPolicyRule(rule.strip())
+    if self._rules is None:
+      rules = []
+      is_all_accept, is_all_reject = True, True
 
-      if rule.is_accept:
-        is_all_reject = False
+      if isinstance(self._input_rules, bytes):
+        decompressed_rules = zlib.decompress(self._input_rules).split(b',')
       else:
-        is_all_accept = False
+        decompressed_rules = self._input_rules
 
-      rules.append(rule)
+      for rule in decompressed_rules:
+        if isinstance(rule, bytes):
+          rule = stem.util.str_tools._to_unicode(rule)
 
-      if rule.is_address_wildcard() and rule.is_port_wildcard():
-        break  # this is a catch-all, no reason to include more
+        if isinstance(rule, unicode):
+          rule = ExitPolicyRule(rule.strip())
 
-    # If we only have one kind of entry *and* end with a wildcard then
-    # we might as well use the simpler version. For instance...
-    #
-    #   reject *:80, reject *:443, reject *:*
-    #
-    # ... could also be represented as simply...
-    #
-    #   reject *:*
-    #
-    # This mostly comes up with reject-all policies because the
-    # 'reject private:*' appends an extra seven rules that have no
-    # effect.
-
-    if rules and (rules[-1].is_address_wildcard() and rules[-1].is_port_wildcard()):
-      if is_all_accept:
-        rules = [ExitPolicyRule("accept *:*")]
-      elif is_all_reject:
-        rules = [ExitPolicyRule("reject *:*")]
-
-    self._input_rules = None
-    return rules
+        if rule.is_accept:
+          is_all_reject = False
+        else:
+          is_all_accept = False
+
+        rules.append(rule)
+
+        if rule.is_address_wildcard() and rule.is_port_wildcard():
+          break  # this is a catch-all, no reason to include more
+
+      # If we only have one kind of entry *and* end with a wildcard then
+      # we might as well use the simpler version. For instance...
+      #
+      #   reject *:80, reject *:443, reject *:*
+      #
+      # ... could also be represented as simply...
+      #
+      #   reject *:*
+      #
+      # This mostly comes up with reject-all policies because the
+      # 'reject private:*' appends an extra seven rules that have no
+      # effect.
+
+      if rules and (rules[-1].is_address_wildcard() and rules[-1].is_port_wildcard()):
+        if is_all_accept:
+          rules = [ExitPolicyRule("accept *:*")]
+        elif is_all_reject:
+          rules = [ExitPolicyRule("reject *:*")]
+
+      self._rules = rules
+      self._input_rules = None
+
+    return self._rules
 
   def __iter__(self):
     for rule in self._get_rules():
diff --git a/test/unit/exit_policy/policy.py b/test/unit/exit_policy/policy.py
index 161bae0..041f81e 100644
--- a/test/unit/exit_policy/policy.py
+++ b/test/unit/exit_policy/policy.py
@@ -2,6 +2,7 @@
 Unit tests for the stem.exit_policy.ExitPolicy class.
 """
 
+import pickle
 import unittest
 
 from stem.exit_policy import get_config_policy, \
@@ -216,3 +217,18 @@ class TestExitPolicy(unittest.TestCase):
 
     for test_input in test_inputs:
       self.assertRaises(ValueError, get_config_policy, test_input)
+
+  def test_pickleability(self):
+    """
+    Checks that we can unpickle ExitPolicy instances.
+    """
+
+    policy = ExitPolicy("accept *:80", "accept *:443", "reject *:*")
+    self.assertTrue(policy.can_exit_to('74.125.28.106', 80))
+
+    encoded_policy = pickle.dumps(policy)
+    restored_policy = pickle.loads(encoded_policy)
+
+    self.assertEqual(policy, restored_policy)
+    self.assertTrue(restored_policy.is_exiting_allowed())
+    self.assertTrue(restored_policy.can_exit_to('74.125.28.106', 80))