v1.0.0

internal/ssh/tunnel.go

@@ -0,0 +1,91 @@
+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
+}