Improve export download file handling (#350)

Author: xuanyu66

internal/util/download.go

@@ -21,6 +21,9 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"runtime"
+	"strings"
+	"unicode"
 
 	"github.com/go-resty/resty/v2"
 )
@@ -60,17 +63,74 @@ func GetResponse(url string, debug bool) (*http.Response, error) {
 
 // CreateFile creates a file if it does not exist
 func CreateFile(path, fileName string) (*os.File, error) {
-	filePath := filepath.Join(path, fileName)
-	if _, err := os.Stat(filePath); err == nil {
-		return nil, fmt.Errorf("file already exists")
+	filePath, err := safeDownloadPath(fileName)
+	if err != nil {
+		return nil, err
 	}
-	file, err := os.Create(filePath)
+
+	if path == "" {
+		path = "."
+	}
+	root, err := os.OpenRoot(path)
 	if err != nil {
 		return nil, err
 	}
+	defer root.Close()
+
+	file, err := root.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
+	if err != nil {
+		if os.IsExist(err) {
+			return nil, fmt.Errorf("file already exists")
+		}
+		return nil, fmt.Errorf("download destination %q cannot be created: %w", fileName, err)
+	}
 	return file, nil
 }
 
+func safeDownloadPath(fileName string) (string, error) {
+	if fileName == "" {
+		return "", downloadDestinationError(fileName, "cannot be empty")
+	}
+	if fileName == "." || fileName == ".." {
+		return "", downloadDestinationError(fileName, "must refer to a file")
+	}
+	if filepath.IsAbs(fileName) || isWindowsAbs(fileName) {
+		return "", downloadDestinationError(fileName, "must be relative to the output path")
+	}
+	if containsControlCharacter(fileName) {
+		return "", downloadDestinationError(fileName, "contains unsupported characters")
+	}
+
+	clean := filepath.Clean(fileName)
+	if clean == "." || clean == ".." {
+		return "", downloadDestinationError(fileName, "must refer to a file")
+	}
+	if strings.HasPrefix(clean, ".."+string(os.PathSeparator)) || filepath.IsAbs(clean) {
+		return "", downloadDestinationError(fileName, "is outside the output path")
+	}
+
+	return clean, nil
+}
+
+func downloadDestinationError(fileName, reason string) error {
+	return fmt.Errorf("download destination %q %s", fileName, reason)
+}
+
+func containsControlCharacter(s string) bool {
+	return strings.ContainsRune(s, 0) || strings.IndexFunc(s, unicode.IsControl) >= 0
+}
+
+func isWindowsAbs(path string) bool {
+	if runtime.GOOS == "windows" {
+		return false
+	}
+	if len(path) >= 3 && path[1] == ':' && (path[2] == '\\' || path[2] == '/') {
+		c := path[0]
+		return ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z')
+	}
+	return strings.HasPrefix(path, `\\`)
+}
+
 // CreateFolder creates a folder if it does not exist
 func CreateFolder(path string) error {
 	if path == "" {

internal/util/download_test.go

@@ -0,0 +1,175 @@
+// Copyright 2026 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package util
+
+import (
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+)
+
+func TestCreateFileAllowsRelativeDestinations(t *testing.T) {
+	tests := []struct {
+		name     string
+		fileName string
+	}{
+		{
+			name:     "normal file name",
+			fileName: "a.sql.gz",
+		},
+		{
+			name:     "subdirectory",
+			fileName: filepath.Join("folder", "a.sql.gz"),
+		},
+		{
+			name:     "nested subdirectory",
+			fileName: filepath.Join("a", "b", "c.sql.gz"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			base := t.TempDir()
+			parent := filepath.Dir(filepath.Join(base, tt.fileName))
+			if err := os.MkdirAll(parent, 0755); err != nil {
+				t.Fatalf("MkdirAll() error = %v", err)
+			}
+
+			file, err := CreateFile(base, tt.fileName)
+			if err != nil {
+				t.Fatalf("CreateFile() error = %v", err)
+			}
+			file.Close()
+
+			if _, err := os.Stat(filepath.Join(base, tt.fileName)); err != nil {
+				t.Fatalf("expected file inside base: %v", err)
+			}
+		})
+	}
+}
+
+func TestCreateFileRejectsUnsupportedDestinations(t *testing.T) {
+	base := t.TempDir()
+
+	tests := []struct {
+		name     string
+		fileName string
+	}{
+		{
+			name:     "path outside output path",
+			fileName: filepath.Join("..", "..", "tmp", "pwned"),
+		},
+		{
+			name:     "absolute path",
+			fileName: filepath.Join(string(os.PathSeparator), "tmp", "pwned"),
+		},
+		{
+			name:     "empty file name",
+			fileName: "",
+		},
+		{
+			name:     "current directory",
+			fileName: ".",
+		},
+		{
+			name:     "parent directory",
+			fileName: "..",
+		},
+		{
+			name:     "cleans to current directory",
+			fileName: filepath.Join("a", ".."),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			file, err := CreateFile(base, tt.fileName)
+			if err == nil {
+				file.Close()
+				t.Fatalf("CreateFile() succeeded for %q", tt.fileName)
+			}
+		})
+	}
+}
+
+func TestCreateFileDoesNotOverwriteExistingDestination(t *testing.T) {
+	base := t.TempDir()
+	path := filepath.Join(base, "a.sql.gz")
+	if err := os.WriteFile(path, []byte("existing"), 0644); err != nil {
+		t.Fatalf("WriteFile() error = %v", err)
+	}
+
+	file, err := CreateFile(base, "a.sql.gz")
+	if err == nil {
+		file.Close()
+		t.Fatalf("CreateFile() succeeded for existing destination")
+	}
+
+	content, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("ReadFile() error = %v", err)
+	}
+	if string(content) != "existing" {
+		t.Fatalf("existing destination was overwritten: %q", string(content))
+	}
+}
+
+func TestCreateFileUsesCurrentDirectoryWhenBaseIsEmpty(t *testing.T) {
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Fatalf("Getwd() error = %v", err)
+	}
+	base := t.TempDir()
+	if err := os.Chdir(base); err != nil {
+		t.Fatalf("Chdir() error = %v", err)
+	}
+	t.Cleanup(func() {
+		if err := os.Chdir(wd); err != nil {
+			t.Fatalf("restore working directory: %v", err)
+		}
+	})
+
+	file, err := CreateFile("", "a.sql.gz")
+	if err != nil {
+		t.Fatalf("CreateFile() error = %v", err)
+	}
+	file.Close()
+
+	if _, err := os.Stat(filepath.Join(base, "a.sql.gz")); err != nil {
+		t.Fatalf("expected file in current directory: %v", err)
+	}
+}
+
+func TestCreateFileRejectsSymlinkedParentOutsideBase(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("symlink behavior requires additional privileges on Windows")
+	}
+
+	base := t.TempDir()
+	outside := t.TempDir()
+	if err := os.Symlink(outside, filepath.Join(base, "link")); err != nil {
+		t.Fatalf("Symlink() error = %v", err)
+	}
+
+	file, err := CreateFile(base, filepath.Join("link", "export.sql.gz"))
+	if err == nil {
+		file.Close()
+		t.Fatalf("CreateFile() succeeded through symlinked parent")
+	}
+	if _, statErr := os.Stat(filepath.Join(outside, "export.sql.gz")); !os.IsNotExist(statErr) {
+		t.Fatalf("expected no file outside base, stat error = %v", statErr)
+	}
+}