dnshelper/hostfile/hostfile_test.go

package hostfile_test

import (
	"ekdns/hostfile"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

// ---------------------------------------------------------------------------
// fakeFileSystem — implements hostfile.FileSystem for unit tests
// ---------------------------------------------------------------------------

type fakeFileInfo struct {
	mode os.FileMode
}

func (f fakeFileInfo) Name() string       { return "" }
func (f fakeFileInfo) Size() int64        { return 0 }
func (f fakeFileInfo) Mode() os.FileMode  { return f.mode }
func (f fakeFileInfo) ModTime() time.Time { return time.Time{} }
func (f fakeFileInfo) IsDir() bool        { return false }
func (f fakeFileInfo) Sys() any           { return nil }

type writeCall struct {
	path string
	data []byte
	perm os.FileMode
}

type renameCall struct {
	src string
	dst string
}

type fakeFileSystem struct {
	readFiles  map[string][]byte
	readErrors map[string]error
	statMode   os.FileMode
	statErr    error
	renameErr  error
	removeErr  error

	// tracking
	createTempDir string
	lastTempFile  *os.File
	writeCalls    []writeCall
	renameCalls   []renameCall
	removedPaths  []string
	chmodPath     string
	chmodMode     os.FileMode

	// real temp dir on disk for CreateTemp to use
	realTempDir string
}

func newFakeFS(t *testing.T) *fakeFileSystem {
	t.Helper()
	return &fakeFileSystem{
		readFiles:   make(map[string][]byte),
		readErrors:  make(map[string]error),
		realTempDir: t.TempDir(),
		statMode:    0644,
	}
}

func (f *fakeFileSystem) ReadFile(path string) ([]byte, error) {
	if err, ok := f.readErrors[path]; ok {
		return nil, err
	}
	if data, ok := f.readFiles[path]; ok {
		return data, nil
	}
	return nil, fmt.Errorf("file not found in fake fs: %s", path)
}

func (f *fakeFileSystem) WriteFile(path string, data []byte, perm os.FileMode) error {
	f.writeCalls = append(f.writeCalls, writeCall{path: path, data: data, perm: perm})
	return nil
}

func (f *fakeFileSystem) Stat(path string) (os.FileInfo, error) {
	if f.statErr != nil {
		return nil, f.statErr
	}
	return fakeFileInfo{mode: f.statMode}, nil
}

func (f *fakeFileSystem) CreateTemp(dir, pattern string) (*os.File, error) {
	f.createTempDir = dir
	tf, err := os.CreateTemp(f.realTempDir, pattern)
	if err == nil {
		f.lastTempFile = tf
	}
	return tf, err
}

func (f *fakeFileSystem) Rename(oldpath, newpath string) error {
	f.renameCalls = append(f.renameCalls, renameCall{src: oldpath, dst: newpath})
	return f.renameErr
}

func (f *fakeFileSystem) Remove(path string) error {
	f.removedPaths = append(f.removedPaths, path)
	return f.removeErr
}

func (f *fakeFileSystem) Chmod(path string, mode os.FileMode) error {
	f.chmodPath = path
	f.chmodMode = mode
	return nil
}

// containsRemoved returns true if path was passed to Remove.
func (f *fakeFileSystem) containsRemoved(path string) bool {
	for _, p := range f.removedPaths {
		if p == path {
			return true
		}
	}
	return false
}

// ---------------------------------------------------------------------------
// Read tests
// ---------------------------------------------------------------------------

func TestRead_ParsesFileCorrectly(t *testing.T) {
	fs := newFakeFS(t)
	content := strings.Join([]string{
		"127.0.0.1 localhost",
		hostfile.StartMarker,
		"1.1.1.1\thost1",
		hostfile.EndMarker,
		"# end",
	}, "\n") + "\n"
	hostsPath := "/fake/hosts"
	fs.readFiles[hostsPath] = []byte(content)

	mgr := hostfile.NewManager(fs)
	hf, err := mgr.Read(hostsPath)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if hf.Path != hostsPath {
		t.Errorf("expected Path=%q, got %q", hostsPath, hf.Path)
	}
	if !hf.HasManagedBlock {
		t.Error("expected HasManagedBlock=true")
	}
	if len(hf.ManagedContent) != 1 || hf.ManagedContent[0] != "1.1.1.1\thost1" {
		t.Errorf("unexpected ManagedContent: %v", hf.ManagedContent)
	}
	if len(hf.PrefixContent) != 1 {
		t.Errorf("expected 1 prefix line, got %v", hf.PrefixContent)
	}
	if len(hf.PostfixContent) != 1 {
		t.Errorf("expected 1 postfix line, got %v", hf.PostfixContent)
	}
}

func TestRead_HandlesMissingFile(t *testing.T) {
	fs := newFakeFS(t)
	hostsPath := "/fake/hosts"
	fs.readErrors[hostsPath] = errors.New("no such file")

	mgr := hostfile.NewManager(fs)
	_, err := mgr.Read(hostsPath)
	if err == nil {
		t.Fatal("expected error for missing file")
	}
	if !strings.Contains(err.Error(), hostsPath) {
		t.Errorf("error should mention the path, got: %v", err)
	}
}

// ---------------------------------------------------------------------------
// Write tests
// ---------------------------------------------------------------------------

// buildTestHostsFile builds a HostsFile where OriginalContent differs from
// the assembled content (so Write proceeds).
func buildTestHostsFile(hostsPath string) *hostfile.HostsFile {
	prefix := []string{"127.0.0.1 localhost"}
	managed := []string{"1.1.1.1\thost1"}
	postfix := []string{"# end"}
	// OriginalContent is different (e.g. empty managed block originally)
	original := prefix
	return &hostfile.HostsFile{
		Path:            hostsPath,
		OriginalContent: original,
		PrefixContent:   prefix,
		ManagedContent:  managed,
		PostfixContent:  postfix,
		HasManagedBlock: false,
	}
}

func TestWrite_CreatesBackupBeforeWriting(t *testing.T) {
	fs := newFakeFS(t)
	hostsPath := filepath.Join(t.TempDir(), "hosts")
	backupDir := t.TempDir()
	hostsContent := "127.0.0.1 localhost\n"
	fs.readFiles[hostsPath] = []byte(hostsContent)

	mgr := hostfile.NewManager(fs)
	hf := buildTestHostsFile(hostsPath)
	if err := mgr.Write(hf, backupDir); err != nil {
		t.Fatalf("Write failed: %v", err)
	}

	// Verify backup was created via WriteFile
	found := false
	for _, wc := range fs.writeCalls {
		if strings.HasPrefix(wc.path, backupDir) && strings.Contains(wc.path, "hosts.bak.") {
			found = true
			break
		}
	}
	if !found {
		t.Errorf("expected backup WriteFile call in backupDir %q, write calls: %v", backupDir, fs.writeCalls)
	}
}

func TestWrite_UsesTempFileInSameDirectory(t *testing.T) {
	fs := newFakeFS(t)
	hostsDir := t.TempDir()
	hostsPath := filepath.Join(hostsDir, "hosts")
	backupDir := t.TempDir()
	fs.readFiles[hostsPath] = []byte("127.0.0.1 localhost\n")

	mgr := hostfile.NewManager(fs)
	hf := buildTestHostsFile(hostsPath)
	if err := mgr.Write(hf, backupDir); err != nil {
		t.Fatalf("Write failed: %v", err)
	}

	if fs.createTempDir != hostsDir {
		t.Errorf("expected CreateTemp dir=%q, got %q", hostsDir, fs.createTempDir)
	}
}

func TestWrite_RenamesTempOverHostsFile(t *testing.T) {
	fs := newFakeFS(t)
	hostsPath := filepath.Join(t.TempDir(), "hosts")
	backupDir := t.TempDir()
	fs.readFiles[hostsPath] = []byte("127.0.0.1 localhost\n")

	mgr := hostfile.NewManager(fs)
	hf := buildTestHostsFile(hostsPath)
	if err := mgr.Write(hf, backupDir); err != nil {
		t.Fatalf("Write failed: %v", err)
	}

	if len(fs.renameCalls) == 0 {
		t.Fatal("expected Rename to be called")
	}
	lastRename := fs.renameCalls[len(fs.renameCalls)-1]
	if lastRename.dst != hostsPath {
		t.Errorf("expected Rename dst=%q, got %q", hostsPath, lastRename.dst)
	}
}

func TestWrite_PreservesOriginalPermissions(t *testing.T) {
	fs := newFakeFS(t)
	fs.statMode = 0640
	hostsPath := filepath.Join(t.TempDir(), "hosts")
	backupDir := t.TempDir()
	fs.readFiles[hostsPath] = []byte("127.0.0.1 localhost\n")

	mgr := hostfile.NewManager(fs)
	hf := buildTestHostsFile(hostsPath)
	if err := mgr.Write(hf, backupDir); err != nil {
		t.Fatalf("Write failed: %v", err)
	}

	if fs.chmodMode != 0640 {
		t.Errorf("expected Chmod mode=0640, got %v", fs.chmodMode)
	}
}

func TestWrite_DeletesBackupOnSuccess(t *testing.T) {
	fs := newFakeFS(t)
	hostsPath := filepath.Join(t.TempDir(), "hosts")
	backupDir := t.TempDir()
	fs.readFiles[hostsPath] = []byte("127.0.0.1 localhost\n")

	mgr := hostfile.NewManager(fs)
	hf := buildTestHostsFile(hostsPath)
	if err := mgr.Write(hf, backupDir); err != nil {
		t.Fatalf("Write failed: %v", err)
	}

	// Find the backup path from writeCalls
	var backupPath string
	for _, wc := range fs.writeCalls {
		if strings.HasPrefix(wc.path, backupDir) && strings.Contains(wc.path, "hosts.bak.") {
			backupPath = wc.path
			break
		}
	}
	if backupPath == "" {
		t.Fatal("backup not found in writeCalls")
	}
	if !fs.containsRemoved(backupPath) {
		t.Errorf("expected backup %q to be removed on success, removed paths: %v", backupPath, fs.removedPaths)
	}
}

func TestWrite_CleansUpTempOnFailure(t *testing.T) {
	fs := newFakeFS(t)
	fs.renameErr = errors.New("rename failed")
	hostsPath := filepath.Join(t.TempDir(), "hosts")
	backupDir := t.TempDir()
	fs.readFiles[hostsPath] = []byte("127.0.0.1 localhost\n")

	mgr := hostfile.NewManager(fs)
	hf := buildTestHostsFile(hostsPath)
	err := mgr.Write(hf, backupDir)
	if err == nil {
		t.Fatal("expected Write to fail when Rename fails")
	}

	if fs.lastTempFile == nil {
		t.Fatal("expected CreateTemp to have been called")
	}
	tempPath := fs.lastTempFile.Name()
	if !fs.containsRemoved(tempPath) {
		t.Errorf("expected temp file %q to be removed on failure, removed: %v", tempPath, fs.removedPaths)
	}
}

func TestWrite_SkipsIfContentUnchanged(t *testing.T) {
	fs := newFakeFS(t)
	hostsPath := filepath.Join(t.TempDir(), "hosts")
	backupDir := t.TempDir()

	prefix := []string{"127.0.0.1 localhost"}
	managed := []string{"1.1.1.1\thost1"}
	postfix := []string{"# end"}
	// OriginalContent matches assembled content exactly
	assembled := hostfile.AssembleContent(prefix, managed, postfix)

	hf := &hostfile.HostsFile{
		Path:            hostsPath,
		OriginalContent: assembled,
		PrefixContent:   prefix,
		ManagedContent:  managed,
		PostfixContent:  postfix,
		HasManagedBlock: true,
	}

	mgr := hostfile.NewManager(fs)
	if err := mgr.Write(hf, backupDir); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if fs.lastTempFile != nil {
		t.Error("expected no CreateTemp call when content is unchanged")
	}
	if len(fs.writeCalls) != 0 {
		t.Errorf("expected no WriteFile calls, got: %v", fs.writeCalls)
	}
}