From 288ed40e8fbadec0c26d3ce4c2c589e77a1fa89f Mon Sep 17 00:00:00 2001
From: Sayantan Das <das.sayantan94@gmail.com>
Date: Wed, 13 May 2020 08:36:43 +0530
Subject: [PATCH] Retry on failed download of plugins from marketplace.
 (#14176)

* Retry on failed download of plugins from marketplace.

* minor polishing.

* Handling http errors during downloading
Adding unit test for DownloadFromURL

* Adding suggested changes from the PR
Close response body before returning error during progressive retry
Remove separate struct for the downloadURL parameter in download_test
Adding comment to clarify the share retries variable
Changing to NoError and Error in the test assertion.

* Added license header to download_test.go

Co-authored-by: mattermod <mattermod@users.noreply.github.com>
---
 app/download.go      | 24 +++++++++++++--
 app/download_test.go | 72 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 2 deletions(-)
 create mode 100644 app/download_test.go

diff --git a/app/download.go b/app/download.go
index ed84b981db..627bee9d25 100644
--- a/app/download.go
+++ b/app/download.go
@@ -4,11 +4,15 @@
 package app
 
 import (
+	"io"
 	"io/ioutil"
+	"net/http"
 	"net/url"
 	"time"
 
 	"github.com/mattermost/mattermost-server/v5/model"
+	"github.com/mattermost/mattermost-server/v5/utils"
+
 	"github.com/pkg/errors"
 )
 
@@ -34,10 +38,26 @@ func (a *App) DownloadFromURL(downloadURL string) ([]byte, error) {
 	client := a.HTTPService().MakeClient(true)
 	client.Timeout = HTTP_REQUEST_TIMEOUT
 
-	resp, err := client.Get(downloadURL)
+	var resp *http.Response
+	err = utils.ProgressiveRetry(func() error {
+		resp, err = client.Get(downloadURL)
+
+		if err != nil {
+			return errors.Wrapf(err, "failed to fetch from %s", downloadURL)
+		}
+
+		if !(resp.StatusCode >= 200 && resp.StatusCode < 300) {
+			_, _ = io.Copy(ioutil.Discard, resp.Body)
+			_ = resp.Body.Close()
+			return errors.Errorf("failed to fetch from %s", downloadURL)
+		}
+
+		return nil
+	})
 	if err != nil {
-		return nil, errors.Wrapf(err, "failed to fetch from %s", downloadURL)
+		return nil, errors.Wrap(err, "download failed after multiple retries.")
 	}
+
 	defer resp.Body.Close()
 
 	return ioutil.ReadAll(resp.Body)
diff --git a/app/download_test.go b/app/download_test.go
new file mode 100644
index 0000000000..0e3d3e9c8b
--- /dev/null
+++ b/app/download_test.go
@@ -0,0 +1,72 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package app
+
+import (
+	"fmt"
+	"github.com/mattermost/mattermost-server/v5/model"
+	"github.com/stretchr/testify/require"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestDownloadFromURL(t *testing.T) {
+	th := Setup(t).InitBasic()
+	defer th.TearDown()
+
+	app := th.App
+	app.Config().PluginSettings.AllowInsecureDownloadUrl = model.NewBool(true)
+
+	// To keep track of how many times an endpoint is retried. This needs to be reset
+	// for each test run.
+	retries := 0
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/succeeds-after-retry", func(w http.ResponseWriter, r *http.Request) {
+		if retries < 2 {
+			http.Error(w, "Request Timed out", http.StatusGatewayTimeout)
+			retries++
+			return
+		}
+
+		_, _ = w.Write([]byte("Your request is successful."))
+	})
+
+	mux.HandleFunc("/fails-forever", func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "This would fail forever", http.StatusInternalServerError)
+	})
+
+	testServer := httptest.NewServer(mux)
+
+	tests := []struct {
+		name        string
+		downloadURL string
+		wantErr     bool
+	}{
+		{
+			name:        "Should succeed after two retries",
+			downloadURL: fmt.Sprintf("%s/succeeds-after-retry", testServer.URL),
+			wantErr:     false,
+		},
+		{
+			name:        "Should not retry forever",
+			downloadURL: fmt.Sprintf("%s/fails-forever", testServer.URL),
+			wantErr:     true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			retries = 0 // reset the retires
+			_, err := th.App.DownloadFromURL(tt.downloadURL)
+
+			if tt.wantErr {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}