Merge pull request #81 from ElijahCapricorn/master

Antizapret fails when several threads update the config

resources/proxy/antizapret.py

@@ -1,69 +1,72 @@
 # -*- coding: utf-8 -*-
 
-import os, re, fnmatch, threading, urllib2
+import os, re, fnmatch, threading, urllib2, time, shelve, anydbm
 from contextlib import contextmanager, closing
 from functions import log, debug, tempdir
 
-LOCKS = {}
 PAC_URL = "http://antizapret.prostovpn.org/proxy.pac"
 CACHE_DIR = tempdir()
 USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_8_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/34.0.1847.116 Safari/537.36"
+CONFIG_LOCK = threading.Lock()
 
 if not os.path.exists(CACHE_DIR):
     os.makedirs(CACHE_DIR)
 
-CACHE = 24 * 3600 # 24 hour caching
-
-#@contextmanager
-def shelf(filename, ttl=0):
-    import shelve
-    filename = os.path.join(CACHE_DIR, filename)
-    with LOCKS.get(filename, threading.RLock()):
-#        with closing(shelve.open(filename, writeback=True)) as d:
-        d = shelve.open(filename, writeback=True)
-        #try:
-        if True:
-            import time
-            if not dict(d):
-                log('xxx')
-                d.update({
-                    "created_at": time.time(),
-                    "data": {},
-                })
-            elif ttl > 0 and (time.time() - d["created_at"]) > ttl:
-                log('xxx')
-                d["created_at"] = time.time()
-                d["data"] = {}
-            return d
-        #except:
-        #    d.close()
-        #    raise
-
-_config = {}
+CACHE_LIFETIME = 24 * 3600 # 24 hour caching
 
 def config():
-    global _config
-    if not _config:
-#        with shelf("antizapret.pac_config", ttl=CACHE) as pac_config:
-        d = shelf("antizapret.pac_config2", ttl=CACHE)
-        pac_config = d['data']
-        if not pac_config:
-            log("[antizapret]: Fetching Antizapret PAC file on %s" %PAC_URL)
-            try:
-                pac_data = urllib2.urlopen(PAC_URL).read()
-            except:
-                pac_data = ""
+    shelf = None
+    try:
+        CONFIG_LOCK.acquire()
+        filename = os.path.join(CACHE_DIR, "antizapret.pac_config2")
+        try:
+            shelf = shelve.open(filename)
+        except anydbm.error:
+            os.remove(filename)
+            shelf = shelve.open(filename)
 
-            r = re.search(r"\"PROXY (.*); DIRECT", pac_data)
-            if r:
-                pac_config["server"] = r.group(1)
-                pac_config["domains"] = map(lambda x: x.replace(r"\Z(?ms)", "").replace("\\", ""), map(fnmatch.translate, re.findall(r"\"(.*?)\",", pac_data)))
-            else:
-                pac_config["server"] = None
-                pac_config["domains"] = []
-        d.close()
-        _config = pac_config
-    return _config
+        created_at = 0
+        data = {}
+
+        if 'created_at' in shelf:
+            created_at = shelf['created_at']
+
+        if 'data' in shelf:
+            data = shelf['data']
+
+        if((time.time() - created_at) <= CACHE_LIFETIME
+                and 'domains' in data
+                and len(data['domains']) > 0):
+            return data
+
+        log("[antizapret]: Fetching Antizapret PAC file on %s" %PAC_URL)
+        try:
+            pac_data = urllib2.urlopen(PAC_URL).read()
+        except:
+            pac_data = ""
+
+        r = re.search(r"\"PROXY (.*); DIRECT", pac_data)
+        if r:
+            data["server"] = r.group(1)
+            data["domains"] = map(lambda x: x.replace(r"\Z(?ms)", "").replace("\\", ""), map(fnmatch.translate, re.findall(r"\"(.*?)\",", pac_data)))
+        else:
+            data["server"] = None
+            data["domains"] = []
+
+        shelf.clear()
+        shelf.update({
+            "created_at": time.time(),
+            "data": data,
+        })
+        return data
+    except Exception as ex:
+        debug("[antizapret]: " + str(ex))
+        raise
+    finally:
+        if shelf:
+            shelf.close()
+        if CONFIG_LOCK.locked():
+            CONFIG_LOCK.release()
 
 class AntizapretProxyHandler(urllib2.ProxyHandler, object):
     def __init__(self):
@@ -76,7 +79,8 @@ class AntizapretProxyHandler(urllib2.ProxyHandler, object):
     def proxy_open(self, req, proxy, type):
         import socket
 
-        if socket.gethostbyname(req.get_host().split(":")[0]) in self.config["domains"]:
+        hostname = req.get_host().split(":")[0]
+        if socket.gethostbyname(hostname) in self.config["domains"] or hostname in self.config["domains"]:
             debug("[antizapret]: Pass request through proxy " + self.config["server"])
             return urllib2.ProxyHandler.proxy_open(self, req, self.config["server"], type)
 
@@ -108,4 +112,3 @@ def url_get(url, params={}, headers={}, post = None):
     except urllib2.HTTPError as e:
         log("[antizapret]: HTTP Error(%s): %s" % (e.errno, e.strerror))
         return None
-