navi/internal/protocols/naive/engine.go

package naive

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

	"vpnclient/internal/config"
)

// Engine runs the official Chromium-based naive binary.
// NaiveProxy cannot be reimplemented in pure Go: TLS/H2 fingerprints
// must match Chrome, which requires the upstream naive.exe.
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.ProtocolNaive }

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("naive: already running")
	}

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

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

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

	runCtx, cancel := context.WithCancel(context.Background())
	cmd := exec.CommandContext(runCtx, bin, 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("naive: 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, "naive: process ended: %v\n", err)
		}
		e.mu.Lock()
		e.cleanupLocked()
		e.mu.Unlock()
		close(done)
	}()

	// Release lock while waiting so the Wait goroutine can clean up.
	e.mu.Unlock()
	timer := time.NewTimer(800 * time.Millisecond)
	defer timer.Stop()
	var startErr error
	select {
	case <-done:
		startErr = fmt.Errorf("naive: process exited immediately; check proxy URI, credentials, 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("naive: 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()
}

type runtimeConfig struct {
	Listen              any    `json:"listen"`
	Proxy               string `json:"proxy"`
	InsecureConcurrency int    `json:"insecure-concurrency,omitempty"`
	ExtraHeaders        string `json:"extra-headers,omitempty"`
	HostResolverRules   string `json:"host-resolver-rules,omitempty"`
	Log                 string `json:"log,omitempty"`
	NoPostQuantum       bool   `json:"no-post-quantum,omitempty"`
}

func writeRuntimeConfig(path string, profile config.Profile) error {
	listen := profile.DefaultListen()
	var listenField any = listen
	if len(listen) == 1 {
		listenField = listen[0]
	}
	proxy, err := NormalizeProxyURI(profile.Proxy)
	if err != nil {
		return err
	}
	rc := runtimeConfig{
		Listen:              listenField,
		Proxy:               proxy,
		InsecureConcurrency: profile.InsecureConcurrency,
		ExtraHeaders:        profile.ExtraHeaders,
		HostResolverRules:   profile.HostResolverRules,
		Log:                 profile.Log,
		NoPostQuantum:       profile.NoPostQuantum,
	}
	data, err := json.MarshalIndent(rc, "", "  ")
	if err != nil {
		return err
	}
	data = append(data, '\n')
	return os.WriteFile(path, data, 0o600)
}