Files

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)
}
}