mysql-tunnel/cmd/tunnel/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	"mysql-tunnel/internal/app"
	"mysql-tunnel/internal/config"
	"mysql-tunnel/internal/gui"
	"mysql-tunnel/internal/logger"
	"mysql-tunnel/internal/winsvc"
)

const version = "0.5.0"

func main() {
	os.Exit(run(os.Args[1:]))
}

func run(args []string) int {
	fs := flag.NewFlagSet("mysql-tunnel", flag.ContinueOnError)
	fs.SetOutput(os.Stderr)
	configPath := fs.String("config", "config.yaml", "path to config.yaml")
	listenHost := fs.String("host", "", "override listen.host")
	listenPort := fs.Int("port", 0, "override listen.port")
	doTest := fs.Bool("test", false, "POST actn=C and SELECT 1 through the HTTP tunnel")
	showVersion := fs.Bool("version", false, "print version and exit")
	doGUI := fs.Bool("gui", false, "open the configuration window (Windows)")
	doConsole := fs.Bool("console", false, "listen in the foreground without a window")
	doInstall := fs.Bool("install", false, "install Windows service (Administrator)")
	doUninstall := fs.Bool("uninstall", false, "uninstall Windows service (Administrator)")
	doStart := fs.Bool("start", false, "start Windows service")
	doStop := fs.Bool("stop", false, "stop Windows service")
	doRestart := fs.Bool("restart", false, "restart Windows service")
	if err := fs.Parse(args); err != nil {
		return 2
	}

	if *showVersion {
		fmt.Println("mysql-tunnel", version)
		return 0
	}
	if *doInstall {
		return cmdInstall(*configPath)
	}
	if *doUninstall {
		return report(winsvc.Uninstall())
	}
	if *doStart {
		return report(winsvc.Start())
	}
	if *doStop {
		return report(winsvc.Stop())
	}
	if *doRestart {
		_ = winsvc.Stop()
		return report(winsvc.Start())
	}

	if winsvc.IsService() {
		return runHeadless(*configPath, *listenHost, *listenPort, false)
	}

	if wantGUI(fs, *doGUI, *doConsole, *doTest) {
		return gui.Run(*configPath)
	}

	return runHeadless(*configPath, *listenHost, *listenPort, *doTest)
}

func wantGUI(fs *flag.FlagSet, gui, console, doTest bool) bool {
	if console || doTest {
		return false
	}
	if gui {
		return true
	}
	if runtime.GOOS != "windows" {
		return false
	}
	// Double-click / no flags → GUI. `--config` without `--gui` stays CLI.
	if flagWasSet(fs, "config") || flagWasSet(fs, "host") || flagWasSet(fs, "port") {
		return false
	}
	return true
}

func flagWasSet(fs *flag.FlagSet, name string) bool {
	found := false
	fs.Visit(func(f *flag.Flag) {
		if f.Name == name {
			found = true
		}
	})
	return found
}

func runHeadless(configPath, listenHost string, listenPort int, doTest bool) int {
	resolved, err := config.ResolveConfigPath(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "config error: %s\n", err)
		return 1
	}
	cfg, err := config.Load(resolved, config.Overrides{
		ListenHost: listenHost,
		ListenPort: listenPort,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "config error: %s\n", logger.RedactError(err))
		return 1
	}
	cfg.ResolvePaths()

	log, err := logger.New(cfg.LogFile, cfg.LogLevel, cfg.LogSQL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "logger error: %v\n", err)
		return 1
	}
	defer log.Close()

	for _, w := range cfg.Warnings {
		log.Warn(w)
	}

	if doTest {
		fmt.Println("mysql-tunnel --test")
		r := app.Test(context.Background(), cfg, log)
		fmt.Println(strings.ReplaceAll(r.Text(), "\r\n", "\n"))
		if !r.OK() {
			return 1
		}
		return 0
	}

	runSrv := func(ctx context.Context) error {
		return app.Serve(ctx, cfg, log)
	}

	if winsvc.IsService() {
		if err := winsvc.Run(runSrv); err != nil {
			log.Error("service stopped", logger.RedactError(err))
			return 1
		}
		return 0
	}

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
	defer cancel()
	fmt.Printf("mysql-tunnel listening on %s:%d\n", cfg.ListenHost, cfg.ListenPort)
	fmt.Println("Yingdao: Host=127.0.0.1 Port=" + strconv.Itoa(cfg.ListenPort) + " User=" + cfg.LocalUsername)
	fmt.Println("Tunnel:", config.RedactedTunnelURL(cfg.TunnelURL))
	fmt.Println("Config:", cfg.Path)
	fmt.Println("Log:", cfg.LogFile)
	if err := runSrv(ctx); err != nil {
		log.Error("server stopped", logger.RedactError(err))
		return 1
	}
	log.Info("server stopped")
	return 0
}

func cmdInstall(configFlag string) int {
	exe, err := winsvc.ExePath()
	if err != nil {
		fmt.Fprintf(os.Stderr, "install: cannot locate exe: %v\n", err)
		return 1
	}
	src, err := config.ResolveConfigPath(configFlag)
	if err != nil {
		fmt.Fprintf(os.Stderr, "install: %s\n", err)
		fmt.Fprintln(os.Stderr, "put a filled config.yaml next to mysql-tunnel.exe or pass --config")
		return 1
	}
	if _, err := config.Load(src, config.Overrides{}); err != nil {
		fmt.Fprintf(os.Stderr, "install: invalid config: %s\n", logger.RedactError(err))
		return 1
	}
	dest := filepath.Join(filepath.Dir(exe), "config.yaml")
	srcAbs, _ := filepath.Abs(src)
	destAbs, _ := filepath.Abs(dest)
	if filepath.Clean(srcAbs) != filepath.Clean(destAbs) {
		if err := winsvc.CopyFile(srcAbs, destAbs); err != nil {
			fmt.Fprintf(os.Stderr, "install: copy config: %v\n", err)
			return 1
		}
		fmt.Println("copied config.yaml ->", destAbs)
	}
	if err := winsvc.Install(exe, destAbs); err != nil {
		fmt.Fprintf(os.Stderr, "install: %v\n", err)
		return 1
	}
	if err := winsvc.Start(); err != nil {
		fmt.Fprintf(os.Stderr, "service installed but start failed: %v\n", err)
		fmt.Println("try: mysql-tunnel.exe --start")
		return 1
	}
	fmt.Println("Windows service installed and started:", winsvc.Name)
	fmt.Println("  exe:   ", exe)
	fmt.Println("  config:", destAbs)
	fmt.Println("  stop:    mysql-tunnel.exe --stop")
	fmt.Println("  uninstall: mysql-tunnel.exe --uninstall")
	return 0
}

func report(err error) int {
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	fmt.Println("ok")
	return 0
}