func (d *Downloader) getPublishedTarball() (io.ReadCloser, error) {
log.Debugf("Downloading Amazon EC2 Container Service Agent from %s", config.AgentRemoteTarball())
resp, err := d.getter.Get(config.AgentRemoteTarball())
if err != nil {
return nil, err
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("unexpected response code %d", resp.StatusCode)
}
return resp.Body, nil
}
func TestDownloadAgentMD5Mismatch(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
md5response := &http.Response{}
md5sum := "md5sum"
tarballResponse := &http.Response{
StatusCode: 200,
Body: ioutil.NopCloser(&bytes.Buffer{}),
}
mockFS := NewMockfileSystem(mockCtrl)
mockFS.EXPECT().MkdirAll(config.CacheDirectory(), os.ModeDir|0700)
mockGetter := NewMockhttpGetter(mockCtrl)
mockGetter.EXPECT().Get(config.AgentRemoteTarballMD5()).Return(md5response, nil)
mockFS.EXPECT().ReadAll(md5response.Body).Return([]byte(md5sum), nil)
mockGetter.EXPECT().Get(config.AgentRemoteTarball()).Return(tarballResponse, nil)
tempfile, err := ioutil.TempFile("", "test")
if err != nil {
t.Fail()
}
defer tempfile.Close()
mockFS.EXPECT().TempFile("", "ecs-agent.tar").Return(tempfile, nil)
mockFS.EXPECT().TeeReader(tarballResponse.Body, gomock.Any())
mockFS.EXPECT().Copy(tempfile, gomock.Any()).Return(int64(0), nil)
mockFS.EXPECT().Remove(tempfile.Name())
d := &Downloader{
getter: mockGetter,
fs: mockFS,
}
d.DownloadAgent()
}
func TestDownloadAgentTempFile(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
md5response := &http.Response{}
md5sum := "md5sum"
tarballResponse := &http.Response{
StatusCode: 200,
Body: ioutil.NopCloser(&bytes.Buffer{}),
}
mockFS := NewMockfileSystem(mockCtrl)
mockFS.EXPECT().MkdirAll(config.CacheDirectory(), os.ModeDir|0700)
mockGetter := NewMockhttpGetter(mockCtrl)
mockGetter.EXPECT().Get(config.AgentRemoteTarballMD5()).Return(md5response, nil)
mockFS.EXPECT().ReadAll(md5response.Body).Return([]byte(md5sum), nil)
mockGetter.EXPECT().Get(config.AgentRemoteTarball()).Return(tarballResponse, nil)
mockFS.EXPECT().TempFile("", "ecs-agent.tar").Return(nil, errors.New("test error"))
d := &Downloader{
getter: mockGetter,
fs: mockFS,
}
d.DownloadAgent()
}
// DownloadAgent downloads a fresh copy of the Agent and performs an
// integrity check on the downloaded image
func (d *Downloader) DownloadAgent() error {
err := d.fs.MkdirAll(config.CacheDirectory(), os.ModeDir|orwPerm)
if err != nil {
return err
}
publishedMd5Sum, err := d.getPublishedMd5Sum()
if err != nil {
return err
}
publishedTarballReader, err := d.getPublishedTarball()
if err != nil {
return err
}
defer publishedTarballReader.Close()
md5hash := md5.New()
tempFile, err := d.fs.TempFile("", "ecs-agent.tar")
if err != nil {
return err
}
log.Debugf("Temp file %s", tempFile.Name())
defer func() {
if err != nil {
log.Debugf("Removing temp file %s", tempFile.Name())
d.fs.Remove(tempFile.Name())
}
}()
defer tempFile.Close()
teeReader := d.fs.TeeReader(publishedTarballReader, md5hash)
_, err = d.fs.Copy(tempFile, teeReader)
if err != nil {
return err
}
calculatedMd5Sum := md5hash.Sum(nil)
calculatedMd5SumString := fmt.Sprintf("%x", calculatedMd5Sum)
log.Debugf("Expected %s", publishedMd5Sum)
log.Debugf("Calculated %s", calculatedMd5SumString)
if publishedMd5Sum != calculatedMd5SumString {
err = fmt.Errorf("mismatched md5sum while downloading %s", config.AgentRemoteTarball())
return err
}
log.Debugf("Attempting to rename %s to %s", tempFile.Name(), config.AgentTarball())
return d.fs.Rename(tempFile.Name(), config.AgentTarball())
}
func TestDownloadAgentDownloadTarballFailure(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
md5response := &http.Response{}
md5sum := "md5sum"
mockFS := NewMockfileSystem(mockCtrl)
mockFS.EXPECT().MkdirAll(config.CacheDirectory(), os.ModeDir|0700)
mockGetter := NewMockhttpGetter(mockCtrl)
mockGetter.EXPECT().Get(config.AgentRemoteTarballMD5()).Return(md5response, nil)
mockFS.EXPECT().ReadAll(md5response.Body).Return([]byte(md5sum), nil)
mockGetter.EXPECT().Get(config.AgentRemoteTarball()).Return(nil, errors.New("test error"))
d := &Downloader{
getter: mockGetter,
fs: mockFS,
}
d.DownloadAgent()
}
func TestDownloadAgentSuccess(t *testing.T) {
mockCtrl := gomock.NewController(t)
defer mockCtrl.Finish()
tarballContents := "tarball contents"
tarballResponse := &http.Response{
StatusCode: 200,
Body: ioutil.NopCloser(bytes.NewBufferString(tarballContents)),
}
expectedMd5Sum := fmt.Sprintf("%x\n", md5.Sum([]byte(tarballContents)))
md5response := &http.Response{}
mockFS := NewMockfileSystem(mockCtrl)
mockFS.EXPECT().MkdirAll(config.CacheDirectory(), os.ModeDir|0700)
mockGetter := NewMockhttpGetter(mockCtrl)
mockGetter.EXPECT().Get(config.AgentRemoteTarballMD5()).Return(md5response, nil)
mockFS.EXPECT().ReadAll(md5response.Body).Return([]byte(expectedMd5Sum), nil)
mockGetter.EXPECT().Get(config.AgentRemoteTarball()).Return(tarballResponse, nil)
tempfile, err := ioutil.TempFile("", "test")
if err != nil {
t.Fail()
}
defer tempfile.Close()
mockFS.EXPECT().TempFile("", "ecs-agent.tar").Return(tempfile, nil)
mockFS.EXPECT().TeeReader(tarballResponse.Body, gomock.Any()).Do(func(reader io.Reader, writer io.Writer) {
_, err = io.Copy(writer, reader)
if err != nil {
t.Fail()
}
})
mockFS.EXPECT().Copy(tempfile, gomock.Any()).Return(int64(0), nil)
mockFS.EXPECT().Rename(tempfile.Name(), config.AgentTarball())
d := &Downloader{
getter: mockGetter,
fs: mockFS,
}
d.DownloadAgent()
}