261 lines
6.4 KiB
Go
261 lines
6.4 KiB
Go
// Copyright 2026 Alibaba Group
|
|
// Licensed under the Apache License, Version 2.0
|
|
|
|
package upgrade
|
|
|
|
import (
|
|
"archive/zip"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
)
|
|
|
|
func TestExtractZip_BasicFiles(t *testing.T) {
|
|
dir := t.TempDir()
|
|
zipPath := filepath.Join(dir, "test.zip")
|
|
targetDir := filepath.Join(dir, "extracted")
|
|
|
|
createTestZip(t, zipPath, map[string]string{
|
|
"README.md": "# Hello",
|
|
"dws/SKILL.md": "# Skill",
|
|
"dws/data.txt": "some data",
|
|
})
|
|
|
|
if err := ExtractZip(zipPath, targetDir); err != nil {
|
|
t.Fatalf("ExtractZip() error = %v", err)
|
|
}
|
|
|
|
assertFileContent(t, filepath.Join(targetDir, "README.md"), "# Hello")
|
|
assertFileContent(t, filepath.Join(targetDir, "dws", "SKILL.md"), "# Skill")
|
|
assertFileContent(t, filepath.Join(targetDir, "dws", "data.txt"), "some data")
|
|
}
|
|
|
|
func TestExtractZip_ZipSlipProtection(t *testing.T) {
|
|
dir := t.TempDir()
|
|
zipPath := filepath.Join(dir, "evil.zip")
|
|
targetDir := filepath.Join(dir, "extracted")
|
|
|
|
f, err := os.Create(zipPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := zip.NewWriter(f)
|
|
// Attempt path traversal
|
|
ew, _ := w.Create("../../etc/passwd")
|
|
ew.Write([]byte("root:x:0:0"))
|
|
// Normal file
|
|
nw, _ := w.Create("normal.txt")
|
|
nw.Write([]byte("ok"))
|
|
w.Close()
|
|
f.Close()
|
|
|
|
if err := ExtractZip(zipPath, targetDir); err != nil {
|
|
t.Fatalf("ExtractZip() error = %v", err)
|
|
}
|
|
|
|
// Path-traversal entry should be skipped
|
|
evilPath := filepath.Join(dir, "..", "etc", "passwd")
|
|
if _, err := os.Stat(evilPath); err == nil {
|
|
t.Error("zip-slip protection failed: traversal file was extracted")
|
|
}
|
|
|
|
// Normal file should exist
|
|
assertFileContent(t, filepath.Join(targetDir, "normal.txt"), "ok")
|
|
}
|
|
|
|
func TestExtractZip_InvalidZip(t *testing.T) {
|
|
dir := t.TempDir()
|
|
badPath := filepath.Join(dir, "bad.zip")
|
|
os.WriteFile(badPath, []byte("not a zip"), 0644)
|
|
|
|
err := ExtractZip(badPath, filepath.Join(dir, "out"))
|
|
if err == nil {
|
|
t.Error("ExtractZip() expected error for invalid zip")
|
|
}
|
|
}
|
|
|
|
func TestExtractZip_EmptyZip(t *testing.T) {
|
|
dir := t.TempDir()
|
|
zipPath := filepath.Join(dir, "empty.zip")
|
|
createTestZip(t, zipPath, map[string]string{})
|
|
|
|
targetDir := filepath.Join(dir, "extracted")
|
|
if err := ExtractZip(zipPath, targetDir); err != nil {
|
|
t.Fatalf("ExtractZip() error = %v for empty zip", err)
|
|
}
|
|
}
|
|
|
|
func TestFindBinaryInDir_TopLevel(t *testing.T) {
|
|
dir := t.TempDir()
|
|
dwsPath := filepath.Join(dir, "dws")
|
|
os.WriteFile(dwsPath, []byte("#!/bin/sh"), 0755)
|
|
|
|
got := FindBinaryInDir(dir)
|
|
if got != dwsPath {
|
|
t.Errorf("FindBinaryInDir() = %q, want %q", got, dwsPath)
|
|
}
|
|
}
|
|
|
|
func TestFindBinaryInDir_Nested(t *testing.T) {
|
|
dir := t.TempDir()
|
|
nested := filepath.Join(dir, "subdir")
|
|
os.MkdirAll(nested, 0755)
|
|
dwsPath := filepath.Join(nested, "dws")
|
|
os.WriteFile(dwsPath, []byte("#!/bin/sh"), 0755)
|
|
|
|
got := FindBinaryInDir(dir)
|
|
if got != dwsPath {
|
|
t.Errorf("FindBinaryInDir() = %q, want %q", got, dwsPath)
|
|
}
|
|
}
|
|
|
|
func TestFindBinaryInDir_Windows(t *testing.T) {
|
|
dir := t.TempDir()
|
|
exePath := filepath.Join(dir, "dws.exe")
|
|
os.WriteFile(exePath, []byte("MZ"), 0755)
|
|
|
|
got := FindBinaryInDir(dir)
|
|
if got != exePath {
|
|
t.Errorf("FindBinaryInDir() = %q, want %q", got, exePath)
|
|
}
|
|
}
|
|
|
|
func TestFindBinaryInDir_NotFound(t *testing.T) {
|
|
dir := t.TempDir()
|
|
os.WriteFile(filepath.Join(dir, "other.bin"), []byte("x"), 0755)
|
|
|
|
got := FindBinaryInDir(dir)
|
|
if got != "" {
|
|
t.Errorf("FindBinaryInDir() = %q, want empty", got)
|
|
}
|
|
}
|
|
|
|
func TestCopyFile(t *testing.T) {
|
|
dir := t.TempDir()
|
|
src := filepath.Join(dir, "src")
|
|
dst := filepath.Join(dir, "dst")
|
|
|
|
os.WriteFile(src, []byte("copy me"), 0644)
|
|
|
|
if err := copyFile(src, dst, 0755); err != nil {
|
|
t.Fatalf("copyFile() error = %v", err)
|
|
}
|
|
|
|
got, _ := os.ReadFile(dst)
|
|
if string(got) != "copy me" {
|
|
t.Errorf("dst content = %q, want %q", string(got), "copy me")
|
|
}
|
|
|
|
}
|
|
|
|
func TestCopyFile_SrcNotExist(t *testing.T) {
|
|
dir := t.TempDir()
|
|
err := copyFile(filepath.Join(dir, "missing"), filepath.Join(dir, "dst"), 0644)
|
|
if err == nil {
|
|
t.Error("copyFile() expected error for missing src")
|
|
}
|
|
}
|
|
|
|
func TestReplaceExeFile_AtomicRename(t *testing.T) {
|
|
dir := t.TempDir()
|
|
src := filepath.Join(dir, "new-binary")
|
|
dst := filepath.Join(dir, "dws")
|
|
|
|
os.WriteFile(dst, []byte("old"), 0755)
|
|
os.WriteFile(src, []byte("new"), 0755)
|
|
|
|
if err := replaceExeFile(src, dst); err != nil {
|
|
t.Fatalf("replaceExeFile() error = %v", err)
|
|
}
|
|
|
|
got, _ := os.ReadFile(dst)
|
|
if string(got) != "new" {
|
|
t.Errorf("dst content = %q, want %q", string(got), "new")
|
|
}
|
|
|
|
// src should no longer exist (it was renamed)
|
|
if _, err := os.Stat(src); err == nil {
|
|
t.Error("src should have been renamed away")
|
|
}
|
|
}
|
|
|
|
func TestWindowsReplace_Basic(t *testing.T) {
|
|
dir := t.TempDir()
|
|
src := filepath.Join(dir, "new-binary")
|
|
dst := filepath.Join(dir, "dws.exe")
|
|
|
|
os.WriteFile(dst, []byte("old-exe"), 0755)
|
|
os.WriteFile(src, []byte("new-exe"), 0755)
|
|
|
|
if err := windowsReplace(src, dst); err != nil {
|
|
t.Fatalf("windowsReplace() error = %v", err)
|
|
}
|
|
|
|
got, _ := os.ReadFile(dst)
|
|
if string(got) != "new-exe" {
|
|
t.Errorf("dst content = %q, want %q", string(got), "new-exe")
|
|
}
|
|
|
|
// .old should have been cleaned up
|
|
if _, err := os.Stat(dst + ".old"); err == nil {
|
|
t.Error(".old file should have been cleaned up")
|
|
}
|
|
}
|
|
|
|
func TestWindowsReplace_CleansUpStaleOld(t *testing.T) {
|
|
dir := t.TempDir()
|
|
src := filepath.Join(dir, "new-binary")
|
|
dst := filepath.Join(dir, "dws.exe")
|
|
oldPath := dst + ".old"
|
|
|
|
os.WriteFile(dst, []byte("current"), 0755)
|
|
os.WriteFile(src, []byte("newer"), 0755)
|
|
os.WriteFile(oldPath, []byte("stale-old"), 0755)
|
|
|
|
if err := windowsReplace(src, dst); err != nil {
|
|
t.Fatalf("windowsReplace() error = %v", err)
|
|
}
|
|
|
|
got, _ := os.ReadFile(dst)
|
|
if string(got) != "newer" {
|
|
t.Errorf("dst content = %q, want %q", string(got), "newer")
|
|
}
|
|
}
|
|
|
|
func TestCleanupStaleFiles(t *testing.T) {
|
|
// Just verify it doesn't panic; the actual files to clean depend on os.Executable()
|
|
CleanupStaleFiles()
|
|
}
|
|
|
|
// --- helpers ---
|
|
|
|
func createTestZip(t *testing.T, zipPath string, files map[string]string) {
|
|
t.Helper()
|
|
f, err := os.Create(zipPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := zip.NewWriter(f)
|
|
for name, content := range files {
|
|
fw, err := w.Create(name)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
fw.Write([]byte(content))
|
|
}
|
|
w.Close()
|
|
f.Close()
|
|
}
|
|
|
|
func assertFileContent(t *testing.T, path, want string) {
|
|
t.Helper()
|
|
got, err := os.ReadFile(path)
|
|
if err != nil {
|
|
t.Errorf("cannot read %s: %v", path, err)
|
|
return
|
|
}
|
|
if string(got) != want {
|
|
t.Errorf("content of %s = %q, want %q", path, string(got), want)
|
|
}
|
|
}
|