Add debug headers when downloading plugins (#92579)

pkg/api/plugins.go

@@ -465,7 +465,8 @@ func (hs *HTTPServer) InstallPlugin(c *contextmodel.ReqContext) response.Respons
 	}
 
 	compatOpts := plugins.NewCompatOpts(hs.Cfg.BuildVersion, runtime.GOOS, runtime.GOARCH)
-	err := hs.pluginInstaller.Add(c.Req.Context(), pluginID, dto.Version, compatOpts)
+	ctx := repo.WithRequestOrigin(c.Req.Context(), "api")
+	err := hs.pluginInstaller.Add(ctx, pluginID, dto.Version, compatOpts)
 	if err != nil {
 		var dupeErr plugins.DuplicateError
 		if errors.As(err, &dupeErr) {

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

@@ -146,6 +146,7 @@ func doInstallPlugin(ctx context.Context, pluginID, version string, o pluginInst
 			return err
 		}
 	} else {
+		ctx = repo.WithRequestOrigin(ctx, "cli")
 		archiveInfo, err := repository.GetPluginArchiveInfo(ctx, pluginID, version, compatOpts)
 		if err != nil {
 			return err

pkg/plugins/manager/fakes/fakes.go

@@ -231,7 +231,7 @@ type FakePluginRepo struct {
 	GetPluginArchiveFunc      func(_ context.Context, pluginID, version string, _ repo.CompatOpts) (*repo.PluginArchive, error)
 	GetPluginArchiveByURLFunc func(_ context.Context, archiveURL string, _ repo.CompatOpts) (*repo.PluginArchive, error)
 	GetPluginArchiveInfoFunc  func(_ context.Context, pluginID, version string, _ repo.CompatOpts) (*repo.PluginArchiveInfo, error)
-	PluginVersionFunc         func(pluginID, version string, compatOpts repo.CompatOpts) (repo.VersionData, error)
+	PluginVersionFunc         func(_ context.Context, pluginID, version string, compatOpts repo.CompatOpts) (repo.VersionData, error)
 }
 
 // GetPluginArchive fetches the requested plugin archive.
@@ -260,9 +260,9 @@ func (r *FakePluginRepo) GetPluginArchiveInfo(ctx context.Context, pluginID, ver
 	return &repo.PluginArchiveInfo{}, nil
 }
 
-func (r *FakePluginRepo) PluginVersion(pluginID, version string, compatOpts repo.CompatOpts) (repo.VersionData, error) {
+func (r *FakePluginRepo) PluginVersion(ctx context.Context, pluginID, version string, compatOpts repo.CompatOpts) (repo.VersionData, error) {
 	if r.PluginVersionFunc != nil {
-		return r.PluginVersionFunc(pluginID, version, compatOpts)
+		return r.PluginVersionFunc(ctx, pluginID, version, compatOpts)
 	}
 	return repo.VersionData{}, nil
 }

pkg/plugins/repo/client.go

@@ -35,7 +35,15 @@ func NewClient(skipTLSVerify bool, logger log.PrettyLogger) *Client {
 	}
 }
 
-func (c *Client) Download(_ context.Context, pluginZipURL, checksum string, compatOpts CompatOpts) (*PluginArchive, error) {
+type requestOrigin struct{}
+
+// WithRequestOrigin adds the request origin to the context which is used
+// to set the `grafana-origin` header in the outgoing HTTP request.
+func WithRequestOrigin(ctx context.Context, origin string) context.Context {
+	return context.WithValue(ctx, requestOrigin{}, origin)
+}
+
+func (c *Client) Download(ctx context.Context, pluginZipURL, checksum string, compatOpts CompatOpts) (*PluginArchive, error) {
 	// Create temp file for downloading zip file
 	tmpFile, err := os.CreateTemp("", "*.zip")
 	if err != nil {
@@ -49,7 +57,7 @@ func (c *Client) Download(_ context.Context, pluginZipURL, checksum string, comp
 
 	c.log.Debugf("Installing plugin from %s", pluginZipURL)
 
-	err = c.downloadFile(tmpFile, pluginZipURL, checksum, compatOpts)
+	err = c.downloadFile(ctx, tmpFile, pluginZipURL, checksum, compatOpts)
 	if err != nil {
 		if err := tmpFile.Close(); err != nil {
 			c.log.Warn("Failed to close file", "error", err)
@@ -65,8 +73,8 @@ func (c *Client) Download(_ context.Context, pluginZipURL, checksum string, comp
 	return &PluginArchive{File: rc}, nil
 }
 
-func (c *Client) SendReq(url *url.URL, compatOpts CompatOpts) ([]byte, error) {
+func (c *Client) SendReq(ctx context.Context, url *url.URL, compatOpts CompatOpts) ([]byte, error) {
-	req, err := c.createReq(url, compatOpts)
+	req, err := c.createReq(ctx, url, compatOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -87,7 +95,7 @@ func (c *Client) SendReq(url *url.URL, compatOpts CompatOpts) ([]byte, error) {
 	return io.ReadAll(bodyReader)
 }
 
-func (c *Client) downloadFile(tmpFile *os.File, pluginURL, checksum string, compatOpts CompatOpts) (err error) {
+func (c *Client) downloadFile(ctx context.Context, tmpFile *os.File, pluginURL, checksum string, compatOpts CompatOpts) (err error) {
 	// Try handling URL as a local file path first
 	if _, err := os.Stat(pluginURL); err == nil {
 		// TODO re-verify
@@ -110,13 +118,11 @@ func (c *Client) downloadFile(tmpFile *os.File, pluginURL, checksum string, comp
 		return nil
 	}
 
-	c.retryCount = 0
-
 	defer func() {
 		if r := recover(); r != nil {
 			c.retryCount++
 			if c.retryCount < 3 {
-				c.log.Debug("Failed downloading. Will retry once.")
+				c.log.Debug("Failed downloading. Will retry.")
 				err = tmpFile.Truncate(0)
 				if err != nil {
 					return
@@ -125,7 +131,7 @@ func (c *Client) downloadFile(tmpFile *os.File, pluginURL, checksum string, comp
 				if err != nil {
 					return
 				}
-				err = c.downloadFile(tmpFile, pluginURL, checksum, compatOpts)
+				err = c.downloadFile(ctx, tmpFile, pluginURL, checksum, compatOpts)
 			} else {
 				c.retryCount = 0
 				failure := fmt.Sprintf("%v", r)
@@ -145,8 +151,13 @@ func (c *Client) downloadFile(tmpFile *os.File, pluginURL, checksum string, comp
 
 	// Using no timeout as some plugin archives make take longer to fetch due to size, network performance, etc.
 	// Note: This is also used as part of the grafana plugin install CLI operation
-	bodyReader, err := c.sendReqNoTimeout(u, compatOpts)
+	bodyReader, err := c.sendReqNoTimeout(ctx, u, compatOpts)
 	if err != nil {
+		if c.retryCount < 3 {
+			c.retryCount++
+			c.log.Debug("Failed downloading. Will retry.")
+			err = c.downloadFile(ctx, tmpFile, pluginURL, checksum, compatOpts)
+		}
 		return err
 	}
 	defer func() {
@@ -166,11 +177,14 @@ func (c *Client) downloadFile(tmpFile *os.File, pluginURL, checksum string, comp
 	if len(checksum) > 0 && checksum != fmt.Sprintf("%x", h.Sum(nil)) {
 		return ErrChecksumMismatch(pluginURL)
 	}
+
+	c.retryCount = 0
+
 	return nil
 }
 
-func (c *Client) sendReqNoTimeout(url *url.URL, compatOpts CompatOpts) (io.ReadCloser, error) {
+func (c *Client) sendReqNoTimeout(ctx context.Context, url *url.URL, compatOpts CompatOpts) (io.ReadCloser, error) {
-	req, err := c.createReq(url, compatOpts)
+	req, err := c.createReq(ctx, url, compatOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -182,7 +196,7 @@ func (c *Client) sendReqNoTimeout(url *url.URL, compatOpts CompatOpts) (io.ReadC
 	return c.handleResp(res, compatOpts)
 }
 
-func (c *Client) createReq(url *url.URL, compatOpts CompatOpts) (*http.Request, error) {
+func (c *Client) createReq(ctx context.Context, url *url.URL, compatOpts CompatOpts) (*http.Request, error) {
 	req, err := http.NewRequest(http.MethodGet, url.String(), nil)
 	if err != nil {
 		return nil, err
@@ -201,6 +215,14 @@ func (c *Client) createReq(url *url.URL, compatOpts CompatOpts) (*http.Request,
 		req.Header.Set("grafana-arch", sysArch)
 	}
 
+	if c.retryCount > 0 {
+		req.Header.Set("grafana-retrycount", fmt.Sprintf("%d", c.retryCount))
+	}
+
+	if orig := ctx.Value(requestOrigin{}); orig != nil {
+		req.Header.Set("grafana-origin", orig.(string))
+	}
+
 	return req, err
 }
 

pkg/plugins/repo/client_test.go

@@ -0,0 +1,72 @@
+package repo
+
+import (
+	"archive/zip"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/plugins/log"
+	"github.com/stretchr/testify/require"
+)
+
+func writeFakeZip(w http.ResponseWriter) error {
+	ww := zip.NewWriter(w)
+	_, err := ww.Create("test.txt")
+	if err != nil {
+		return err
+	}
+	return ww.Close()
+}
+
+func Test_Download(t *testing.T) {
+	t.Run("it should download a file", func(t *testing.T) {
+		fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			err := writeFakeZip(w)
+			require.NoError(t, err)
+		}))
+		defer fakeServer.Close()
+		cli := fakeServer.Client()
+		repo := Client{httpClient: *cli, httpClientNoTimeout: *cli, log: log.NewPrettyLogger("test")}
+		_, err := repo.Download(context.Background(), fakeServer.URL, "", CompatOpts{})
+		require.NoError(t, err)
+	})
+
+	t.Run("it should set the origin header", func(t *testing.T) {
+		var origin string
+		fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			origin = r.Header.Get("grafana-origin")
+			err := writeFakeZip(w)
+			require.NoError(t, err)
+		}))
+		defer fakeServer.Close()
+		cli := fakeServer.Client()
+		repo := Client{httpClient: *cli, httpClientNoTimeout: *cli, log: log.NewPrettyLogger("test")}
+		ctx := WithRequestOrigin(context.Background(), "test")
+		_, err := repo.Download(ctx, fakeServer.URL, "", CompatOpts{})
+		require.NoError(t, err)
+		require.Equal(t, "test", origin, "origin header should be set")
+	})
+
+	t.Run("it should retry on error", func(t *testing.T) {
+		var count int
+		fakeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			count++
+			if count < 2 {
+				http.Error(w, "error", http.StatusInternalServerError)
+				return
+			}
+			retryCount := r.Header.Get("grafana-retrycount")
+			require.Equal(t, "2", retryCount, "retry count should be set")
+			err := writeFakeZip(w)
+			require.NoError(t, err)
+		}))
+		defer fakeServer.Close()
+		cli := fakeServer.Client()
+		repo := Client{httpClient: *cli, httpClientNoTimeout: *cli, log: log.NewPrettyLogger("test"), retryCount: 1}
+		_, err := repo.Download(context.Background(), fakeServer.URL, "", CompatOpts{})
+		require.NoError(t, err)
+		require.Equal(t, 2, count, "should retry on error")
+	})
+}

pkg/plugins/repo/ifaces.go

@@ -15,7 +15,7 @@ type Service interface {
 	// GetPluginArchiveInfo fetches information needed for downloading the requested plugin.
 	GetPluginArchiveInfo(ctx context.Context, pluginID, version string, opts CompatOpts) (*PluginArchiveInfo, error)
 	// PluginVersion will return plugin version based on the requested information.
-	PluginVersion(pluginID, version string, compatOpts CompatOpts) (VersionData, error)
+	PluginVersion(ctx context.Context, pluginID, version string, compatOpts CompatOpts) (VersionData, error)
 }
 
 type CompatOpts struct {

pkg/plugins/repo/service.go

@@ -63,8 +63,8 @@ func (m *Manager) GetPluginArchiveByURL(ctx context.Context, pluginZipURL string
 }
 
 // GetPluginArchiveInfo returns the options for downloading the requested plugin (with optional `version`)
-func (m *Manager) GetPluginArchiveInfo(_ context.Context, pluginID, version string, compatOpts CompatOpts) (*PluginArchiveInfo, error) {
+func (m *Manager) GetPluginArchiveInfo(ctx context.Context, pluginID, version string, compatOpts CompatOpts) (*PluginArchiveInfo, error) {
-	v, err := m.PluginVersion(pluginID, version, compatOpts)
+	v, err := m.PluginVersion(ctx, pluginID, version, compatOpts)
 	if err != nil {
 		return nil, err
 	}
@@ -77,8 +77,8 @@ func (m *Manager) GetPluginArchiveInfo(_ context.Context, pluginID, version stri
 }
 
 // PluginVersion will return plugin version based on the requested information
-func (m *Manager) PluginVersion(pluginID, version string, compatOpts CompatOpts) (VersionData, error) {
+func (m *Manager) PluginVersion(ctx context.Context, pluginID, version string, compatOpts CompatOpts) (VersionData, error) {
-	versions, err := m.grafanaCompatiblePluginVersions(pluginID, compatOpts)
+	versions, err := m.grafanaCompatiblePluginVersions(ctx, pluginID, compatOpts)
 	if err != nil {
 		return VersionData{}, err
 	}
@@ -103,7 +103,7 @@ func (m *Manager) downloadURL(pluginID, version string) string {
 }
 
 // grafanaCompatiblePluginVersions will get version info from /api/plugins/$pluginID/versions
-func (m *Manager) grafanaCompatiblePluginVersions(pluginID string, compatOpts CompatOpts) ([]Version, error) {
+func (m *Manager) grafanaCompatiblePluginVersions(ctx context.Context, pluginID string, compatOpts CompatOpts) ([]Version, error) {
 	u, err := url.Parse(m.baseURL)
 	if err != nil {
 		return nil, err
@@ -111,7 +111,7 @@ func (m *Manager) grafanaCompatiblePluginVersions(pluginID string, compatOpts Co
 
 	u.Path = path.Join(u.Path, pluginID, "versions")
 
-	body, err := m.client.SendReq(u, compatOpts)
+	body, err := m.client.SendReq(ctx, u, compatOpts)
 	if err != nil {
 		return nil, err
 	}

pkg/services/pluginsintegration/plugininstaller/service.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/plugins/repo"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/pluginsintegration/pluginstore"
 	"github.com/grafana/grafana/pkg/setting"
@@ -107,6 +108,7 @@ func (s *Service) installPlugins(ctx context.Context) error {
 
 		s.log.Info("Installing plugin", "pluginId", installPlugin.ID, "version", installPlugin.Version)
 		start := time.Now()
+		ctx = repo.WithRequestOrigin(ctx, "preinstall")
 		err := s.pluginInstaller.Add(ctx, installPlugin.ID, installPlugin.Version, compatOpts)
 		if err != nil {
 			var dupeErr plugins.DuplicateError