Merge pull request #381 from hashicorp/f-getter-refactor

client/drivers: Refactor to use Getter wrapper

client/driver/exec.go

@@ -2,17 +2,15 @@ package driver
 
 import (
 	"fmt"
-	"log"
-	"path"
 	"path/filepath"
 	"runtime"
 	"syscall"
 	"time"
 
-	"github.com/hashicorp/go-getter"
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/client/driver/executor"
+	"github.com/hashicorp/nomad/client/getter"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -55,29 +53,24 @@ func (d *ExecDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle,
 		return nil, fmt.Errorf("missing command for exec driver")
 	}
 
+	// Create a location to download the artifact.
+	taskDir, ok := ctx.AllocDir.TaskDirs[d.DriverContext.taskName]
+	if !ok {
+		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
+	}
+
 	// Check if an artificat is specified and attempt to download it
 	source, ok := task.Config["artifact_source"]
 	if ok && source != "" {
 		// Proceed to download an artifact to be executed.
-		// We use go-getter to support a variety of protocols, but need to change
-		// file permissions of the resulted download to be executable
-
-		// Create a location to download the artifact.
-		taskDir, ok := ctx.AllocDir.TaskDirs[d.DriverContext.taskName]
-		if !ok {
-			return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
-		}
-		destDir := filepath.Join(taskDir, allocdir.TaskLocal)
-
-		artifactName := path.Base(source)
-		artifactFile := filepath.Join(destDir, artifactName)
-		if err := getter.GetFile(artifactFile, source); err != nil {
-			return nil, fmt.Errorf("Error downloading artifact for Exec driver: %s", err)
-		}
-
-		// Add execution permissions to the newly downloaded artifact
-		if err := syscall.Chmod(artifactFile, 0755); err != nil {
-			log.Printf("[ERR] driver.exec: Error making artifact executable: %s", err)
+		_, err := getter.GetArtifact(
+			filepath.Join(taskDir, allocdir.TaskLocal),
+			task.Config["artifact_source"],
+			task.Config["checksum"],
+			d.logger,
+		)
+		if err != nil {
+			return nil, err
 		}
 	}
 

client/driver/exec_test.go

@@ -123,11 +123,12 @@ func TestExecDriver_Start_Wait(t *testing.T) {
 func TestExecDriver_Start_Artifact_basic(t *testing.T) {
 	ctestutils.ExecCompatible(t)
 	file := "hi_linux_amd64"
+	checksum := "sha256:6f99b4c5184726e601ecb062500aeb9537862434dfe1898dbe5c68d9f50c179c"
 
 	task := &structs.Task{
 		Name: "sleep",
 		Config: map[string]string{
-			"artifact_source": fmt.Sprintf("https://dl.dropboxusercontent.com/u/47675/jar_thing/%s", file),
+			"artifact_source": fmt.Sprintf("https://dl.dropboxusercontent.com/u/47675/jar_thing/%s?checksum=%s", file, checksum),
 			"command":         filepath.Join("$NOMAD_TASK_DIR", file),
 		},
 		Resources: basicResources,

client/driver/java.go

@@ -4,17 +4,16 @@ import (
 	"bytes"
 	"fmt"
 	"os/exec"
-	"path"
 	"path/filepath"
 	"runtime"
 	"strings"
 	"syscall"
 	"time"
 
-	"github.com/hashicorp/go-getter"
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/client/driver/executor"
+	"github.com/hashicorp/nomad/client/getter"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -89,26 +88,24 @@ func (d *JavaDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool,
 }
 
 func (d *JavaDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
-	// Get the jar source
-	source, ok := task.Config["jar_source"]
-	if !ok || source == "" {
-		return nil, fmt.Errorf("missing jar source for Java Jar driver")
-	}
-
 	taskDir, ok := ctx.AllocDir.TaskDirs[d.DriverContext.taskName]
 	if !ok {
 		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
 	}
 
-	destDir := filepath.Join(taskDir, allocdir.TaskLocal)
-
-	// Create a location to download the binary.
-	jarName := path.Base(source)
-	jarPath := filepath.Join(destDir, jarName)
-	if err := getter.GetFile(jarPath, source); err != nil {
-		return nil, fmt.Errorf("Error downloading source for Java driver: %s", err)
+	// Proceed to download an artifact to be executed.
+	path, err := getter.GetArtifact(
+		filepath.Join(taskDir, allocdir.TaskLocal),
+		task.Config["artifact_source"],
+		task.Config["checksum"],
+		d.logger,
+	)
+	if err != nil {
+		return nil, err
 	}
 
+	jarName := filepath.Base(path)
+
 	// Get the environment variables.
 	envVars := TaskEnvironmentVariables(ctx, task)
 

client/driver/java_test.go

@@ -97,10 +97,9 @@ func TestJavaDriver_Start_Wait(t *testing.T) {
 	task := &structs.Task{
 		Name: "demo-app",
 		Config: map[string]string{
-			"jar_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/demoapp.jar",
-			// "jar_source": "https://s3-us-west-2.amazonaws.com/java-jar-thing/demoapp.jar",
-			// "args": "-d64",
-			"jvm_options": "-Xmx2048m -Xms256m",
+			"artifact_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/demoapp.jar",
+			"jvm_options":     "-Xmx2048m -Xms256m",
+			"checksum":        "sha256:58d6e8130308d32e197c5108edd4f56ddf1417408f743097c2e662df0f0b17c8",
 		},
 		Resources: basicResources,
 	}
@@ -145,9 +144,7 @@ func TestJavaDriver_Start_Kill_Wait(t *testing.T) {
 	task := &structs.Task{
 		Name: "demo-app",
 		Config: map[string]string{
-			"jar_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/demoapp.jar",
-			// "jar_source": "https://s3-us-west-2.amazonaws.com/java-jar-thing/demoapp.jar",
-			// "args": "-d64",
+			"artifact_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/demoapp.jar",
 		},
 		Resources: basicResources,
 	}

client/driver/qemu.go

@@ -2,11 +2,8 @@ package driver
 
 import (
 	"bytes"
-	"crypto/sha256"
-	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log"
 	"os"
 	"os/exec"
@@ -17,9 +14,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/hashicorp/go-getter"
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
+	"github.com/hashicorp/nomad/client/getter"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -82,7 +79,7 @@ func (d *QemuDriver) Fingerprint(cfg *config.Config, node *structs.Node) (bool,
 // image and save it to the Drivers Allocation Dir
 func (d *QemuDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle, error) {
 	// Get the image source
-	source, ok := task.Config["image_source"]
+	source, ok := task.Config["artifact_source"]
 	if !ok || source == "" {
 		return nil, fmt.Errorf("Missing source image Qemu driver")
 	}
@@ -99,34 +96,18 @@ func (d *QemuDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandle,
 		return nil, fmt.Errorf("Could not find task directory for task: %v", d.DriverContext.taskName)
 	}
 
-	// Create a location to download the binary.
-	destDir := filepath.Join(taskDir, allocdir.TaskLocal)
-	vmID := fmt.Sprintf("qemu-vm-%s-%s", structs.GenerateUUID(), filepath.Base(source))
-	vmPath := filepath.Join(destDir, vmID)
-	if err := getter.GetFile(vmPath, source); err != nil {
-		return nil, fmt.Errorf("Error downloading artifact for Qemu driver: %s", err)
+	// Proceed to download an artifact to be executed.
+	vmPath, err := getter.GetArtifact(
+		filepath.Join(taskDir, allocdir.TaskLocal),
+		task.Config["artifact_source"],
+		task.Config["checksum"],
+		d.logger,
+	)
+	if err != nil {
+		return nil, err
 	}
 
-	// compute and check checksum
-	if check, ok := task.Config["checksum"]; ok {
-		d.logger.Printf("[DEBUG] Running checksum on (%s)", vmID)
-		hasher := sha256.New()
-		file, err := os.Open(vmPath)
-		if err != nil {
-			return nil, fmt.Errorf("Failed to open file for checksum")
-		}
-
-		defer file.Close()
-		io.Copy(hasher, file)
-
-		sum := hex.EncodeToString(hasher.Sum(nil))
-		if sum != check {
-			return nil, fmt.Errorf(
-				"Error in Qemu: checksums did not match.\nExpected (%s), got (%s)",
-				check,
-				sum)
-		}
-	}
+	vmID := filepath.Base(vmPath)
 
 	// Parse configuration arguments
 	// Create the base arguments

client/driver/qemu_test.go

@@ -54,10 +54,10 @@ func TestQemuDriver_Start(t *testing.T) {
 	task := &structs.Task{
 		Name: "linux",
 		Config: map[string]string{
-			"image_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/linux-0.2.img",
-			"checksum":     "a5e836985934c3392cbbd9b26db55a7d35a8d7ae1deb7ca559dd9c0159572544",
-			"accelerator":  "tcg",
-			"guest_ports":  "22,8080",
+			"artifact_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/linux-0.2.img",
+			"checksum":        "sha256:a5e836985934c3392cbbd9b26db55a7d35a8d7ae1deb7ca559dd9c0159572544",
+			"accelerator":     "tcg",
+			"guest_ports":     "22,8080",
 		},
 		Resources: &structs.Resources{
 			MemoryMB: 512,
@@ -103,11 +103,11 @@ func TestQemuDriver_RequiresMemory(t *testing.T) {
 	task := &structs.Task{
 		Name: "linux",
 		Config: map[string]string{
-			"image_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/linux-0.2.img",
-			"accelerator":  "tcg",
-			"host_port":    "8080",
-			"guest_port":   "8081",
-			"checksum":     "a5e836985934c3392cbbd9b26db55a7d35a8d7ae1deb7ca559dd9c0159572544",
+			"artifact_source": "https://dl.dropboxusercontent.com/u/47675/jar_thing/linux-0.2.img",
+			"accelerator":     "tcg",
+			"host_port":       "8080",
+			"guest_port":      "8081",
+			"checksum":        "sha256:a5e836985934c3392cbbd9b26db55a7d35a8d7ae1deb7ca559dd9c0159572544",
 			// ssh u/p would be here
 		},
 	}

client/driver/raw_exec.go

@@ -2,21 +2,18 @@ package driver
 
 import (
 	"fmt"
-	"log"
 	"os"
 	"os/exec"
-	"path"
 	"path/filepath"
 	"runtime"
 	"strconv"
 	"strings"
-	"syscall"
 	"time"
 
-	"github.com/hashicorp/go-getter"
 	"github.com/hashicorp/nomad/client/allocdir"
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/client/driver/args"
+	"github.com/hashicorp/nomad/client/getter"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -83,23 +80,14 @@ func (d *RawExecDriver) Start(ctx *ExecContext, task *structs.Task) (DriverHandl
 	source, ok := task.Config["artifact_source"]
 	if ok && source != "" {
 		// Proceed to download an artifact to be executed.
-		// We use go-getter to support a variety of protocols, but need to change
-		// file permissions of the resulted download to be executable
-
-		// Create a location to download the artifact.
-		destDir := filepath.Join(taskDir, allocdir.TaskLocal)
-
-		artifactName := path.Base(source)
-		artifactFile := filepath.Join(destDir, artifactName)
-		if err := getter.GetFile(artifactFile, source); err != nil {
-			return nil, fmt.Errorf("Error downloading artifact for Raw Exec driver: %s", err)
-		}
-
-		// Add execution permissions to the newly downloaded artifact
-		if runtime.GOOS != "windows" {
-			if err := syscall.Chmod(artifactFile, 0755); err != nil {
-				log.Printf("[ERR] driver.raw_exec: Error making artifact executable: %s", err)
-			}
+		_, err := getter.GetArtifact(
+			filepath.Join(taskDir, allocdir.TaskLocal),
+			task.Config["artifact_source"],
+			task.Config["checksum"],
+			d.logger,
+		)
+		if err != nil {
+			return nil, err
 		}
 	}
 

client/driver/raw_exec_test.go

@@ -94,19 +94,22 @@ func TestRawExecDriver_StartOpen_Wait(t *testing.T) {
 }
 
 func TestRawExecDriver_Start_Artifact_basic(t *testing.T) {
-	var file string
+	var file, checksum string
 	switch runtime.GOOS {
 	case "darwin":
 		file = "hi_darwin_amd64"
+		checksum = "md5:d7f2fdb13b36dcb7407721d78926b335"
 	default:
 		file = "hi_linux_amd64"
+		checksum = "md5:a9b14903a8942748e4f8474e11f795d3"
 	}
 
 	task := &structs.Task{
 		Name: "sleep",
 		Config: map[string]string{
 			"artifact_source": fmt.Sprintf("https://dl.dropboxusercontent.com/u/47675/jar_thing/%s", file),
 			"command":         filepath.Join("$NOMAD_TASK_DIR", file),
+			"checksum":        checksum,
 		},
 	}
 	driverCtx := testDriverContext(task.Name)

client/getter/getter.go

@@ -0,0 +1,43 @@
+package getter
+
+import (
+	"fmt"
+	"log"
+	"net/url"
+	"path"
+	"path/filepath"
+	"runtime"
+	"strings"
+	"syscall"
+
+	gg "github.com/hashicorp/go-getter"
+)
+
+func GetArtifact(destDir, source, checksum string, logger *log.Logger) (string, error) {
+	if source == "" {
+		return "", fmt.Errorf("Source url is empty in Artifact Getter")
+	}
+	u, err := url.Parse(source)
+	if err != nil {
+		return "", err
+	}
+
+	// if checksum is seperate, apply to source
+	if checksum != "" {
+		source = strings.Join([]string{source, fmt.Sprintf("checksum=%s", checksum)}, "?")
+		logger.Printf("[DEBUG] client.getter: Applying checksum to Artifact Source URL, new url: %s", source)
+	}
+
+	artifactFile := filepath.Join(destDir, path.Base(u.Path))
+	if err := gg.GetFile(artifactFile, source); err != nil {
+		return "", fmt.Errorf("Error downloading artifact: %s", err)
+	}
+
+	// Add execution permissions to the newly downloaded artifact
+	if runtime.GOOS != "windows" {
+		if err := syscall.Chmod(artifactFile, 0755); err != nil {
+			logger.Printf("[ERR] driver.raw_exec: Error making artifact executable: %s", err)
+		}
+	}
+	return artifactFile, nil
+}