CLI: Reduce memory usage for plugin installation (#19639)

* grafana-cli: use tmp file when downloading plugin install file

pkg/cmd/grafana-cli/commands/commandstest/fake_api_client.go

@@ -1,12 +1,14 @@
 package commandstest
 
 import (
+	"os"
+
 	"github.com/grafana/grafana/pkg/cmd/grafana-cli/models"
 )
 
 type FakeGrafanaComClient struct {
 	GetPluginFunc      func(pluginId, repoUrl string) (models.Plugin, error)
-	DownloadFileFunc   func(pluginName, filePath, url string, checksum string) (content []byte, err error)
+	DownloadFileFunc   func(pluginName string, tmpFile *os.File, url string, checksum string) (err error)
 	ListAllPluginsFunc func(repoUrl string) (models.PluginRepo, error)
 }
 
@@ -18,12 +20,12 @@ func (client *FakeGrafanaComClient) GetPlugin(pluginId, repoUrl string) (models.
 	return models.Plugin{}, nil
 }
 
-func (client *FakeGrafanaComClient) DownloadFile(pluginName, filePath, url string, checksum string) (content []byte, err error) {
+func (client *FakeGrafanaComClient) DownloadFile(pluginName string, tmpFile *os.File, url string, checksum string) (err error) {
 	if client.DownloadFileFunc != nil {
-		return client.DownloadFileFunc(pluginName, filePath, url, checksum)
+		return client.DownloadFileFunc(pluginName, tmpFile, url, checksum)
 	}
 
-	return make([]byte, 0), nil
+	return nil
 }
 
 func (client *FakeGrafanaComClient) ListAllPlugins(repoUrl string) (models.PluginRepo, error) {

pkg/cmd/grafana-cli/commands/install_command.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"os"
 	"path"
 	"path/filepath"
@@ -107,12 +108,24 @@ func InstallPlugin(pluginName, version string, c utils.CommandLine) error {
 	logger.Infof("into: %v\n", pluginFolder)
 	logger.Info("\n")
 
-	content, err := c.ApiClient().DownloadFile(pluginName, pluginFolder, downloadURL, checksum)
+	// Create temp file for downloading zip file
+	tmpFile, err := ioutil.TempFile("", "*.zip")
 	if err != nil {
+		return errutil.Wrap("Failed to create temporary file", err)
+	}
+	defer os.Remove(tmpFile.Name())
+
+	err = c.ApiClient().DownloadFile(pluginName, tmpFile, downloadURL, checksum)
+	if err != nil {
+		tmpFile.Close()
 		return errutil.Wrap("Failed to download plugin archive", err)
 	}
+	err = tmpFile.Close()
+	if err != nil {
+		return errutil.Wrap("Failed to close tmp file", err)
+	}
 
-	err = extractFiles(content, pluginName, pluginFolder, isInternal)
+	err = extractFiles(tmpFile.Name(), pluginName, pluginFolder, isInternal)
 	if err != nil {
 		return errutil.Wrap("Failed to extract plugin archive", err)
 	}
@@ -197,8 +210,10 @@ func RemoveGitBuildFromName(pluginName, filename string) string {
 
 var permissionsDeniedMessage = "Could not create %s. Permission denied. Make sure you have write access to plugindir"
 
-func extractFiles(body []byte, pluginName string, filePath string, allowSymlinks bool) error {
-	r, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
+func extractFiles(archiveFile string, pluginName string, filePath string, allowSymlinks bool) error {
+	logger.Debugf("Extracting archive %v to %v...\n", archiveFile, filePath)
+
+	r, err := zip.OpenReader(archiveFile)
 	if err != nil {
 		return err
 	}

pkg/cmd/grafana-cli/commands/install_command_test.go

@@ -2,7 +2,7 @@ package commands
 
 import (
 	"fmt"
-	"io/ioutil"
+	"io"
 	"os"
 	"runtime"
 	"testing"
@@ -52,10 +52,8 @@ func TestExtractFiles(t *testing.T) {
 		pluginDir, del := setupFakePluginsDir(t)
 		defer del()
 
-		body, err := ioutil.ReadFile("testdata/grafana-simple-json-datasource-ec18fa4da8096a952608a7e4c7782b4260b41bcf.zip")
-		assert.Nil(t, err)
-
-		err = extractFiles(body, "grafana-simple-json-datasource", pluginDir, false)
+		archive := "testdata/grafana-simple-json-datasource-ec18fa4da8096a952608a7e4c7782b4260b41bcf.zip"
+		err := extractFiles(archive, "grafana-simple-json-datasource", pluginDir, false)
 		assert.Nil(t, err)
 
 		//File in zip has permissions 755
@@ -83,10 +81,7 @@ func TestExtractFiles(t *testing.T) {
 		pluginDir, del := setupFakePluginsDir(t)
 		defer del()
 
-		body, err := ioutil.ReadFile("testdata/plugin-with-symlink.zip")
-		assert.Nil(t, err)
-
-		err = extractFiles(body, "plugin-with-symlink", pluginDir, false)
+		err := extractFiles("testdata/plugin-with-symlink.zip", "plugin-with-symlink", pluginDir, false)
 		assert.Nil(t, err)
 
 		_, err = os.Stat(pluginDir + "/plugin-with-symlink/text.txt")
@@ -100,10 +95,7 @@ func TestExtractFiles(t *testing.T) {
 		pluginDir, del := setupFakePluginsDir(t)
 		defer del()
 
-		body, err := ioutil.ReadFile("testdata/plugin-with-symlink.zip")
-		assert.Nil(t, err)
-
-		err = extractFiles(body, "plugin-with-symlink", pluginDir, true)
+		err := extractFiles("testdata/plugin-with-symlink.zip", "plugin-with-symlink", pluginDir, true)
 		assert.Nil(t, err)
 
 		_, err = os.Stat(pluginDir + "/plugin-with-symlink/symlink_to_txt")
@@ -228,13 +220,15 @@ func setupPluginInstallCmd(t *testing.T, pluginDir string) utils.CommandLine {
 		return plugin, nil
 	}
 
-	client.DownloadFileFunc = func(pluginName, filePath, url string, checksum string) (content []byte, err error) {
+	client.DownloadFileFunc = func(pluginName string, tmpFile *os.File, url string, checksum string) (err error) {
 		assert.Equal(t, "test-plugin-panel", pluginName)
 		assert.Equal(t, "/test-plugin-panel/versions/1.0.0/download", url)
 		assert.Equal(t, "test", checksum)
-		body, err := ioutil.ReadFile("testdata/grafana-simple-json-datasource-ec18fa4da8096a952608a7e4c7782b4260b41bcf.zip")
+		f, err := os.Open("testdata/grafana-simple-json-datasource-ec18fa4da8096a952608a7e4c7782b4260b41bcf.zip")
 		assert.Nil(t, err)
-		return body, nil
+		_, err = io.Copy(tmpFile, f)
+		assert.Nil(t, err)
+		return nil
 	}
 
 	cmd.Client = client

pkg/cmd/grafana-cli/services/api_client.go

@@ -1,9 +1,11 @@
 package services
 
 import (
+	"bufio"
 	"crypto/md5"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/url"
@@ -23,7 +25,7 @@ type GrafanaComClient struct {
 
 func (client *GrafanaComClient) GetPlugin(pluginId, repoUrl string) (models.Plugin, error) {
 	logger.Debugf("getting plugin metadata from: %v pluginId: %v \n", repoUrl, pluginId)
-	body, err := sendRequest(HttpClient, repoUrl, "repo", pluginId)
+	body, err := sendRequestGetBytes(HttpClient, repoUrl, "repo", pluginId)
 
 	if err != nil {
 		if err == ErrNotFoundError {
@@ -42,14 +44,18 @@ func (client *GrafanaComClient) GetPlugin(pluginId, repoUrl string) (models.Plug
 	return data, nil
 }
 
-func (client *GrafanaComClient) DownloadFile(pluginName, filePath, url string, checksum string) (content []byte, err error) {
+func (client *GrafanaComClient) DownloadFile(pluginName string, tmpFile *os.File, url string, checksum string) (err error) {
 	// Try handling url like local file path first
 	if _, err := os.Stat(url); err == nil {
-		bytes, err := ioutil.ReadFile(url)
+		f, err := os.Open(url)
 		if err != nil {
-			return nil, errutil.Wrap("Failed to read file", err)
+			return errutil.Wrap("Failed to read plugin archive", err)
 		}
-		return bytes, nil
+		_, err = io.Copy(tmpFile, f)
+		if err != nil {
+			return errutil.Wrap("Failed to copy plugin archive", err)
+		}
+		return nil
 	}
 
 	client.retryCount = 0
@@ -59,7 +65,15 @@ func (client *GrafanaComClient) DownloadFile(pluginName, filePath, url string, c
 			client.retryCount++
 			if client.retryCount < 3 {
 				logger.Info("Failed downloading. Will retry once.")
-				content, err = client.DownloadFile(pluginName, filePath, url, checksum)
+				err = tmpFile.Truncate(0)
+				if err != nil {
+					return
+				}
+				_, err = tmpFile.Seek(0, 0)
+				if err != nil {
+					return
+				}
+				err = client.DownloadFile(pluginName, tmpFile, url, checksum)
 			} else {
 				client.retryCount = 0
 				failure := fmt.Sprintf("%v", r)
@@ -72,23 +86,28 @@ func (client *GrafanaComClient) DownloadFile(pluginName, filePath, url string, c
 		}
 	}()
 
-	// TODO: this would be better if it was streamed file by file instead of buffered.
 	// Using no timeout here as some plugins can be bigger and smaller timeout would prevent to download a plugin on
 	// slow network. As this is CLI operation hanging is not a big of an issue as user can just abort.
-	body, err := sendRequest(HttpClientNoTimeout, url)
-
+	bodyReader, err := sendRequest(HttpClientNoTimeout, url)
 	if err != nil {
-		return nil, errutil.Wrap("Failed to send request", err)
+		return errutil.Wrap("Failed to send request", err)
 	}
+	defer bodyReader.Close()
 
-	if len(checksum) > 0 && checksum != fmt.Sprintf("%x", md5.Sum(body)) {
-		return nil, xerrors.New("Expected MD5 checksum does not match the downloaded archive. Please contact security@grafana.com.")
+	w := bufio.NewWriter(tmpFile)
+	h := md5.New()
+	if _, err = io.Copy(w, io.TeeReader(bodyReader, h)); err != nil {
+		return errutil.Wrap("Failed to compute MD5 checksum", err)
 	}
-	return body, nil
+	w.Flush()
+	if len(checksum) > 0 && checksum != fmt.Sprintf("%x", h.Sum(nil)) {
+		return xerrors.New("Expected MD5 checksum does not match the downloaded archive. Please contact security@grafana.com.")
+	}
+	return nil
 }
 
 func (client *GrafanaComClient) ListAllPlugins(repoUrl string) (models.PluginRepo, error) {
-	body, err := sendRequest(HttpClient, repoUrl, "repo")
+	body, err := sendRequestGetBytes(HttpClient, repoUrl, "repo")
 
 	if err != nil {
 		logger.Info("Failed to send request", "error", err)
@@ -105,45 +124,61 @@ func (client *GrafanaComClient) ListAllPlugins(repoUrl string) (models.PluginRep
 	return data, nil
 }
 
-func sendRequest(client http.Client, repoUrl string, subPaths ...string) ([]byte, error) {
+func sendRequestGetBytes(client http.Client, repoUrl string, subPaths ...string) ([]byte, error) {
+	bodyReader, err := sendRequest(client, repoUrl, subPaths...)
+	if err != nil {
+		return []byte{}, err
+	}
+	defer bodyReader.Close()
+	return ioutil.ReadAll(bodyReader)
+}
+
+func sendRequest(client http.Client, repoUrl string, subPaths ...string) (io.ReadCloser, error) {
+	req, err := createRequest(repoUrl, subPaths...)
+	if err != nil {
+		return nil, err
+	}
+
+	res, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	return handleResponse(res)
+}
+
+func createRequest(repoUrl string, subPaths ...string) (*http.Request, error) {
 	u, _ := url.Parse(repoUrl)
 	for _, v := range subPaths {
 		u.Path = path.Join(u.Path, v)
 	}
 
 	req, err := http.NewRequest(http.MethodGet, u.String(), nil)
+	if err != nil {
+		return nil, err
+	}
 
 	req.Header.Set("grafana-version", grafanaVersion)
 	req.Header.Set("grafana-os", runtime.GOOS)
 	req.Header.Set("grafana-arch", runtime.GOARCH)
 	req.Header.Set("User-Agent", "grafana "+grafanaVersion)
 
-	if err != nil {
-		return []byte{}, err
-	}
-
-	res, err := client.Do(req)
-	if err != nil {
-		return []byte{}, err
-	}
-	return handleResponse(res)
+	return req, err
 }
 
-func handleResponse(res *http.Response) ([]byte, error) {
+func handleResponse(res *http.Response) (io.ReadCloser, error) {
 	if res.StatusCode == 404 {
-		return []byte{}, ErrNotFoundError
+		return nil, ErrNotFoundError
 	}
 
 	if res.StatusCode/100 != 2 && res.StatusCode/100 != 4 {
-		return []byte{}, fmt.Errorf("Api returned invalid status: %s", res.Status)
+		return nil, fmt.Errorf("Api returned invalid status: %s", res.Status)
 	}
 
-	body, err := ioutil.ReadAll(res.Body)
-	defer res.Body.Close()
-
 	if res.StatusCode/100 == 4 {
-		if len(body) == 0 {
-			return []byte{}, &BadRequestError{Status: res.Status}
+		body, err := ioutil.ReadAll(res.Body)
+		defer res.Body.Close()
+		if err != nil || len(body) == 0 {
+			return nil, &BadRequestError{Status: res.Status}
 		}
 		var message string
 		var jsonBody map[string]string
@@ -153,8 +188,8 @@ func handleResponse(res *http.Response) ([]byte, error) {
 		} else {
 			message = jsonBody["message"]
 		}
-		return []byte{}, &BadRequestError{Status: res.Status, Message: message}
+		return nil, &BadRequestError{Status: res.Status, Message: message}
 	}
 
-	return body, err
+	return res.Body, nil
 }

pkg/cmd/grafana-cli/services/api_client_test.go

@@ -12,8 +12,10 @@ import (
 
 func TestHandleResponse(t *testing.T) {
 	t.Run("Returns body if status == 200", func(t *testing.T) {
-		body, err := handleResponse(makeResponse(200, "test"))
-		assert.Nil(t, err)
+		bodyReader, err := handleResponse(makeResponse(200, "test"))
+		assert.NoError(t, err)
+		body, err := ioutil.ReadAll(bodyReader)
+		assert.NoError(t, err)
 		assert.Equal(t, "test", string(body))
 	})
 
@@ -24,25 +26,25 @@ func TestHandleResponse(t *testing.T) {
 
 	t.Run("Returns message from body if status == 400", func(t *testing.T) {
 		_, err := handleResponse(makeResponse(400, "{ \"message\": \"error_message\" }"))
-		assert.NotNil(t, err)
+		assert.Error(t, err)
 		assert.Equal(t, "error_message", asBadRequestError(t, err).Message)
 	})
 
 	t.Run("Returns body if status == 400 and no message key", func(t *testing.T) {
 		_, err := handleResponse(makeResponse(400, "{ \"test\": \"test_message\"}"))
-		assert.NotNil(t, err)
+		assert.Error(t, err)
 		assert.Equal(t, "{ \"test\": \"test_message\"}", asBadRequestError(t, err).Message)
 	})
 
 	t.Run("Returns Bad request error if status == 400 and no body", func(t *testing.T) {
 		_, err := handleResponse(makeResponse(400, ""))
-		assert.NotNil(t, err)
+		assert.Error(t, err)
 		_ = asBadRequestError(t, err)
 	})
 
 	t.Run("Returns error with invalid status if status == 500", func(t *testing.T) {
 		_, err := handleResponse(makeResponse(500, ""))
-		assert.NotNil(t, err)
+		assert.Error(t, err)
 		assert.Contains(t, err.Error(), "invalid status")
 	})
 }

pkg/cmd/grafana-cli/utils/command_line.go

@@ -1,6 +1,8 @@
 package utils
 
 import (
+	"os"
+
 	"github.com/codegangsta/cli"
 	"github.com/grafana/grafana/pkg/cmd/grafana-cli/models"
 	"github.com/grafana/grafana/pkg/cmd/grafana-cli/services"
@@ -27,7 +29,7 @@ type CommandLine interface {
 
 type ApiClient interface {
 	GetPlugin(pluginId, repoUrl string) (models.Plugin, error)
-	DownloadFile(pluginName, filePath, url string, checksum string) (content []byte, err error)
+	DownloadFile(pluginName string, tmpFile *os.File, url string, checksum string) (err error)
 	ListAllPlugins(repoUrl string) (models.PluginRepo, error)
 }