From 486bef21b88383511015c56cb3962cea1a72cab6 Mon Sep 17 00:00:00 2001
From: dengkaipeng <dengkaipeng@baidu.com>
Date: Fri, 28 Aug 2020 07:38:24 +0000
Subject: [PATCH] check md5 after weights download

---
 ppdet/utils/download.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/ppdet/utils/download.py b/ppdet/utils/download.py
index a7a999ebc..5697e14b6 100644
--- a/ppdet/utils/download.py
+++ b/ppdet/utils/download.py
@@ -22,6 +22,8 @@ import shutil
 import requests
 import tqdm
 import hashlib
+import binascii
+import base64
 import tarfile
 import zipfile
 
@@ -291,9 +293,19 @@ def _download(url, path, md5sum=None):
                 for chunk in req.iter_content(chunk_size=1024):
                     if chunk:
                         f.write(chunk)
-        shutil.move(tmp_fullname, fullname)
 
-    return fullname
+        # check md5 after download in Content-MD5 in req.headers
+        content_md5 = req.headers.get('content-md5')
+        if not content_md5 or _md5check(
+                tmp_fullname,
+                binascii.hexlify(base64.b64decode(content_md5.strip('"')))):
+            shutil.move(tmp_fullname, fullname)
+            return fullname
+        else:
+            logger.warn(
+                "Download from url imcomplete, try downloading again...")
+            os.remove(tmp_fullname)
+            continue
 
 
 def _md5check(fullname, md5sum=None):
-- 
GitLab