bugfix

ahserver/auth_api.py

@@ -1,13 +1,16 @@
 import time
+import uuid
 from aiohttp_auth import auth
 from aiohttp_auth.auth.ticket_auth import TktAuthentication
+from aiohttp_session.redis_storage import RedisStorage
 from os import urandom
 from aiohttp import web
 import aiohttp_session
 import aioredis
 import base64
+import binascii
 
-from aiohttp_session import get_session, session_middleware
+from aiohttp_session import get_session, session_middleware, Session
 from aiohttp_session.cookie_storage import EncryptedCookieStorage
 from aiohttp_session.redis_storage import RedisStorage
 
@@ -22,6 +25,35 @@ def get_client_ip(obj, request):
 	request['client_ip'] = ip
 	return ip
 
+class MyRedisStorage(RedisStorage):
+	def key_gen(self, request):
+		uuid = request.headers.get('client_uuid')
+		if not uuid:
+			return uuid.uuid4().hex
+		b = uuid.encode('utf-8')
+		return binascii.hexlify(b)
+		
+	async def save_session(self, request: web.Request, 
+				response: web.StreamResponse, 
+				session: Session) -> None:
+		key = session.identity
+		if key is None:
+			key = self.key_gen(request)
+			self.save_cookie(response, key, max_age=session.max_age)
+		else:
+			if session.empty:
+				self.save_cookie(response, "", max_age=session.max_age)
+			else:
+				key = str(key)
+				self.save_cookie(response, key, max_age=session.max_age)
+
+		data_str = self._encoder(self._get_session_data(session))
+		await self._redis.set(
+			self.cookie_name + "_" + key,
+			data_str,
+			ex=session.max_age,
+		)
+
 class AuthAPI(AppLogger):
 	def __init__(self):
 		super().__init__()
@@ -44,13 +76,15 @@ class AuthAPI(AppLogger):
 
 	async def setupAuth(self,app):
 		# setup session middleware in aiohttp fashion
-			
+		secret = b'iqwertyuiopasdfghjklzxcvbnm12345'
-		storage = EncryptedCookieStorage(urandom(32))
+		if self.conf.website.cookie_secret:
+			secret = self.conf.website.cookie_secret.encode('utf-8')
+		storage = EncryptedCookieStorage(secret)
 		if self.conf.website.session_redis:
 			url = self.conf.website.session_redis.url
 			# redis = await aioredis.from_url("redis://127.0.0.1:6379")
 			redis = await aioredis.from_url(url)
-			storage = aiohttp_session.redis_storage.RedisStorage(redis)
+			storage = MyRedisStorage(redis)
 		aiohttp_session.setup(app, storage)
 
 		# Create an auth ticket mechanism that expires after 1 minute (60
@@ -66,22 +100,23 @@ class AuthAPI(AppLogger):
 		def _new_ticket(self, request, user_id):
 			client_uuid = request.headers.get('client_uuid')
 			ip = self._get_ip(request)
-			if not ip:
-				ip = request.remote
 			valid_until = int(time.time()) + self._max_age
 			print(f'hack: my _new_ticket() called ...remote {ip=}, {client_uuid=}')
-			return self._ticket.new(user_id, valid_until=valid_until, client_ip=ip, user_data=client_uuid)
+			return self._ticket.new(user_id, 
+							valid_until=valid_until, 
+							client_ip=ip, 
+							user_data=client_uuid)
 
 		TktAuthentication._get_ip = get_client_ip
 		TktAuthentication._new_ticket = _new_ticket
-		policy = auth.SessionTktAuthentication(urandom(32), session_max_time,
+		policy = auth.SessionTktAuthentication(secret, 
+										session_max_time,
 										reissue_time=session_reissue_time,
 									   include_ip=True)
 
 		# setup aiohttp_auth.auth middleware in aiohttp fashion
 		# print('policy = ', policy)
 		auth.setup(app, policy)
-		print('add auth middleware ....................')
 		app.middlewares.append(self.checkAuth)
 
 	@web.middleware
@@ -100,6 +135,7 @@ class AuthAPI(AppLogger):
 				self.info(f'timecost=client({ip}) {user} access {path} cost {t3-t1}, ({t2-t1})')
 				return ret
 			except Exception as e:
+				t3 = time.time()
 				self.info(f'timecost=client({ip}) {user} access {path} cost {t3-t1}, ({t2-t1}), except={e}')
 				raise e