github.com/manicqin/nomad@v0.9.5/client/allocrunner/taskrunner/getter/getter.go (about)

     1  package getter
     2  
     3  import (
     4  	"fmt"
     5  	"net/url"
     6  	"path/filepath"
     7  	"strings"
     8  	"sync"
     9  
    10  	gg "github.com/hashicorp/go-getter"
    11  	"github.com/hashicorp/nomad/nomad/structs"
    12  )
    13  
    14  var (
    15  	// getters is the map of getters suitable for Nomad. It is initialized once
    16  	// and the lock is used to guard access to it.
    17  	getters map[string]gg.Getter
    18  	lock    sync.Mutex
    19  
    20  	// supported is the set of download schemes supported by Nomad
    21  	supported = []string{"http", "https", "s3", "hg", "git", "gcs"}
    22  )
    23  
    24  const (
    25  	// gitSSHPrefix is the prefix for downloading via git using ssh
    26  	gitSSHPrefix = "git@github.com:"
    27  )
    28  
    29  // EnvReplacer is an interface which can interpolate environment variables and
    30  // is usually satisfied by taskenv.TaskEnv.
    31  type EnvReplacer interface {
    32  	ReplaceEnv(string) string
    33  }
    34  
    35  // getClient returns a client that is suitable for Nomad downloading artifacts.
    36  func getClient(src string, mode gg.ClientMode, dst string) *gg.Client {
    37  	lock.Lock()
    38  	defer lock.Unlock()
    39  
    40  	// Return the pre-initialized client
    41  	if getters == nil {
    42  		getters = make(map[string]gg.Getter, len(supported))
    43  		for _, getter := range supported {
    44  			if impl, ok := gg.Getters[getter]; ok {
    45  				getters[getter] = impl
    46  			}
    47  		}
    48  	}
    49  
    50  	return &gg.Client{
    51  		Src:     src,
    52  		Dst:     dst,
    53  		Mode:    mode,
    54  		Getters: getters,
    55  		Umask:   060000000,
    56  	}
    57  }
    58  
    59  // getGetterUrl returns the go-getter URL to download the artifact.
    60  func getGetterUrl(taskEnv EnvReplacer, artifact *structs.TaskArtifact) (string, error) {
    61  	source := taskEnv.ReplaceEnv(artifact.GetterSource)
    62  
    63  	// Handle an invalid URL when given a go-getter url such as
    64  	// git@github.com:hashicorp/nomad.git
    65  	gitSSH := false
    66  	if strings.HasPrefix(source, gitSSHPrefix) {
    67  		gitSSH = true
    68  		source = source[len(gitSSHPrefix):]
    69  	}
    70  
    71  	u, err := url.Parse(source)
    72  	if err != nil {
    73  		return "", fmt.Errorf("failed to parse source URL %q: %v", artifact.GetterSource, err)
    74  	}
    75  
    76  	// Build the url
    77  	q := u.Query()
    78  	for k, v := range artifact.GetterOptions {
    79  		q.Add(k, taskEnv.ReplaceEnv(v))
    80  	}
    81  	u.RawQuery = q.Encode()
    82  
    83  	// Add the prefix back
    84  	url := u.String()
    85  	if gitSSH {
    86  		url = fmt.Sprintf("%s%s", gitSSHPrefix, url)
    87  	}
    88  
    89  	return url, nil
    90  }
    91  
    92  // GetArtifact downloads an artifact into the specified task directory.
    93  func GetArtifact(taskEnv EnvReplacer, artifact *structs.TaskArtifact, taskDir string) error {
    94  	url, err := getGetterUrl(taskEnv, artifact)
    95  	if err != nil {
    96  		return newGetError(artifact.GetterSource, err, false)
    97  	}
    98  
    99  	// Download the artifact
   100  	dest := filepath.Join(taskDir, artifact.RelativeDest)
   101  
   102  	// Convert from string getter mode to go-getter const
   103  	mode := gg.ClientModeAny
   104  	switch artifact.GetterMode {
   105  	case structs.GetterModeFile:
   106  		mode = gg.ClientModeFile
   107  	case structs.GetterModeDir:
   108  		mode = gg.ClientModeDir
   109  	}
   110  
   111  	if err := getClient(url, mode, dest).Get(); err != nil {
   112  		return newGetError(url, err, true)
   113  	}
   114  
   115  	return nil
   116  }
   117  
   118  // GetError wraps the underlying artifact fetching error with the URL. It
   119  // implements the RecoverableError interface.
   120  type GetError struct {
   121  	URL         string
   122  	Err         error
   123  	recoverable bool
   124  }
   125  
   126  func newGetError(url string, err error, recoverable bool) *GetError {
   127  	return &GetError{
   128  		URL:         url,
   129  		Err:         err,
   130  		recoverable: recoverable,
   131  	}
   132  }
   133  
   134  func (g *GetError) Error() string {
   135  	return g.Err.Error()
   136  }
   137  
   138  func (g *GetError) IsRecoverable() bool {
   139  	return g.recoverable
   140  }