navi/internal/protocols/hysteria2/engine.go

package hysteria2

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"vpnclient/internal/config"
)

// Engine runs the official apernet/hysteria client (Hysteria 2).
type Engine struct {
	mu      sync.Mutex
	cmd     *exec.Cmd
	cancel  context.CancelFunc
	done    chan struct{}
	workdir string
	profile config.Profile
	stderr  io.Writer
	running bool
}

func New(stderr io.Writer) *Engine {
	if stderr == nil {
		stderr = os.Stderr
	}
	return &Engine{stderr: stderr}
}

func (e *Engine) Protocol() config.Protocol { return config.ProtocolHysteria2 }

func (e *Engine) Start(ctx context.Context, profile config.Profile, binDir string) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.running {
		return fmt.Errorf("hysteria2: already running")
	}

	bin, err := ResolveBinary(binDir)
	if err != nil {
		return err
	}

	workdir, err := os.MkdirTemp("", "vpnclient-hy2-*")
	if err != nil {
		return fmt.Errorf("hysteria2: temp dir: %w", err)
	}

	cfgPath := filepath.Join(workdir, "config.yaml")
	if err := writeRuntimeConfig(cfgPath, profile); err != nil {
		os.RemoveAll(workdir)
		return err
	}

	runCtx, cancel := context.WithCancel(context.Background())
	cmd := exec.CommandContext(runCtx, bin, "-c", cfgPath)
	cmd.Dir = workdir
	cmd.Stdout = e.stderr
	cmd.Stderr = e.stderr
	applySysProcAttr(cmd)
	env := os.Environ()
	for k, v := range profile.Env {
		env = append(env, k+"="+v)
	}
	cmd.Env = env

	if err := cmd.Start(); err != nil {
		cancel()
		os.RemoveAll(workdir)
		return fmt.Errorf("hysteria2: start %s: %w", bin, err)
	}

	done := make(chan struct{})
	e.cmd = cmd
	e.cancel = cancel
	e.done = done
	e.workdir = workdir
	e.profile = profile
	e.running = true

	go func() {
		err := cmd.Wait()
		if err != nil {
			fmt.Fprintf(e.stderr, "hysteria2: process ended: %v\n", err)
		}
		e.mu.Lock()
		e.cleanupLocked()
		e.mu.Unlock()
		close(done)
	}()

	e.mu.Unlock()
	timer := time.NewTimer(900 * time.Millisecond)
	defer timer.Stop()
	var startErr error
	select {
	case <-done:
		startErr = fmt.Errorf("hysteria2: process exited immediately; check URI/password and binary")
	case <-ctx.Done():
		_ = e.Stop()
		startErr = ctx.Err()
	case <-timer.C:
		e.mu.Lock()
		alive := e.running
		e.mu.Unlock()
		if !alive {
			startErr = fmt.Errorf("hysteria2: process exited during startup")
		}
	}
	e.mu.Lock()
	return startErr
}

func (e *Engine) Stop() error {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.stopLocked()
}

func (e *Engine) stopLocked() error {
	if !e.running || e.cmd == nil {
		return nil
	}
	done := e.done
	if e.cancel != nil {
		e.cancel()
	}
	proc := e.cmd.Process
	e.mu.Unlock()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		if proc != nil {
			_ = proc.Kill()
		}
		<-done
	}
	e.mu.Lock()
	return nil
}

func (e *Engine) cleanupLocked() {
	if e.workdir != "" {
		_ = os.RemoveAll(e.workdir)
		e.workdir = ""
	}
	e.cmd = nil
	e.cancel = nil
	e.running = false
}

func (e *Engine) Running() bool {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.running
}

func (e *Engine) LocalHTTPProxy() (string, bool) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if !e.running {
		return "", false
	}
	return e.profile.HTTPListenHostPort()
}

func (e *Engine) LocalSOCKSProxy() (string, bool) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if !e.running {
		return "", false
	}
	return e.profile.SOCKSListenHostPort()
}

func writeRuntimeConfig(path string, profile config.Profile) error {
	server, _, err := NormalizeShareLink(profile.Proxy)
	if err != nil {
		return err
	}
	socks := "127.0.0.1:1080"
	httpListen := "127.0.0.1:1081"
	if hp, ok := profile.SOCKSListenHostPort(); ok {
		socks = hp
	}
	if hp, ok := profile.HTTPListenHostPort(); ok {
		httpListen = hp
	}

	var b strings.Builder
	b.WriteString("# generated by Navis — do not edit\n")
	b.WriteString("server: " + yamlQuote(server) + "\n")
	b.WriteString("socks5:\n")
	b.WriteString("  listen: " + yamlQuote(socks) + "\n")
	b.WriteString("http:\n")
	b.WriteString("  listen: " + yamlQuote(httpListen) + "\n")
	return os.WriteFile(path, []byte(b.String()), 0o600)
}

func yamlQuote(s string) string {
	// Always quote — URIs and host:port can confuse YAML.
	escaped := strings.ReplaceAll(s, `\`, `\\`)
	escaped = strings.ReplaceAll(escaped, `"`, `\"`)
	return `"` + escaped + `"`
}