Checksum verification while downloading files.

utils/download.go

@@ -14,6 +14,7 @@ import (
 // Downloader is parallel HTTP fetcher
 type Downloader interface {
 	Download(url string, destination string, result chan<- error)
+	DownloadWithChecksum(url string, destination string, result chan<- error, expected ChecksumInfo)
 	Pause()
 	Resume()
 	Shutdown()
@@ -39,6 +40,7 @@ type downloadTask struct {
 	url         string
 	destination string
 	result      chan<- error
+	expected    ChecksumInfo
 }
 
 // NewDownloader creates new instance of Downloader which specified number
@@ -88,7 +90,12 @@ func (downloader *downloaderImpl) Resume() {
 
 // Download starts new download task
 func (downloader *downloaderImpl) Download(url string, destination string, result chan<- error) {
-	downloader.queue <- &downloadTask{url: url, destination: destination, result: result}
+	downloader.DownloadWithChecksum(url, destination, result, ChecksumInfo{Size: -1})
+}
+
+// DownloadWithChecksum starts new download task with checksum verification
+func (downloader *downloaderImpl) DownloadWithChecksum(url string, destination string, result chan<- error, expected ChecksumInfo) {
+	downloader.queue <- &downloadTask{url: url, destination: destination, result: result, expected: expected}
 }
 
 // handleTask processes single download task
@@ -122,13 +129,43 @@ func (downloader *downloaderImpl) handleTask(task *downloadTask) {
 	}
 	defer outfile.Close()
 
-	_, err = io.Copy(outfile, resp.Body)
+	var w io.Writer
+
+	checksummer := NewChecksumWriter()
+
+	if task.expected.Size != -1 {
+		w = io.MultiWriter(outfile, checksummer)
+	} else {
+		w = outfile
+	}
+
+	_, err = io.Copy(w, resp.Body)
 	if err != nil {
 		os.Remove(temppath)
 		task.result <- err
 		return
 	}
 
+	if task.expected.Size != -1 {
+		actual := checksummer.Sum()
+
+		if actual.Size != task.expected.Size {
+			err = fmt.Errorf("%s: size check mismatch %d != %d", task.url, actual.Size, task.expected.Size)
+		} else if task.expected.MD5 != "" && actual.MD5 != task.expected.MD5 {
+			err = fmt.Errorf("%s: md5 hash mismatch %#v != %#v", task.url, actual.MD5, task.expected.MD5)
+		} else if task.expected.SHA1 != "" && actual.SHA1 != task.expected.SHA1 {
+			err = fmt.Errorf("%s: sha1 hash mismatch %#v != %#v", task.url, actual.SHA1, task.expected.SHA1)
+		} else if task.expected.SHA256 != "" && actual.SHA256 != task.expected.SHA256 {
+			err = fmt.Errorf("%s: sha256 hash mismatch %#v != %#v", task.url, actual.SHA256, task.expected.SHA256)
+		}
+
+		if err != nil {
+			os.Remove(temppath)
+			task.result <- err
+			return
+		}
+	}
+
 	err = os.Rename(temppath, task.destination)
 	if err != nil {
 		os.Remove(temppath)

utils/download_test.go

@@ -86,6 +86,43 @@ func (s *DownloaderSuite) TestDownloadOK(c *C) {
 	c.Assert(res, IsNil)
 }
 
+func (s *DownloaderSuite) TestDownloadWithChecksum(c *C) {
+	d := NewDownloader(2)
+	defer d.Shutdown()
+	ch := make(chan error)
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{})
+	res := <-ch
+	c.Assert(res, ErrorMatches, ".*size check mismatch 12 != 0")
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{Size: 12, MD5: "abcdef"})
+	res = <-ch
+	c.Assert(res, ErrorMatches, ".*md5 hash mismatch \"a1acb0fe91c7db45ec4d775192ec5738\" != \"abcdef\"")
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738"})
+	res = <-ch
+	c.Assert(res, IsNil)
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738", SHA1: "abcdef"})
+	res = <-ch
+	c.Assert(res, ErrorMatches, ".*sha1 hash mismatch \"921893bae6ad6fd818401875d6779254ef0ff0ec\" != \"abcdef\"")
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
+		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec"})
+	res = <-ch
+	c.Assert(res, IsNil)
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
+		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec", SHA256: "abcdef"})
+	res = <-ch
+	c.Assert(res, ErrorMatches, ".*sha256 hash mismatch \"b3c92ee1246176ed35f6e8463cd49074f29442f5bbffc3f8591cde1dcc849dac\" != \"abcdef\"")
+
+	d.DownloadWithChecksum(s.url+"/test", s.tempfile.Name(), ch, ChecksumInfo{Size: 12, MD5: "a1acb0fe91c7db45ec4d775192ec5738",
+		SHA1: "921893bae6ad6fd818401875d6779254ef0ff0ec", SHA256: "b3c92ee1246176ed35f6e8463cd49074f29442f5bbffc3f8591cde1dcc849dac"})
+	res = <-ch
+	c.Assert(res, IsNil)
+}
+
 func (s *DownloaderSuite) TestDownload404(c *C) {
 	d := NewDownloader(2)
 	defer d.Shutdown()

utils/fake.go

@@ -2,6 +2,7 @@ package utils
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
 )
@@ -47,18 +48,18 @@ func (f *FakeDownloader) Empty() bool {
 	return len(f.expected) == 0
 }
 
-// Download performs fake download by matching against first expectation in the queue
+// DownloadWithChecksum performs fake download by matching against first expectation in the queue, with cheksum verification
-func (f *FakeDownloader) Download(url string, filename string, result chan<- error) {
+func (f *FakeDownloader) DownloadWithChecksum(url string, filename string, result chan<- error, expected ChecksumInfo) {
 	if len(f.expected) == 0 || f.expected[0].URL != url {
 		result <- fmt.Errorf("unexpected request for %s", url)
 		return
 	}
 
-	expected := f.expected[0]
+	expectation := f.expected[0]
 	f.expected = f.expected[1:]
 
-	if expected.Err != nil {
+	if expectation.Err != nil {
-		result <- expected.Err
+		result <- expectation.Err
 		return
 	}
 
@@ -75,16 +76,31 @@ func (f *FakeDownloader) Download(url string, filename string, result chan<- err
 	}
 	defer outfile.Close()
 
-	_, err = outfile.Write([]byte(expected.Response))
+	cks := NewChecksumWriter()
+	w := io.MultiWriter(outfile, cks)
+
+	_, err = w.Write([]byte(expectation.Response))
 	if err != nil {
 		result <- err
 		return
 	}
 
+	if expected.MD5 != "" {
+		if expected != cks.Sum() {
+			result <- fmt.Errorf("checksums don't match: %#v != %#v", expected, cks.Sum())
+			return
+		}
+	}
+
 	result <- nil
 	return
 }
 
+// Download performs fake download by matching against first expectation in the queue
+func (f *FakeDownloader) Download(url string, filename string, result chan<- error) {
+	f.DownloadWithChecksum(url, filename, result, ChecksumInfo{})
+}
+
 // Shutdown does nothing
 func (f *FakeDownloader) Shutdown() {
 }