// Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.

//

// Licensed under the Apache License, Version 2.0 (the "License"). You may

// not use this file except in compliance with the License. A copy of the

// License is located at

//

// http://aws.amazon.com/apache2.0/

//

// or in the "license" file accompanying this file. This file is distributed

// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either

// express or implied. See the License for the specific language governing

// permissions and limitations under the License.

// Package cache provides functionality for working with an on-disk cache of

// the ECS Agent image.

package cache

import (

"bufio"

"crypto/md5"

"fmt"

"io"

"os"

"strings"

"github.com/aws/amazon-ecs-init/ecs-init/config"

"github.com/aws/aws-sdk-go/aws/ec2metadata"

"github.com/aws/aws-sdk-go/aws/session"

log "github.com/cihub/seelog"

"github.com/pkg/errors"

)

const (

orwPerm = 0700

regionalBucketFormat = "%s-%s"

)

// CacheStatus represents the status of the on-disk cache for agent

// tarballs in the cache directory. This status may be used to

// determine what the appropriate actions are based on the

// availability of agent images and communicates advice from the

// cache's state file.

type CacheStatus uint8

const (

// StatusUncached indicates that there is not an already downloaded

// and cached agent that is suitable for loading.

StatusUncached CacheStatus = 0

// StatusCached indicates that there is an agent downloaded and

// cached. This should be taken to mean that the image is suitable

// for load if agent isn't already loaded.

StatusCached CacheStatus = 1

// StatusReloadNeeded indicates that the cached image should take

// precedence over the already loaded agent image. This may be

// specified by the packaging to cause ecs-init to load a package

// distributed cached image on package installation, upgrades, or

// downgrades.

StatusReloadNeeded CacheStatus = 2

)

// Downloader is responsible for cache operations relating to downloading the agent

type Downloader struct {

s3Downloader s3DownloaderAPI

fs fileSystem

metadata instanceMetadata

region string

}

// NewDownloader returns a Downloader with default dependencies

func NewDownloader() (*Downloader, error) {

downloader := &Downloader{

fs: &standardFS{},

}

if config.RunningInExternal() {

downloader.metadata = &blackholeInstanceMetadata{}

} else {

sessionInstance, err := session.NewSession()

if err != nil {

// metadata client is only used for retrieving the user's region.

// If it cannot be initialized, the region field is populated with the default value to prevent future

// calls to retrieve the region from metadata.

log.Warnf("Got error when initializing session for metadata client: %v. Use default region: %s",

err, config.DefaultRegionName)

downloader.region = config.DefaultRegionName

} else {

downloader.metadata = ec2metadata.New(sessionInstance)

}

}

s3Downloader := &s3Downloader{

bucketDownloaders: make([]*s3BucketDownloader, 0),

cacheDir: config.CacheDirectory(),

fs: downloader.fs,

}

partitionBucketRegion := downloader.getPartitionBucketRegion()

partitionBucket := config.AgentPartitionBucketName

partitionBucketDownloader, err := newS3BucketDownloader(partitionBucketRegion, partitionBucket)

if err != nil {

log.Warnf("Failed to initialize partition bucket downloader: %v", err)

} else {

s3Downloader.addBucketDownloader(partitionBucketDownloader)

}

region := downloader.getRegion()

regionalBucket := fmt.Sprintf(regionalBucketFormat, partitionBucket, region)

regionalBucketDownloader, err := newS3BucketDownloader(region, regionalBucket)

if err != nil {

log.Warnf("Failed to initialize regional bucket downloader: %v", err)

} else {

s3Downloader.addBucketDownloader(regionalBucketDownloader)

}

if len(s3Downloader.bucketDownloaders) == 0 {

log.Error("Failed to initialize s3 downloader for either partition bucket or regional bucket. Downloader initialization fails.")

return nil, errors.New("failed to initialize downloader")

}

downloader.s3Downloader = s3Downloader

return downloader, nil

}

// AgentCacheStatus inspects the on-disk cache and returns its

// status. See `CacheStatus` for possible cache statuses and

// scenarios.

func (d *Downloader) AgentCacheStatus() CacheStatus {

stateFile := config.CacheState()

// State file and tarball must be non-zero to report status on

uncached := !(d.fileNotEmpty(stateFile) && d.fileNotEmpty(config.AgentTarball()))

if uncached {

return StatusUncached

}

file, err := d.fs.Open(stateFile)

if err != nil {

return StatusUncached

}

var status CacheStatus

_, err = fmt.Fscanf(file, "%d", &status)

if err != nil {

return StatusUncached

}

return status

}

// IsAgentCached returns true if there is a cached copy of the Agent present

// and a cache state file is not empty (no validation is performed on the

// tarball or cache state file contents)

func (d *Downloader) IsAgentCached() bool {

switch d.AgentCacheStatus() {

case StatusUncached:

return false

}

return true

}

func (d *Downloader) fileNotEmpty(filename string) bool {

fileinfo, err := d.fs.Stat(filename)

if err != nil {

return false

}

return fileinfo.Size() > 0

}

// getRegion finds region from metadata and caches for the life of downloader

func (d *Downloader) getRegion() string {

if d.region != "" {

return d.region

}

defaultRegion := config.DefaultRegionName

if config.RunningInExternal() {

region := os.Getenv(config.DefaultRegionEnvVar)

if region == "" {

log.Warnf("%s is not specified while running in external (non-EC2) environment. Using default region: %s",

config.DefaultRegionEnvVar, defaultRegion)

}

d.region = defaultRegion

return d.region

}

region, err := d.metadata.Region()

if err != nil {

log.Warn("Could not retrieve the region from EC2 Instance Metadata. Error: %s", err.Error())

region = defaultRegion

}

d.region = region

return d.region

}

// getBucketRegion returns a region that contains the agent's bucket

func (d *Downloader) getPartitionBucketRegion() string {

region := d.getRegion()

destination, err := config.GetAgentPartitionBucketRegion(region)

if err != nil {

log.Warnf("Current region not supported, using the default region (%s) for downloader, err: %v", config.DefaultRegionName, err)

return config.DefaultRegionName

}

return destination

}

// DownloadAgent downloads a copy of the Agent and performs an

// integrity check of 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

}

tempFileName, err := d.getPublishedTarball()

if err != nil {

return err

}

defer func() { // clean up temp file

if _, err := d.fs.Stat(tempFileName); err == nil { // if temp file exists, remove it

log.Debugf("Removing temp file %s", tempFileName)

d.fs.Remove(tempFileName)

}

}()

publishedTarballReader, err := d.fs.Open(tempFileName)

if err != nil {

return err

}

defer publishedTarballReader.Close()

md5hash := md5.New()

_, err = d.fs.Copy(md5hash, publishedTarballReader)

if err != nil {

return err

}

calculatedMd5Sum := md5hash.Sum(nil)

calculatedMd5SumString := fmt.Sprintf("%x", calculatedMd5Sum)

log.Debugf("Expected MD5 %q", publishedMd5Sum)

log.Debugf("Calculated MD5 %q", calculatedMd5SumString)

if publishedMd5Sum != calculatedMd5SumString {

agentTarballName, err := config.AgentRemoteTarballKey()

if err != nil {

return errors.New("downloaded agent does not match expected checksum")

}

return errors.Errorf("downloaded agent %q does not match expected checksum", agentTarballName)

}

log.Debugf("Attempting to rename %s to %s", tempFileName, config.AgentTarball())

return d.fs.Rename(tempFileName, config.AgentTarball())

}

func (d *Downloader) getPublishedMd5Sum() (string, error) {

objectKey, err := config.AgentRemoteTarballMD5Key()

if err != nil {

return "", errors.Wrap(err, "failed to determine md5 file for download")

}

tempMd5FileName, err := d.s3Downloader.downloadFile(objectKey)

if err != nil {

return "", errors.Wrap(err, "failed to download md5 file for published tarball")

}

tempMd5File, err := d.fs.Open(tempMd5FileName)

if err != nil {

return "", errors.Wrap(err, "failed to open temporary md5 file")

}

defer func() { // clean up temp file

log.Debugf("Removing temp file %s", tempMd5FileName)

d.fs.Remove(tempMd5FileName)

}()

body, err := d.fs.ReadAll(tempMd5File)

if err != nil {

return "", errors.Wrap(err, "failed to read from temporary md5 file")

}

return strings.TrimSpace(string(body)), nil

}

func (d *Downloader) getPublishedTarball() (string, error) {

objectKey, err := config.AgentRemoteTarballKey()

if err != nil {

return "", errors.Wrap(err, "failed to determine download tarball")

}

tempAgentFileName, err := d.s3Downloader.downloadFile(objectKey)

if err != nil {

return "", errors.Wrap(err, "failed to download published tarball")

}

return tempAgentFileName, nil

}

// LoadCachedAgent returns an io.ReadCloser of the Agent from the cache

func (d *Downloader) LoadCachedAgent() (io.ReadCloser, error) {

return d.fs.Open(config.AgentTarball())

}

// RecordCachedAgent writes the StatusCached state to disk to record a newly

// cached or loaded agent image; this prevents StatusReloadNeeded from

// being interpreted after the reload.

func (d *Downloader) RecordCachedAgent() error {

data := []byte(fmt.Sprintf("%d", StatusCached))

return d.fs.WriteFile(config.CacheState(), data, orwPerm)

}

// LoadDesiredAgent returns an io.ReadCloser of the Agent indicated by the desiredImageLocatorFile

// (/var/cache/ecs/desired-image). The desiredImageLocatorFile must contain as the beginning of the file the name of

// the file containing the desired image (interpreted as a basename) and ending in a newline. Only the first line is

// read, with the rest of the file reserved for future use.

func (d *Downloader) LoadDesiredAgent() (io.ReadCloser, error) {

desiredImageFile, err := d.getDesiredImageFile()

if err != nil {

return nil, err

}

return d.fs.Open(desiredImageFile)

}

func (d *Downloader) getDesiredImageFile() (string, error) {

file, err := d.fs.Open(config.DesiredImageLocatorFile())

if err != nil {

return "", err

}

defer file.Close()

reader := bufio.NewReader(file)

desiredImageString, err := reader.ReadString('\n')

if err != nil {

return "", err

}

desiredImageFile := strings.TrimSpace(config.CacheDirectory() + "/" + d.fs.Base(desiredImageString))

return desiredImageFile, nil

}