@automattic/vip 4.1.4 → 4.1.6

package/internal/wpstream/run.go

@@ -1,364 +0,0 @@
-package wpstream
-
-import (
-	"context"
-	"encoding/json/v2"
-	"errors"
-	"fmt"
-	"io"
-	"math"
-	"math/rand/v2"
-	"net/http"
-	"sync"
-	"sync/atomic"
-	"time"
-
-	"github.com/fatih/color"
-
-	"github.com/Automattic/vip/internal/debuglog"
-)
-
-const (
-	nonTTYColumns = 100 // NON_TTY_COLUMNS (vip-wp.js:42)
-	nonTTYRows    = 15  // NON_TTY_ROWS    (vip-wp.js:43)
-)
-
-// errRunDone is injected into a stdout IOStream to interrupt a blocked Read
-// when run() is about to return (exit/cancel received before stdout EOF).
-var errRunDone = errors.New("wpstream: run done")
-
-// Options configure a single Run.
-type Options struct {
-	APIHost    string
-	Token      string
-	GUID       string
-	InputToken string
-	Columns    int
-	Rows       int
-	IsTTY      bool // CR→LF stdin normalization when true (vip-wp.js:51)
-
-	Stdin  io.Reader
-	Stdout io.Writer
-	Stderr io.Writer
-}
-
-// Result carries the terminal outcome (the caller maps ExitCode to os.Exit).
-type Result struct {
-	ExitCode int
-}
-
-// Run connects and executes one WP-CLI command over socket.io.
-// It implements the reconnect/offset loop: on disconnect (before an exit event
-// arrives) it re-dials with exponential backoff and resumes from offset.
-//
-// C1/C2 fix: the loop is entirely self-contained; every Engine is explicitly
-// closed before the next attempt or before returning. No goroutine is launched
-// that outlives its engine.
-func Run(ctx context.Context, opts Options) (Result, error) {
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-
-	var offset atomic.Int64
-	backoff := time.Second
-	const maxBackoff = 5 * time.Second
-	first := true
-
-	for attempt := 1; ; attempt++ {
-		debuglog.Printf(ctx, "@automattic/vip:wp", "socket: connect attempt=%d", attempt)
-		eng, err := Dial(ctx, DialOptions{
-			BaseURL: opts.APIHost,
-			Header:  bearerHeader(opts.Token),
-		})
-		if err != nil {
-			debuglog.Printf(ctx, "@automattic/vip:wp", "socket: connect_error stage=transport")
-			if first {
-				return Result{}, err
-			}
-			if werr := waitBackoff(ctx, &backoff, maxBackoff); werr != nil {
-				return Result{}, werr
-			}
-			continue
-		}
-
-		cli := NewClient(eng, "/wp-cli")
-		ss := NewStreamSocket(ctx, cli)
-
-		// Install a retry handler: the server sends "retry" to signal it wants us to
-		// reconnect. We close the engine after 5 s to force the disconnect.
-		cli.On("retry", func(args []any) {
-			debuglog.Printf(ctx, "@automattic/vip:wp", "socket: retry")
-			go func() {
-				select {
-				case <-time.After(5 * time.Second):
-				case <-ctx.Done():
-					return
-				}
-				eng.Close()
-			}()
-		})
-
-		if err := cli.Connect(ctx); err != nil {
-			debuglog.Printf(ctx, "@automattic/vip:wp", "socket: connect_error stage=namespace")
-			eng.Close()
-			if first {
-				return Result{}, err
-			}
-			if werr := waitBackoff(ctx, &backoff, maxBackoff); werr != nil {
-				return Result{}, werr
-			}
-			continue
-		}
-		debuglog.Printf(ctx, "@automattic/vip:wp", "socket: connected")
-		first = false
-
-		res, clean := runOnce(ctx, opts, cli, ss, &offset, offset.Load())
-		eng.Close() // C2: ALWAYS close engine before next attempt or return
-
-		if ctx.Err() != nil {
-			return Result{}, ctx.Err()
-		}
-		if clean {
-			return res, nil
-		}
-
-		debuglog.Printf(ctx, "@automattic/vip:wp", "socket.io: reconnect")
-		// disconnected mid-command → wait then reconnect from offset
-		if werr := waitBackoff(ctx, &backoff, maxBackoff); werr != nil {
-			return Result{}, werr
-		}
-		backoff = time.Second // reset after a successful (even if interrupted) attempt
-	}
-}
-
-// waitBackoff sleeps for a jittered backoff duration, then doubles backoff up
-// to max. Returns ctx.Err() if the context is cancelled during the wait.
-func waitBackoff(ctx context.Context, backoff *time.Duration, max time.Duration) error {
-	jitter := time.Duration(float64(*backoff) * (0.5 + rand.Float64()*0.5))
-	select {
-	case <-time.After(jitter):
-	case <-ctx.Done():
-		return ctx.Err()
-	}
-	*backoff = time.Duration(math.Min(float64(*backoff*2), float64(max)))
-	return nil
-}
-
-// runOnce executes one attempt: registers handlers, launches the command,
-// pipes stdio, and waits for exit, stdout EOF, disconnect, or ctx cancel.
-//
-// Returns (result, true) on clean exit, or (Result{}, false) on disconnect
-// so the caller can reconnect. ctx cancel is treated as clean=true and the
-// caller checks ctx.Err() afterwards.
-func runOnce(ctx context.Context, opts Options, cli *Client, ss *StreamSocket, offset *atomic.Int64, resumeAt int64) (Result, bool) {
-	// exitCh is buffered: handlers run in goroutines and MUST NOT block on send.
-	exitCh := make(chan int, 4)
-	signalExit := func(code int) {
-		select {
-		case exitCh <- code:
-		default:
-		}
-	}
-
-	cli.On("unauthorized", func(args []any) {
-		fmt.Fprintln(opts.Stdout, "There was an error with the authentication:", errMessage(args))
-	})
-	cli.On("cancel", func(args []any) {
-		fmt.Fprintf(opts.Stdout, "Cancel received from server: %s\n", strArg(args))
-		signalExit(1)
-	})
-	cli.On("error", func(args []any) {
-		if strArg(args) == "Rate limit exceeded" {
-			fmt.Fprintln(opts.Stdout, color.RedString("\nError:"),
-				"Rate limit exceeded: Please wait a moment and try again.")
-			return
-		}
-		fmt.Fprintln(opts.Stdout, strArg(args))
-	})
-	cli.On("exit", func(args []any) {
-		code, msg := parseExit(args)
-		debuglog.Printf(ctx, "@automattic/vip:wp", "socket: exit code=%d", code)
-		if msg != "" {
-			fmt.Fprintln(opts.Stdout, msg)
-		}
-		signalExit(code)
-	})
-
-	disconnected := make(chan struct{}, 1)
-	cli.On("disconnect", func(args []any) {
-		select {
-		case disconnected <- struct{}{}:
-		default:
-		}
-	})
-
-	stdoutDone := make(chan struct{})
-	// runDone is closed when runOnce() is about to return, allowing the stdout
-	// goroutine to exit even if the server never closes the stdout stream.
-	runDone := make(chan struct{})
-	var runDoneOnce sync.Once
-	closeRunDone := func() { runDoneOnce.Do(func() { close(runDone) }) }
-
-	stdinStream := ss.CreateStream()
-	stdoutStream := ss.CreateStream()
-
-	cols := opts.Columns
-	if cols == 0 {
-		cols = nonTTYColumns
-	}
-	rows := opts.Rows
-	if rows == 0 {
-		rows = nonTTYRows
-	}
-	data := map[string]any{
-		"guid":       opts.GUID,
-		"inputToken": opts.InputToken,
-		"columns":    cols,
-		"rows":       rows,
-		"offset":     resumeAt,
-	}
-	_ = ss.Emit(ctx, "cmd", []any{data, stdinStream, stdoutStream}, nil)
-
-	// Pipe stdin → stdinStream.
-	// When stdinStream is aborted (C3) Write returns an error, io.Copy stops,
-	// and the goroutine exits — no leak.
-	go func() {
-		src := opts.Stdin
-		if opts.IsTTY && src != nil {
-			src = crToLF{src}
-		}
-		if src != nil {
-			_, _ = io.Copy(stdinStream, src)
-		}
-		_ = stdinStream.Close()
-	}()
-
-	// Background watcher: inject errRunDone into stdoutStream when runOnce is
-	// about to return, unblocking the stdout goroutine.
-	go func() {
-		select {
-		case <-runDone:
-			stdoutStream.abort(errRunDone)
-		case <-stdoutDone:
-		}
-	}()
-	// Pipe stdoutStream → opts.Stdout, tracking byte offset.
-	go func() {
-		buf := make([]byte, 32*1024)
-		for {
-			n, rerr := stdoutStream.Read(buf)
-			if n > 0 {
-				_, _ = opts.Stdout.Write(buf[:n])
-				if offset != nil {
-					offset.Add(int64(n))
-				}
-			}
-			if rerr != nil {
-				close(stdoutDone)
-				return
-			}
-		}
-	}()
-
-	defer closeRunDone()
-
-	var exitCode int
-	select {
-	case exitCode = <-exitCh:
-		// Drain stdout: signal runDone so the watcher goroutine aborts the
-		// stream, unblocking the stdout goroutine.
-		closeRunDone()
-		<-stdoutDone
-		return Result{ExitCode: exitCode}, true
-	case <-stdoutDone:
-		// stdout EOF before any exit event — the server always sends an 'exit'
-		// event after streaming completes (vip-wp.js:261). Wait briefly for it
-		// so a non-zero exit code is not silently lost.
-		select {
-		case exitCode = <-exitCh:
-		case <-time.After(2 * time.Second):
-		case <-ctx.Done():
-		}
-		return Result{ExitCode: exitCode}, true
-	case <-disconnected:
-		// Transport dropped mid-command.
-		// C3: abort all in-flight streams so blocked Read/Write goroutines exit.
-		// We do this here (not in NewStreamSocket) to ensure the disconnected
-		// channel is selected BEFORE stdoutDone can fire — preventing the race
-		// where stdoutDone fires first and runOnce returns clean=true incorrectly.
-		ss.abortAll(io.ErrUnexpectedEOF)
-		closeRunDone() // also aborts stdoutStream via watcher (redundant but safe)
-		return Result{}, false
-	case <-ctx.Done():
-		return Result{}, true // caller checks ctx.Err()
-	}
-}
-
-// run is the internal single-attempt function used by the unit tests (run_test.go).
-// The public API uses runOnce via Run. Kept for backward compatibility with tests.
-func run(ctx context.Context, opts Options, cli *Client, ss *StreamSocket, offset *atomic.Int64, resumeAt int64) (Result, error) {
-	res, clean := runOnce(ctx, opts, cli, ss, offset, resumeAt)
-	if !clean {
-		// disconnect treated as context cancellation for unit-test callers
-		return Result{}, ctx.Err()
-	}
-	if ctx.Err() != nil {
-		return Result{}, ctx.Err()
-	}
-	return res, nil
-}
-
-// crToLF replaces '\r' with '\n' (normalizeNewlineStream, vip-wp.js:51).
-type crToLF struct{ r io.Reader }
-
-func (c crToLF) Read(p []byte) (int, error) {
-	n, err := c.r.Read(p)
-	for i := 0; i < n; i++ {
-		if p[i] == '\r' {
-			p[i] = '\n'
-		}
-	}
-	return n, err
-}
-
-func bearerHeader(token string) http.Header {
-	return http.Header{"Authorization": {"Bearer " + token}}
-}
-
-func parseExit(args []any) (int, string) {
-	if len(args) == 0 {
-		return 0, ""
-	}
-	m, ok := args[0].(map[string]any)
-	if !ok {
-		return 0, ""
-	}
-	code := 0
-	if c, ok := m["exitCode"].(float64); ok {
-		code = int(c)
-	}
-	msg, _ := m["message"].(string)
-	return code, msg
-}
-
-func strArg(args []any) string {
-	if len(args) == 0 {
-		return ""
-	}
-	if s, ok := args[0].(string); ok {
-		return s
-	}
-	b, _ := json.Marshal(args[0])
-	return string(b)
-}
-
-func errMessage(args []any) string {
-	if len(args) == 0 {
-		return ""
-	}
-	if m, ok := args[0].(map[string]any); ok {
-		if s, ok := m["message"].(string); ok {
-			return s
-		}
-	}
-	return strArg(args)
-}

package/internal/wpstream/run_test.go

@@ -1,184 +0,0 @@
-package wpstream
-
-import (
-	"bytes"
-	"context"
-	"strings"
-	"sync/atomic"
-	"testing"
-	"time"
-)
-
-// loopbackPair holds both sides of an in-memory loopback for run tests.
-type loopbackPair struct {
-	// client (A) side — passed to run()
-	ssA  *StreamSocket
-	cliA *Client
-	// server (B) side — scripted in tests
-	ssB  *StreamSocket
-	cliB *Client
-}
-
-// newLoopbackPair builds a cross-linked loopback and returns both sides.
-func newLoopbackPair(t *testing.T) loopbackPair {
-	t.Helper()
-
-	ta := &pipeTransport{in: make(chan Packet, 64)}
-	tb := &pipeTransport{in: make(chan Packet, 64)}
-	ta.peer = tb
-	tb.peer = ta
-
-	cliA := NewClient(ta, "/wp-cli")
-	cliB := NewClient(tb, "/wp-cli")
-
-	ctx, cancel := context.WithCancel(context.Background())
-	t.Cleanup(cancel)
-
-	go cliA.readLoop(ctx)
-	go cliB.readLoop(ctx)
-
-	ssA := NewStreamSocket(ctx, cliA)
-	ssB := NewStreamSocket(ctx, cliB)
-
-	return loopbackPair{ssA: ssA, cliA: cliA, ssB: ssB, cliB: cliB}
-}
-
-// script is a function that scripts the "server" (B) side when it receives a "cmd" event.
-// args are the decoded arguments: [data, stdinStream, stdoutStream].
-type script func(ctx context.Context, args []any, cliB *Client, ssB *StreamSocket)
-
-// runWithScript runs run() on the A side while B executes the given script.
-// Returns the Result and the captured stdout (combined with opts.Stdout).
-func runWithScript(t *testing.T, s script) (Result, string) {
-	t.Helper()
-	t.Setenv("NO_COLOR", "1")
-
-	pair := newLoopbackPair(t)
-
-	var buf bytes.Buffer
-	opts := Options{
-		Stdin:  strings.NewReader(""), // empty stdin
-		Stdout: &buf,
-		Stderr: &buf,
-	}
-
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	t.Cleanup(cancel)
-
-	// Script the server side: register a "cmd" handler on ssB.
-	// The handler will be called in its own goroutine (per StreamSocket.On dispatch).
-	pair.ssB.On("cmd", func(args []any, ackID *int) {
-		s(ctx, args, pair.cliB, pair.ssB)
-	})
-
-	// Run the inner function with the A-side client + StreamSocket.
-	var offset atomic.Int64
-	resCh := make(chan Result, 1)
-	errCh := make(chan error, 1)
-	go func() {
-		res, err := run(ctx, opts, pair.cliA, pair.ssA, &offset, 0)
-		if err != nil {
-			errCh <- err
-			return
-		}
-		resCh <- res
-	}()
-
-	select {
-	case res := <-resCh:
-		return res, buf.String()
-	case err := <-errCh:
-		t.Fatalf("run() returned error: %v", err)
-		return Result{}, ""
-	case <-time.After(4 * time.Second):
-		t.Fatal("timeout waiting for run() to complete")
-		return Result{}, ""
-	}
-}
-
-// scriptExit returns a script that emits an exit event with the given code and message.
-func scriptExit(code int, message string) script {
-	return func(ctx context.Context, args []any, cliB *Client, ssB *StreamSocket) {
-		_ = cliB.Emit(ctx, "exit", []any{map[string]any{
-			"exitCode": float64(code),
-			"message":  message,
-		}}, nil)
-	}
-}
-
-// scriptCancel returns a script that emits a cancel event with the given message.
-func scriptCancel(message string) script {
-	return func(ctx context.Context, args []any, cliB *Client, ssB *StreamSocket) {
-		_ = cliB.Emit(ctx, "cancel", []any{message}, nil)
-	}
-}
-
-// scriptError returns a script that emits an error event with the given message.
-func scriptError(message string) script {
-	return func(ctx context.Context, args []any, cliB *Client, ssB *StreamSocket) {
-		_ = cliB.Emit(ctx, "error", []any{message}, nil)
-		// After error we still need to signal exit so run() terminates.
-		_ = cliB.Emit(ctx, "exit", []any{map[string]any{
-			"exitCode": float64(1),
-		}}, nil)
-	}
-}
-
-// scriptStdout returns a script that writes data to the stdout stream (args[2]),
-// closes it, then emits exit with the given code.
-func scriptStdout(data string, code int) script {
-	return func(ctx context.Context, args []any, cliB *Client, ssB *StreamSocket) {
-		// args: [data(map), stdinStream(*IOStream), stdoutStream(*IOStream)]
-		if len(args) < 3 {
-			return
-		}
-		stdoutStream, ok := args[2].(*IOStream)
-		if !ok {
-			return
-		}
-		// Write data then close the stdout stream.
-		_, _ = stdoutStream.Write([]byte(data))
-		_ = stdoutStream.Close()
-		// Emit exit to be deterministic.
-		_ = cliB.Emit(ctx, "exit", []any{map[string]any{
-			"exitCode": float64(code),
-		}}, nil)
-	}
-}
-
-func TestRunExitEvent(t *testing.T) {
-	res, out := runWithScript(t, scriptExit(5, "done"))
-	if res.ExitCode != 5 {
-		t.Errorf("exit = %d, want 5", res.ExitCode)
-	}
-	if !strings.Contains(out, "done") {
-		t.Errorf("message not printed: %q", out)
-	}
-}
-
-func TestRunCancelEvent(t *testing.T) {
-	res, out := runWithScript(t, scriptCancel("nope"))
-	if res.ExitCode != 1 {
-		t.Errorf("exit = %d, want 1", res.ExitCode)
-	}
-	if !strings.Contains(out, "Cancel received from server: nope") {
-		t.Errorf("out = %q", out)
-	}
-}
-
-func TestRunRateLimitError(t *testing.T) {
-	_, out := runWithScript(t, scriptError("Rate limit exceeded"))
-	if !strings.Contains(out, "Rate limit exceeded: Please wait a moment and try again.") {
-		t.Errorf("out = %q", out)
-	}
-}
-
-func TestRunStdoutStreamed(t *testing.T) {
-	res, out := runWithScript(t, scriptStdout("line1\nline2\n", 0))
-	if res.ExitCode != 0 {
-		t.Errorf("exit = %d", res.ExitCode)
-	}
-	if !strings.Contains(out, "line1") || !strings.Contains(out, "line2") {
-		t.Errorf("stdout not streamed: %q", out)
-	}
-}