feat: support download directory white list (GopeedLab#752)

donetkit · Sep 20, 2024 · 2dc7048 · 2dc7048
1 parent d375046
commit 2dc7048
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 19 deletions.
diff --git a/pkg/base/error.go b/pkg/base/error.go
diff --git a/pkg/download/downloader.go b/pkg/download/downloader.go
@@ -857,6 +857,20 @@ func (d *Downloader) doCreate(f fetcher.Fetcher, opts *base.Options) (taskId str
 		opts.Path = storeConfig.DownloadDir
 	}
 
+	// if enable white download directory, check if the download directory is in the white list
+	if len(d.cfg.DownloadDirWhiteList) > 0 {
+		inWhiteList := false
+		for _, dir := range d.cfg.DownloadDirWhiteList {
+			if match, err := filepath.Match(dir, opts.Path); match && err == nil {
+				inWhiteList = true
+				break
+			}
+		}
+		if !inWhiteList {
+			return "", errors.New("download directory is not in white list")
+		}
+	}
+
 	fm, err := d.parseFm(f.Meta().Req.URL)
 	if err != nil {
 		return

diff --git a/pkg/download/downloader_test.go b/pkg/download/downloader_test.go
@@ -82,7 +82,45 @@ func TestDownloader_Create(t *testing.T) {
 	want := test.FileMd5(test.BuildFile)
 	got := test.FileMd5(test.DownloadFile)
 	if want != got {
-		t.Errorf("Download() got = %v, want %v", got, want)
+		t.Errorf("Downloader_Create() got = %v, want %v", got, want)
+	}
+}
+
+func TestDownloader_CreateNotInWhite(t *testing.T) {
+	listener := test.StartTestFileServer()
+	defer listener.Close()
+
+	downloader := NewDownloader(&DownloaderConfig{
+		DownloadDirWhiteList: []string{"./downloads"},
+	})
+	if err := downloader.Setup(); err != nil {
+		t.Fatal(err)
+	}
+	defer downloader.Clear()
+	req := &base.Request{
+		URL: "http://" + listener.Addr().String() + "/" + test.BuildName,
+	}
+	rr, err := downloader.Resolve(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	downloader.Listener(func(event *Event) {
+		if event.Key == EventKeyDone {
+			wg.Done()
+		}
+	})
+	_, err = downloader.Create(rr.ID, &base.Options{
+		Path: test.Dir,
+		Name: test.DownloadName,
+		Extra: http.OptsExtra{
+			Connections: 4,
+		},
+	})
+	if !strings.Contains(err.Error(), "white") {
+		t.Errorf("TestDownloader_CreateNotInWhite() got = %v, want %v", err.Error(), "not in white list")
 	}
 }
 

diff --git a/pkg/download/model.go b/pkg/download/model.go
@@ -116,9 +116,10 @@ type DownloaderConfig struct {
 	Controller    *controller.Controller
 	FetchManagers []fetcher.FetcherManager
 
-	RefreshInterval int `json:"refreshInterval"` // RefreshInterval time duration to refresh task progress(ms)
-	Storage         Storage
-	StorageDir      string
+	RefreshInterval      int `json:"refreshInterval"` // RefreshInterval time duration to refresh task progress(ms)
+	Storage              Storage
+	StorageDir           string
+	DownloadDirWhiteList []string `json:"downloadDirWhiteList"`
 
 	ProductionMode bool
 

diff --git a/pkg/rest/model/server.go b/pkg/rest/model/server.go
@@ -14,13 +14,14 @@ const (
 )
 
 type StartConfig struct {
-	Network         string                      `json:"network"`
-	Address         string                      `json:"address"`
-	RefreshInterval int                         `json:"refreshInterval"`
-	Storage         Storage                     `json:"storage"`
-	StorageDir      string                      `json:"storageDir"`
-	ApiToken        string                      `json:"apiToken"`
-	DownloadConfig  *base.DownloaderStoreConfig `json:"downloadConfig"`
+	Network           string                      `json:"network"`
+	Address           string                      `json:"address"`
+	RefreshInterval   int                         `json:"refreshInterval"`
+	Storage           Storage                     `json:"storage"`
+	StorageDir        string                      `json:"storageDir"`
+	WhiteDownloadDirs []string                    `json:"whiteDownloadDirs"`
+	ApiToken          string                      `json:"apiToken"`
+	DownloadConfig    *base.DownloaderStoreConfig `json:"downloadConfig"`
 
 	ProductionMode bool