DarthSim пре 8 година
родитељ
комит
69b607cd98
1 измењених фајлова са 37 додато и 17 уклоњено
  1. 37 17
      download.go

+ 37 - 17
download.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"bytes"
 	"errors"
 	"errors"
 	"image"
 	"image"
+	"io"
 	"net/http"
 	"net/http"
 
 
 	_ "image/gif"
 	_ "image/gif"
@@ -11,10 +12,39 @@ import (
 	_ "image/png"
 	_ "image/png"
 )
 )
 
 
-const chunkSize = 4096
+type netReader struct {
+	reader io.Reader
+	buf    *bytes.Buffer
+}
+
+func newNetReader(r io.Reader) *netReader {
+	return &netReader{
+		reader: r,
+		buf:    bytes.NewBuffer([]byte{}),
+	}
+}
+
+func (r *netReader) Read(p []byte) (n int, err error) {
+	n, err = r.reader.Read(p)
+	if err == nil {
+		r.buf.Write(p[:n])
+	}
+	return
+}
+
+func (r *netReader) ReadAll() ([]byte, error) {
+	if _, err := r.buf.ReadFrom(r.reader); err != nil {
+		return []byte{}, err
+	}
+	return r.buf.Bytes(), nil
+}
 
 
-func checkTypeAndDimensions(b []byte) error {
-	imgconf, _, err := image.DecodeConfig(bytes.NewReader(b))
+func (r *netReader) GrowBuf(s int) {
+	r.buf.Grow(s)
+}
+
+func checkTypeAndDimensions(r io.Reader) error {
+	imgconf, _, err := image.DecodeConfig(r)
 	if err != nil {
 	if err != nil {
 		return err
 		return err
 	}
 	}
@@ -25,27 +55,17 @@ func checkTypeAndDimensions(b []byte) error {
 }
 }
 
 
 func readAndCheckImage(res *http.Response) ([]byte, error) {
 func readAndCheckImage(res *http.Response) ([]byte, error) {
-	b := make([]byte, chunkSize)
-	n, err := res.Body.Read(b)
-	if err != nil {
-		return nil, err
-	}
+	nr := newNetReader(res.Body)
 
 
-	if err = checkTypeAndDimensions(b[:n]); err != nil {
+	if err := checkTypeAndDimensions(nr); err != nil {
 		return nil, err
 		return nil, err
 	}
 	}
 
 
-	buf := bytes.NewBuffer(b[:n])
-
 	if res.ContentLength > 0 {
 	if res.ContentLength > 0 {
-		buf.Grow(int(res.ContentLength))
-	}
-
-	if _, err = buf.ReadFrom(res.Body); err != nil {
-		return nil, err
+		nr.GrowBuf(int(res.ContentLength))
 	}
 	}
 
 
-	return buf.Bytes(), nil
+	return nr.ReadAll()
 }
 }
 
 
 func downloadImage(url string) ([]byte, error) {
 func downloadImage(url string) ([]byte, error) {