fix pickle unserialize error

Author: zaxtyson

.gitignore

@@ -0,0 +1,9 @@
+# Byte-compiled / optimized
+__pycache__/
+*.py[cod]
+
+# Pycharm
+.idea/
+
+# Environment
+venv/

README.md

@@ -34,6 +34,10 @@
 
 # 更新日志
 
+## `v2.6.6`
+
+- 修复 pickle 库反序列化导致内存溢出的问题[#65](https://github.com/zaxtyson/LanZouCloud-API/issues/65)
+
 ## `v2.6.5`
 
 - 修复蓝奏云主域名解析异常的问题[#59](https://github.com/zaxtyson/LanZouCloud-API/issues/59) [#60](https://github.com/zaxtyson/LanZouCloud-API/pull/60)

lanzou/api/core.py

@@ -832,7 +832,7 @@ def _upload_big_file(self, file_path, dir_id, *, callback=None, uploaded_handler
         if not os.path.exists(record_file):  # 初始化记录文件
             info = {'name': file_name, 'size': file_size, 'uploaded': 0, 'parts': []}
             with open(record_file, 'wb') as f:
-                pickle.dump(info, f)
+                pickle.dump(info, f, protocol=4)
         else:
             with open(record_file, 'rb') as f:
                 info = pickle.load(f)
@@ -859,7 +859,7 @@ def _close_pwd(fid, is_file):  # 数据块上传后默认关闭提取码
                 info['parts'].append(os.path.basename(data_path))  # 记录已上传的文件名
                 with open(record_file, 'wb') as f:
                     logger.debug(f"Update record file: {uploaded_size}/{file_size}")
-                    pickle.dump(info, f)
+                    pickle.dump(info, f, protocol=4)
             else:
                 logger.debug(f"Upload data file failed: data_path={data_path}")
                 return LanZouCloud.FAILED
@@ -1041,10 +1041,19 @@ def down_file_by_url(self, share_url, pwd='', save_path='./Download', *, callbac
         os.rename(tmp_file_path, file_path)  # 下载完成，改回正常文件名
         if os.path.getsize(file_path) > 512:  # 文件大于 512 bytes 就检查一下
             file_info = None
+            is_protocol_3 = False
             with open(file_path, 'rb') as f:
                 f.seek(-512, os.SEEK_END)
                 last_512_bytes = f.read()
                 file_info = un_serialize(last_512_bytes)
+                # Python3.6 序列化时默认使用 pickle 第三版协议, 
+                # 导致计算时文件尾部多写了 5 字节， 保险期起见处理一下 
+                if not file_info: 
+                    is_protocol_3 = True
+                    f.seek(-517, os.SEEK_END)
+                    last_517_bytes = f.read()
+                    file_info = un_serialize(last_517_bytes)
+                    
 
             # 大文件的记录文件也可以反序列化出 name,但是没有 padding 字段
             if file_info is not None and 'padding' in file_info:
@@ -1059,7 +1068,8 @@ def down_file_by_url(self, share_url, pwd='', save_path='./Download', *, callbac
                 os.rename(file_path, new_file_path)
                 # 截断最后 512 字节隐藏信息, 还原文件
                 with open(new_file_path, 'rb+') as f:
-                    f.seek(-512, os.SEEK_END)
+                    truncate_size = 517 if is_protocol_3 else 512
+                    f.seek(-truncate_size, os.SEEK_END)
                     f.truncate()
                 file_path = new_file_path  # 保存文件重命名后真实路径
 
@@ -1222,7 +1232,7 @@ def _down_big_file(self, name, total_size, file_list, save_path, *, callback=Non
         if not os.path.exists(record_file):  # 初始化记录文件
             info = {'last_ending': 0, 'finished': []}  # 记录上一个数据块结尾地址和已经下载的数据块
             with open(record_file, 'wb') as rf:
-                pickle.dump(info, rf)
+                pickle.dump(info, rf, protocol=4)
         else:  # 读取记录文件，下载续传
             with open(record_file, 'rb') as rf:
                 info = pickle.load(rf)
@@ -1266,7 +1276,7 @@ def _down_big_file(self, name, total_size, file_list, save_path, *, callback=Non
                 finally:
                     info['last_ending'] = file_size_now
                     with open(record_file, 'wb') as rf:
-                        pickle.dump(info, rf)
+                        pickle.dump(info, rf, protocol=4)
                     logger.debug(f"Update download record info: {info}")
             # 全部数据块下载完成, 记录文件可以删除
             logger.debug(f"Delete download record file: {record_file}")

lanzou/api/utils.py

@@ -112,6 +112,15 @@ def is_folder_url(share_url: str) -> bool:
 
 def un_serialize(data: bytes):
     """反序列化文件信息数据"""
+    # https://github.com/zaxtyson/LanZouCloud-API/issues/65
+    is_right_format = False
+    if data.startswith(b"\x80\x04") and data.endswith(b"u."):
+        is_right_format = True
+    if data.startswith(b"\x80\x03") and data.endswith(b"u."):
+        is_right_format = True
+
+    if not is_right_format:
+        return None
     try:
         ret = pickle.loads(data)
         if not isinstance(ret, dict):
@@ -189,9 +198,10 @@ def let_me_upload(file_path):
                 chunk = in_f.read(4096)
         # 构建文件 "报尾" 保存真实文件名,大小 512 字节
         # 追加数据到文件尾部，并不会影响文件的使用，无需修改即可分享给其他人使用，自己下载时则会去除，确保数据无误
-        padding = 512 - len(file_name.encode('utf-8')) - 42  # 序列化后空字典占 42 字节
+        # protocol=4（py3.8默认), 序列化后空字典占 42 字节
+        padding = 512 - len(file_name.encode('utf-8')) - 42
         data = {'name': file_name, 'padding': b'\x00' * padding}
-        data = pickle.dumps(data)
+        data = pickle.dumps(data, protocol=4)
         out_f.write(data)
     return new_file_path