浏览代码

Grow download buffer to Content-Length value

DarthSim 6 年之前
父节点
当前提交
6fa0c539bd
共有 1 个文件被更改,包括 11 次插入4 次删除
  1. 11 4
      download.go

+ 11 - 4
download.go

@@ -103,10 +103,17 @@ func readAndCheckImage(ctx context.Context, res *http.Response) (context.Context
 		return ctx, cancel, err
 	}
 
-	if cls := res.Header.Get("Content-Length"); len(cls) > 0 {
-		if cl, err := strconv.Atoi(cls); err == nil && cl > buf.Len() && cl > buf.Cap() {
-			buf.Grow(cl - buf.Len())
-		}
+	var contentLength int
+
+	if res.ContentLength > 0 {
+		contentLength = int(res.ContentLength)
+	} else {
+		// ContentLength wasn't set properly, trying to parse the header
+		contentLength, _ = strconv.Atoi(res.Header.Get("Content-Length"))
+	}
+
+	if contentLength > buf.Cap() {
+		buf.Grow(contentLength - buf.Len())
 	}
 
 	if _, err = buf.ReadFrom(res.Body); err != nil {