dashboard/cmd/buildlet: support for writing files from tgz URL directly

Client + server support, and gomote flags.

Change-Id: I91320f47731f8c69b84c4961028bfbbdfc85467a
Reviewed-on: https://go-review.googlesource.com/3029
Reviewed-by: Andrew Gerrand <adg@golang.org>

3 files changed, 94 insertions, 29 deletions

diff --git a/dashboard/buildlet/buildletclient.go b/dashboard/buildlet/buildletclient.go
index d9ecbd1..ef9c0a2 100644
--- a/dashboard/buildlet/buildletclient.go
+++ b/dashboard/buildlet/buildletclient.go
@@ -58,25 +58,50 @@ func (c *Client) do(req *http.Request) (*http.Response, error) {
 	return c.httpClient.Do(req)
 }
 
-// PutTarball writes files to the remote buildlet.
-// The Reader must be of a tar.gz file.
-func (c *Client) PutTarball(r io.Reader) error {
-	req, err := http.NewRequest("PUT", c.URL()+"/writetgz", r)
-	if err != nil {
-		return err
-	}
+// doOK sends the request and expects a 200 OK response.
+func (c *Client) doOK(req *http.Request) error {
 	res, err := c.do(req)
 	if err != nil {
 		return err
 	}
 	defer res.Body.Close()
-	if res.StatusCode/100 != 2 {
+	if res.StatusCode != http.StatusOK {
 		slurp, _ := ioutil.ReadAll(io.LimitReader(res.Body, 4<<10))
 		return fmt.Errorf("%v; body: %s", res.Status, slurp)
 	}
 	return nil
 }
 
+// PutTar writes files to the remote buildlet, rooted at the relative
+// directory dir.
+// If dir is empty, they're placed at the root of the buildlet's work directory.
+// The dir is created if necessary.
+// The Reader must be of a tar.gz file.
+func (c *Client) PutTar(r io.Reader, dir string) error {
+	req, err := http.NewRequest("PUT", c.URL()+"/writetgz?dir="+url.QueryEscape(dir), r)
+	if err != nil {
+		return err
+	}
+	return c.doOK(req)
+}
+
+// PutTarFromURL tells the buildlet to download the tar.gz file from tarURL
+// and write it to dir, a relative directory from the workdir.
+// If dir is empty, they're placed at the root of the buildlet's work directory.
+// The dir is created if necessary.
+// The url must be of a tar.gz file.
+func (c *Client) PutTarFromURL(tarURL, dir string) error {
+	form := url.Values{
+		"url": {tarURL},
+	}
+	req, err := http.NewRequest("POST", c.URL()+"/writetgz?dir="+url.QueryEscape(dir), strings.NewReader(form.Encode()))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	return c.doOK(req)
+}
+
 // ExecOpts are options for a remote command invocation.
 type ExecOpts struct {
 	// Output is the output of stdout and stderr.
diff --git a/dashboard/cmd/buildlet/buildlet.go b/dashboard/cmd/buildlet/buildlet.go
index 63ec7f4..b4d3f7e 100644
--- a/dashboard/cmd/buildlet/buildlet.go
+++ b/dashboard/cmd/buildlet/buildlet.go
@@ -26,6 +26,7 @@ import (
 	"log"
 	"net"
 	"net/http"
+	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -183,11 +184,49 @@ func handleRoot(w http.ResponseWriter, r *http.Request) {
 }
 
 func handleWriteTGZ(w http.ResponseWriter, r *http.Request) {
-	if r.Method != "PUT" {
-		http.Error(w, "requires PUT method", http.StatusBadRequest)
+	var tgz io.Reader
+	switch r.Method {
+	case "PUT":
+		tgz = r.Body
+	case "POST":
+		urlStr := r.FormValue("url")
+		if urlStr == "" {
+			http.Error(w, "missing url POST param", http.StatusBadRequest)
+			return
+		}
+		res, err := http.Get(urlStr)
+		if err != nil {
+			http.Error(w, fmt.Sprintf("fetching URL %s: %v", urlStr, err), http.StatusInternalServerError)
+			return
+		}
+		defer res.Body.Close()
+		if res.StatusCode != http.StatusOK {
+			http.Error(w, fmt.Sprintf("fetching provided url: %s", res.Status), http.StatusInternalServerError)
+			return
+		}
+		tgz = res.Body
+	default:
+		http.Error(w, "requires PUT or POST method", http.StatusBadRequest)
 		return
 	}
-	err := untar(r.Body, *scratchDir)
+
+	urlParam, _ := url.ParseQuery(r.URL.RawQuery)
+	baseDir := *scratchDir
+	if dir := urlParam.Get("dir"); dir != "" {
+		dir = filepath.FromSlash(dir)
+		if strings.Contains(dir, "../") {
+			// This is a remote code execution daemon, so security is kinda pointless, but:
+			http.Error(w, "bogus dir", http.StatusBadRequest)
+			return
+		}
+		baseDir = filepath.Join(baseDir, dir)
+		if err := os.MkdirAll(baseDir, 0755); err != nil {
+			http.Error(w, "mkdir of base: "+err.Error(), http.StatusInternalServerError)
+			return
+		}
+	}
+
+	err := untar(tgz, baseDir)
 	if err != nil {
 		status := http.StatusInternalServerError
 		if he, ok := err.(httpStatuser); ok {
diff --git a/dashboard/cmd/gomote/put.go b/dashboard/cmd/gomote/put.go
index 31595c2..3ccc3ff 100644
--- a/dashboard/cmd/gomote/put.go
+++ b/dashboard/cmd/gomote/put.go
@@ -10,7 +10,6 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"net/http"
 	"os"
 )
 
@@ -23,12 +22,23 @@ func putTar(args []string) error {
 		os.Exit(1)
 	}
 	var rev string
-	fs.StringVar(&rev, "gorev", "", "If non-empty, git hash to download from gerrit and put to the buildlet. e.g. 886b02d705ff for Go 1.4.1")
+	fs.StringVar(&rev, "gorev", "", "If non-empty, git hash to download from gerrit and put to the buildlet. e.g. 886b02d705ff for Go 1.4.1. This just maps to the --URL flag, so the two options are mutually exclusive.")
+	var dir string
+	fs.StringVar(&dir, "dir", "", "relative directory from buildlet's work dir to extra tarball into")
+	var tarURL string
+	fs.StringVar(&tarURL, "url", "", "URL of tarball, instead of provided file.")
 
 	fs.Parse(args)
 	if fs.NArg() < 1 || fs.NArg() > 2 {
 		fs.Usage()
 	}
+	if rev != "" {
+		if tarURL != "" {
+			fmt.Fprintln(os.Stderr, "--gorev and --url are mutually exclusive")
+			fs.Usage()
+		}
+		tarURL = "https://go.googlesource.com/go/+archive/" + rev + ".tar.gz"
+	}
 
 	name := fs.Arg(0)
 	bc, err := namedClient(name)
@@ -36,24 +46,15 @@ func putTar(args []string) error {
 		return err
 	}
 
-	var tgz io.Reader = os.Stdin
-	if rev != "" {
+	if tarURL != "" {
 		if fs.NArg() != 1 {
 			fs.Usage()
 		}
-		// TODO(bradfitz): tell the buildlet to do this
-		// itself, to avoid network to & from home networks.
-		// Staying Google<->Google will be much faster.
-		res, err := http.Get("https://go.googlesource.com/go/+archive/" + rev + ".tar.gz")
-		if err != nil {
-			return fmt.Errorf("Error fetching rev %s from Gerrit: %v", rev, err)
-		}
-		defer res.Body.Close()
-		if res.StatusCode != 200 {
-			return fmt.Errorf("Error fetching rev %s from Gerrit: %v", rev, res.Status)
-		}
-		tgz = res.Body
-	} else if fs.NArg() == 2 && fs.Arg(1) != "-" {
+		return bc.PutTarFromURL(tarURL, dir)
+	}
+
+	var tgz io.Reader = os.Stdin
+	if fs.NArg() == 2 && fs.Arg(1) != "-" {
 		f, err := os.Open(fs.Arg(1))
 		if err != nil {
 			return err
@@ -61,7 +62,7 @@ func putTar(args []string) error {
 		defer f.Close()
 		tgz = f
 	}
-	return bc.PutTarball(tgz)
+	return bc.PutTar(tgz, dir)
 }
 
 // put single files