dnshelper/hostfile/hostfile.go

package hostfile

import (
	"fmt"
	"math/rand"
	"path/filepath"
	"strings"
	"time"
)

const (
	StartMarker = "# ekDNSHelper <<-> START CONFIG"
	EndMarker   = "# ekDNSHelper <->> END CONFIG"

	// Legacy markers from the pre-rename "DNSHelper" era.
	// ParseManagedBlock recognises these so that hosts files written by
	// the old binary are detected, parsed, and silently upgraded to the
	// current markers on the next write.
	LegacyStartMarker = "# DNSHelper <<-> START CONFIG"
	LegacyEndMarker   = "# DNSHelper <->> END CONFIG"
)

// DNSEntry represents a single IP-to-hostname mapping.
type DNSEntry struct {
	IP       string
	Hostname string
}

// String returns the hosts file line representation: "IP\tHostname".
func (e DNSEntry) String() string {
	return e.IP + "\t" + e.Hostname
}

// HostsFile represents the parsed content of a hosts file.
type HostsFile struct {
	Path            string
	OriginalContent []string
	PrefixContent   []string
	ManagedContent  []string
	PostfixContent  []string
	HasManagedBlock bool
}

// Manager handles hosts file read and atomic write operations.
type Manager struct {
	fs FileSystem
}

// NewManager creates a Manager with the given FileSystem implementation.
func NewManager(fs FileSystem) *Manager {
	return &Manager{fs: fs}
}

// Read reads and parses the hosts file at the given path.
// Returns an error if the file cannot be read or the managed block is corrupt.
func (m *Manager) Read(path string) (*HostsFile, error) {
	data, err := m.fs.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("reading hosts file %s: %w", path, err)
	}

	lines := strings.Split(string(data), "\n")
	// Remove trailing empty line from Split if file ends with newline.
	if len(lines) > 0 && lines[len(lines)-1] == "" {
		lines = lines[:len(lines)-1]
	}
	// Trim \r from each line to handle Windows \r\n line endings.
	// strings.Split on "\n" does not strip \r, unlike bufio.Scanner.
	for i, line := range lines {
		lines[i] = strings.TrimRight(line, "\r")
	}

	prefix, managed, postfix, hasBlock, err := ParseManagedBlock(lines)
	if err != nil {
		return nil, err
	}

	return &HostsFile{
		Path:            path,
		OriginalContent: lines,
		PrefixContent:   prefix,
		ManagedContent:  managed,
		PostfixContent:  postfix,
		HasManagedBlock: hasBlock,
	}, nil
}

// Write atomically writes the hosts file content back to disk.
// It creates a backup before writing and removes it on success.
// Uses the temp-file-and-rename pattern for atomicity (R-002).
// backupDir is the directory for the backup file (typically the exe directory).
func (m *Manager) Write(hf *HostsFile, backupDir string) error {
	// 1. Assemble new content.
	newContent := AssembleContent(hf.PrefixContent, hf.ManagedContent, hf.PostfixContent)

	// 2. Skip write if content is unchanged.
	if slicesEqual(newContent, hf.OriginalContent) {
		return nil
	}

	// 3. Build output bytes (LF line endings).
	var buf strings.Builder
	for _, line := range newContent {
		buf.WriteString(line)
		buf.WriteString("\n")
	}
	output := []byte(buf.String())

	// 4. Get original file permissions so we can restore them on the temp file.
	info, err := m.fs.Stat(hf.Path)
	if err != nil {
		return fmt.Errorf("checking hosts file permissions: %w", err)
	}
	originalMode := info.Mode()

	// 5. Create backup (copy of current hosts file in backupDir).
	backupPath, err := m.createBackup(hf.Path, backupDir)
	if err != nil {
		return fmt.Errorf("creating backup: %w", err)
	}

	// 6. Write to a temp file in the SAME directory as the hosts file.
	//    This is required so the subsequent rename stays on the same filesystem
	//    partition (avoids EXDEV on Unix).
	success := false
	tempFile, err := m.fs.CreateTemp(filepath.Dir(hf.Path), ".ekdns-tmp-*")
	if err != nil {
		m.fs.Remove(backupPath) //nolint:errcheck
		return fmt.Errorf("creating temp file: %w", err)
	}
	tempPath := tempFile.Name()
	defer func() {
		if !success {
			m.fs.Remove(tempPath)   //nolint:errcheck
			m.fs.Remove(backupPath) //nolint:errcheck
		}
	}()

	if _, err := tempFile.Write(output); err != nil {
		tempFile.Close() //nolint:errcheck
		return fmt.Errorf("writing temp file: %w", err)
	}
	// Sync before close for durability (data survives power loss).
	if err := tempFile.Sync(); err != nil {
		tempFile.Close() //nolint:errcheck
		return fmt.Errorf("syncing temp file: %w", err)
	}
	// Close before rename — required on Windows.
	tempFile.Close() //nolint:errcheck

	// 7. Set permissions on temp file to match original.
	if err := m.fs.Chmod(tempPath, originalMode); err != nil {
		return fmt.Errorf("setting temp file permissions: %w", err)
	}

	// 8. Atomic rename (with retry on Windows sharing violations).
	if err := m.renameWithRetry(tempPath, hf.Path); err != nil {
		return fmt.Errorf("renaming temp file to hosts file: %w", err)
	}

	// 9. Success — delete the backup (original was safely replaced).
	success = true
	m.fs.Remove(backupPath) //nolint:errcheck
	return nil
}

// createBackup copies the hosts file to <backupDir>/hosts.bak.<YYYYMMDD>-<4hex>.
func (m *Manager) createBackup(hostsPath, backupDir string) (string, error) {
	data, err := m.fs.ReadFile(hostsPath)
	if err != nil {
		return "", fmt.Errorf("reading hosts file for backup: %w", err)
	}

	timestamp := time.Now().Format("20060102")
	suffix := fmt.Sprintf("%04x", rand.Uint32()&0xFFFF)
	backupName := fmt.Sprintf("hosts.bak.%s-%s", timestamp, suffix)
	backupPath := filepath.Join(backupDir, backupName)

	if err := m.fs.WriteFile(backupPath, data, 0600); err != nil {
		return "", fmt.Errorf("writing backup file %s: %w", backupPath, err)
	}
	return backupPath, nil
}

// renameWithRetry calls Rename and retries once on failure (for Windows
// sharing violations where the hosts file may be briefly locked).
func (m *Manager) renameWithRetry(src, dst string) error {
	err := m.fs.Rename(src, dst)
	if err == nil {
		return nil
	}
	// Single retry after a short delay.
	time.Sleep(200 * time.Millisecond)
	return m.fs.Rename(src, dst)
}

// slicesEqual reports whether a and b contain the same strings in the same order.
func slicesEqual(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}