hook.go

package main

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"sync/atomic"
	"time"

	"github.com/kkkgo/mini-ppdns/mlog"
	"github.com/kkkgo/mini-ppdns/pkg/cache"
	"codeberg.org/miekg/dns"
)

const (
	hookCheckTimeout = 30 * time.Second
	hookCmdTimeout   = 60 * time.Second
)

type hookMonitor struct {
	cfg      *HookConfig
	failed   *atomic.Bool
	logger   *mlog.Logger
	dnsCache *cache.Cache[CacheKey, *dns.Msg]
}

func (hm *hookMonitor) run(ctx context.Context) {
	sleepTime := time.Duration(hm.cfg.SleepTime) * time.Second
	retryTime := time.Duration(hm.cfg.RetryTime) * time.Second
	failCount := 0
	wasDown := false

	for {
		err := hm.check(ctx)
		if err == nil {
			failCount = 0
			if wasDown {
				hm.failed.Store(false)
				wasDown = false
				hm.logger.Infow("[hook] main DNS recovered, switching back to main DNS")
				runOptionalCmd(hm.logger, hm.cfg.SwitchMainExec)
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(sleepTime):
			}
		} else {
			failCount++
			hm.logger.DebugEventw("[hook] check failed",
				mlog.Int("count", failCount),
				mlog.Int("threshold", hm.cfg.Count),
				mlog.Err(err))
			if failCount >= hm.cfg.Count && !wasDown {
				hm.failed.Store(true)
				wasDown = true
				// Clear DNS cache so stale main-DNS results are purged immediately
				if hm.dnsCache != nil {
					hm.dnsCache.Flush()
					hm.logger.Infow("[hook] DNS cache flushed")
				}
				hm.logger.Warnw("[hook] main DNS marked DOWN, switching to fallback",
					mlog.Int("failures", failCount))
				// Wait retryTime before executing switch_fall_exec so fallback DNS
				// is active and available for the script (e.g. sending notifications)
				if hm.cfg.SwitchFallExec != "" {
					go func(cmd string) {
						select {
						case <-ctx.Done():
							return
						case <-time.After(retryTime / 2):
							if !hm.failed.Load() {
								return
							}
							runOptionalCmd(hm.logger, cmd)
						}
					}(hm.cfg.SwitchFallExec)
				}
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(retryTime):
			}
		}
	}
}

// shellCommand creates an exec.Cmd that runs cmdStr through the system shell.
// Uses SHELL env var on Unix (fallback /bin/sh), cmd.exe on Windows.
func shellCommand(ctx context.Context, cmdStr string) *exec.Cmd {
	if runtime.GOOS == "windows" {
		return exec.CommandContext(ctx, "cmd", "/C", cmdStr)
	}
	shell := os.Getenv("SHELL")
	if shell == "" {
		shell = "/bin/sh"
	}
	return exec.CommandContext(ctx, shell, "-c", cmdStr)
}

func (hm *hookMonitor) check(ctx context.Context) error {
	checkCtx, cancel := context.WithTimeout(ctx, hookCheckTimeout)
	defer cancel()

	cmd := shellCommand(checkCtx, hm.cfg.Exec)
	output, err := cmd.CombinedOutput()

	exitCode := 0
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			exitCode = exitErr.ExitCode()
		} else {
			// Command failed to run (not found, permission denied, timeout, etc.)
			return fmt.Errorf("execution failed: %v", err)
		}
	}

	// If neither exit_code nor keyword is configured, default: success = exit code 0
	if !hm.cfg.ExitCodeSet && hm.cfg.Keyword == "" {
		if exitCode == 0 {
			return nil
		}
		return fmt.Errorf("exit code %d (expected 0)", exitCode)
	}

	if hm.cfg.ExitCodeSet && exitCode != hm.cfg.ExitCode {
		return fmt.Errorf("exit code %d (expected %d)", exitCode, hm.cfg.ExitCode)
	}
	if hm.cfg.Keyword != "" && !strings.Contains(string(output), hm.cfg.Keyword) {
		return fmt.Errorf("output does not contain keyword %q", hm.cfg.Keyword)
	}
	return nil
}

func runOptionalCmd(logger *mlog.Logger, cmdStr string) {
	if cmdStr == "" {
		return
	}
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), hookCmdTimeout)
		defer cancel()
		cmd := shellCommand(ctx, cmdStr)
		out, err := cmd.CombinedOutput()
		if err != nil {
			logger.Warnw("[hook] switch exec failed",
				mlog.String("cmd", cmdStr),
				mlog.Err(err),
				mlog.String("output", string(out)))
		}
	}()
}