est/internal/ssh/tunnel.go

package ssh

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log"
	"os/exec"
	"strings"
	"time"

	"git.scuroneko.dev/ScuroNeko/est/internal/config"
)

var errTunnelExited = errors.New("ssh tunnel exited unexpectedly")

// RunTunnel starts an SSH process for entry and waits for it to exit.
// When SSH exits, it reconnects according to retry.
// The process and any pending retry wait are terminated when ctx is cancelled.
func RunTunnel(ctx context.Context, configPath string, entry config.ProxyEntry, retry config.RetrySettings) error {
	args, err := prepareArgs(configPath, entry)
	if err != nil {
		return err
	}
	return runWithRetry(ctx, retry.Count, retry.Delay, func() error {
		return runTunnelOnce(ctx, args)
	}, func(attempt int, err error) {
		if retry.Count == -1 {
			log.Printf("ssh tunnel %q failed; retrying in %s (attempt %d): %v", entry.Host, retry.Delay, attempt, err)
			return
		}
		log.Printf("ssh tunnel %q failed; retrying in %s (%d/%d): %v", entry.Host, retry.Delay, attempt, retry.Count, err)
	})
}

func runWithRetry(ctx context.Context, retries int, delay time.Duration, run func() error, onRetry func(int, error)) error {
	if err := ctx.Err(); err != nil {
		return err
	}

	for attempt := 0; ; {
		err := run()
		if err == nil || ctx.Err() != nil {
			return err
		}
		if retries != -1 && attempt >= retries {
			return err
		}

		attempt++
		if onRetry != nil {
			onRetry(attempt, err)
		}

		timer := time.NewTimer(delay)
		select {
		case <-ctx.Done():
			if !timer.Stop() {
				select {
				case <-timer.C:
				default:
				}
			}
			return ctx.Err()
		case <-timer.C:
		}
	}
}

func runTunnelOnce(ctx context.Context, args []string) error {
	var stderr bytes.Buffer
	cmd := exec.CommandContext(ctx, "ssh", args...)
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		if ctx.Err() != nil {
			return fmt.Errorf("ssh stopped: %w", ctx.Err())
		}

		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			return fmt.Errorf(
				"ssh exit code %d: %s",
				exitErr.ExitCode(),
				strings.TrimSpace(stderr.String()),
			)
		}
		return fmt.Errorf("can't run ssh: %w", err)
	}
	return errTunnelExited
}