feat: add session token middleware (#253)

aisk · web-flow · commit fec71097fdfb · 2016-10-24T11:38:47.000+08:00
diff --git a/.travis.yml b/.travis.yml
@@ -18,7 +18,7 @@ matrix:
 sudo: false
 
 install:
-  - pip install -r dev-requirements.txt
+  - pip install -e .'[test]'
   - pip install codecov
 
 script:
diff --git a/leancloud/engine/cookie_session.py b/leancloud/engine/cookie_session.py
@@ -0,0 +1,74 @@
+# coding: utf-8
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+from werkzeug import http
+from werkzeug.wrappers import Request
+from werkzeug.contrib.securecookie import SecureCookie
+
+from leancloud.user import User
+
+
+__author__ = 'asaka <lan@leancloud.rocks>'
+
+
+class CookieSessionMiddleware(object):
+    def __init__(self, app, secret, name='leancloud:session', exluded_paths=None, fetch_user=False):
+        if not secret:
+            raise RuntimeError('secret is required')
+        self.fetch_user = fetch_user
+        self.secret = secret
+        self.app = app
+        self.name = name
+        self.exluded_paths = [
+            '/__engine/',
+            '/1/functions/',
+            '/1.1/functions/',
+            '/1/call/',
+            '/1.1/call/',
+        ]
+        if exluded_paths:
+            self.exluded_paths += exluded_paths
+
+    def __call__(self, environ, start_response):
+        self.pre_process(environ)
+        def new_start_response(status, response_headers):
+            self.post_process(response_headers)
+            return start_response(status, response_headers)
+        return self.app(environ, new_start_response)
+
+    def pre_process(self, environ):
+        request = Request(environ)
+        for prefix in self.exluded_paths:
+            if request.path.startswith(prefix):
+                return
+
+        cookie = request.cookies.get(self.name)
+        if not cookie:
+            return
+
+        session = SecureCookie.unserialize(cookie, self.secret)
+
+        if not self.fetch_user:
+            user = User()
+            user._session_token = session['session_token']
+            user.id = session['uid']
+            User.set_current(user)
+        else:
+            user = User.become(session['session_token'])
+            User.set_current(user)
+
+
+    def post_process(self, headers):
+        user = User.get_current()
+        if not user:
+            return
+        cookie = SecureCookie({
+            'uid': user.id,
+            'session_token': user.get_session_token(),
+        }, self.secret)
+        raw = http.dump_cookie(self.name, cookie.serialize())
+        headers.append((b'Set-Cookie', raw))
diff --git a/leancloud/engine/https_redirect_middleware.py b/leancloud/engine/https_redirect_middleware.py
@@ -1,5 +1,9 @@
 # coding: utf-8
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
 import os
 
 from werkzeug.wrappers import Request
diff --git a/leancloud/user.py b/leancloud/user.py
@@ -54,6 +54,10 @@ def create_followee_query(cls, user_id):
     def get_current(cls):
         return getattr(thread_locals, 'current_user', None)
 
+    @classmethod
+    def set_current(cls, user):
+        thread_locals.current_user = user
+
     @classmethod
     def become(cls, session_token):
         response = client.get('/users/me', params={'session_token': session_token})
@@ -95,7 +99,7 @@ def _cleanup_auth_data(self):
 
     def _handle_save_result(self, make_current=False):
         if make_current:
-            thread_locals.current_user = self
+            User.set_current(self)
         self._cleanup_auth_data()
         # self._sync_all_auth_data()
         self._attributes.pop('password', None)
@@ -142,7 +146,7 @@ def logout(self):
         if not self.is_current:
             return
         self._cleanup_auth_data()
-        thread_locals.current_user = None
+        del thread_locals.current_user
 
     @classmethod
     def login_with_mobile_phone(cls, phone_number, password):
diff --git a/setup.py b/setup.py
@@ -5,13 +5,15 @@
 
 here = path.abspath(path.dirname(__file__))
 
-extra_require = {
-        'dev': ['sphinx'],
-        'test': ['nose', 'coverage', 'wsgi_intercept'],
-    }
-
-if sys.version_info.major == 2:
-    extra_require['test'].append('typing')
+install_requires=[
+    'arrow',
+    'iso8601',
+    'qiniu',
+    'requests',
+    'werkzeug',
+]
+if sys.version_info < (3, 5, 0):
+    install_requires.append('typing')
 
 setup(
     name='leancloud-sdk',
@@ -41,13 +43,10 @@
 
     test_suite='nose.collector',
 
-    install_requires=[
-        'arrow',
-        'iso8601',
-        'qiniu',
-        'requests',
-        'werkzeug',
-    ],
+    install_requires=install_requires,
 
-    extras_require=extra_require
+    extras_require = {
+        'dev': ['sphinx'],
+        'test': ['nose', 'wsgi_intercept'],
+    }
 )
diff --git a/tests/test_engine.py b/tests/test_engine.py
@@ -8,9 +8,10 @@
 import time
 import json
 import requests
+import typing
 
-from wsgi_intercept import requests_intercept, add_wsgi_intercept
-
+from wsgi_intercept import requests_intercept
+from wsgi_intercept import add_wsgi_intercept
 
 import leancloud
 from leancloud import Engine
@@ -22,7 +23,7 @@
 
 __author__ = 'asaka <lan@leancloud.rocks>'
 
-env = None
+env = None  # type: typing.Dict[str, str]
 
 TEST_APP_ID = os.environ['APP_ID']
 TEST_APP_KEY = os.environ['APP_KEY']
diff --git a/tests/test_middlewares.py b/tests/test_middlewares.py
@@ -0,0 +1,65 @@
+# coding: utf-8
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import time
+
+import requests
+from wsgi_intercept import add_wsgi_intercept
+from wsgi_intercept import remove_wsgi_intercept
+from wsgi_intercept import requests_intercept
+
+from leancloud import user as user_module
+from leancloud.engine.cookie_session import CookieSessionMiddleware
+
+
+HOST, PORT = 'localhost', 80
+URL = 'http://{}:{}/'.format(HOST, PORT)
+
+FAKE_USER_DATA = {
+  'sessionToken': 'qmdj8pdidnmyzp0c7yqil91oc',
+  'updatedAt': '2015-07-14T02:31:50.100Z',
+  'phone': '18612340000',
+  'objectId': '55a47496e4b05001a7732c5f',
+  'username': 'fool',
+  'createdAt': '2015-07-14T02:31:50.100Z',
+  'emailVerified': False,
+  'mobilePhoneVerified': False,
+}
+
+
+def application(environ, start_response):
+    start_response('200 OK', [('Content-Type', 'text/plain')])
+    return [b'hello!']
+
+
+def setup():
+    requests_intercept.install()
+
+
+def teardown():
+    requests_intercept.uninstall()
+
+
+def test_cookie_session_middleware():
+    user = user_module.User()
+    user._update_data(FAKE_USER_DATA)
+    user_module.thread_locals.current_user = user
+
+    app = CookieSessionMiddleware(application, b'wtf!')
+    add_wsgi_intercept(HOST, PORT, lambda: app)
+
+    response = requests.get(URL)
+    assert response.cookies['leancloud:session']
+
+    del user_module.thread_locals.current_user
+
+    requests.get(URL, cookies=response.cookies)
+    current = user_module.User.get_current()
+    assert current.id == user.id
+    assert current.get_session_token() == user.get_session_token()
+    assert not current._attributes
+    remove_wsgi_intercept()
