Add Request.Verify.

sdg0919 · Apr 26, 2016 · e265fca · e265fca
1 parent 1eb25f2
commit e265fca
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 1 deletion.
diff --git a/downloader/download.go b/downloader/download.go
@@ -22,6 +22,10 @@ type Request struct {
 	// TargetDir is the directory into which the file will be downloaded.
 	// It defaults to os.TempDir().
 	TargetDir string
+
+	// Verify is used to ensure that the download result is correct. If
+	// no func is provided then no verification happens.
+	Verify func(*os.File) error
 }
 
 // Status represents the status of a completed download.
@@ -99,6 +103,13 @@ func (dl *Download) run(req Request) {
 		err = errors.Errorf("cannot download %q: %v", req.URL, err)
 	}
 
+	if req.Verify != nil {
+		err = req.Verify(file)
+		if _, err2 := file.Seek(0, os.SEEK_SET); err2 != nil && err == nil {
+			err = err2
+		}
+	}
+
 	status := Status{
 		File: file,
 		Err:  err,

diff --git a/downloader/download_test.go b/downloader/download_test.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/juju/errors"
 	gitjujutesting "github.com/juju/testing"
 	jc "github.com/juju/testing/checkers"
 	"github.com/juju/utils"
@@ -96,7 +97,7 @@ func (s *DownloadSuite) TestDownloadError(c *gc.C) {
 	c.Assert(status.Err, gc.ErrorMatches, `cannot download ".*": bad http response: 404 Not Found`)
 }
 
-func (s *DownloadSuite) TestStopDownload(c *gc.C) {
+func (s *DownloadSuite) TestStop(c *gc.C) {
 	tmp := c.MkDir()
 	d := downloader.StartDownload(
 		downloader.Request{
@@ -116,6 +117,51 @@ func (s *DownloadSuite) TestStopDownload(c *gc.C) {
 	c.Assert(infos, gc.HasLen, 0)
 }
 
+func (s *DownloadSuite) TestVerifyValid(c *gc.C) {
+	stub := &gitjujutesting.Stub{}
+	tmp := c.MkDir()
+	gitjujutesting.Server.Response(200, nil, []byte("archive"))
+	dl := downloader.StartDownload(
+		downloader.Request{
+			URL:       s.URL(c, "/archive.tgz"),
+			TargetDir: tmp,
+			Verify: func(f *os.File) error {
+				stub.AddCall("Verify", f)
+				return nil
+			},
+		},
+		downloader.NewHTTPBlobOpener(utils.VerifySSLHostnames),
+	)
+	status := <-dl.Done()
+	c.Assert(status.Err, jc.ErrorIsNil)
+
+	stub.CheckCallNames(c, "Verify")
+	stub.CheckCall(c, 0, "Verify", status.File)
+}
+
+func (s *DownloadSuite) TestVerifyInvalid(c *gc.C) {
+	stub := &gitjujutesting.Stub{}
+	tmp := c.MkDir()
+	gitjujutesting.Server.Response(200, nil, []byte("archive"))
+	invalid := errors.NotValidf("oops")
+	dl := downloader.StartDownload(
+		downloader.Request{
+			URL:       s.URL(c, "/archive.tgz"),
+			TargetDir: tmp,
+			Verify: func(f *os.File) error {
+				stub.AddCall("Verify", f)
+				return invalid
+			},
+		},
+		downloader.NewHTTPBlobOpener(utils.VerifySSLHostnames),
+	)
+	status := <-dl.Done()
+
+	c.Check(errors.Cause(status.Err), gc.Equals, invalid)
+	stub.CheckCallNames(c, "Verify")
+	stub.CheckCall(c, 0, "Verify", status.File)
+}
+
 func assertFileContents(c *gc.C, f *os.File, expect string) {
 	got, err := ioutil.ReadAll(f)
 	c.Assert(err, jc.ErrorIsNil)

diff --git a/downloader/utils.go b/downloader/utils.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"os"
 
 	"github.com/juju/errors"
 	"github.com/juju/utils"
@@ -30,3 +31,19 @@ func NewHTTPBlobOpener(hostnameVerification utils.SSLHostnameVerification) func(
 		return resp.Body, nil
 	}
 }
+
+// NewSha256Verifier returns a verifier suitable for Request. The
+// verifier checks the SHA-256 checksum of the file to ensure that it
+// matches the provided one.
+func NewSha256Verifier(expected string) func(*os.File) error {
+	return func(file *os.File) error {
+		actual, _, err := utils.ReadSHA256(file)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if actual != expected {
+			return errors.Errorf("expected sha256 %q, got %q", expected, actual)
+		}
+		return nil
+	}
+}