labctl port-forward --restore

Author: iximiuz

api/playgrounds.go

@@ -25,6 +25,7 @@ type Playground struct {
 	InitTasks      map[string]InitTask `yaml:"initTasks,omitempty" json:"initTasks,omitempty"`
 	InitConditions InitConditions      `yaml:"initConditions,omitempty" json:"initConditions,omitempty"`
 	RegistryAuth   string              `yaml:"registryAuth,omitempty" json:"registryAuth,omitempty"`
+	PortForwards   []PortForward       `yaml:"portForwards,omitempty" json:"portForwards,omitempty"`
 
 	AccessControl PlaygroundAccessControl `yaml:"accessControl" json:"accessControl"`
 	UserAccess    PlaygroundUserAccess    `yaml:"userAccess,omitempty" json:"userAccess,omitempty"`
@@ -251,6 +252,7 @@ type PlaygroundSpec struct {
 	InitTasks      map[string]InitTask `yaml:"initTasks,omitempty" json:"initTasks,omitempty"`
 	InitConditions InitConditions      `yaml:"initConditions,omitempty" json:"initConditions,omitempty"`
 	RegistryAuth   string              `yaml:"registryAuth,omitempty" json:"registryAuth,omitempty"`
+	PortForwards   []PortForward       `yaml:"portForwards,omitempty" json:"portForwards,omitempty"`
 
 	// Deprecated: Use PlaygroundAccessControl instead
 	Access *PlaygroundAccess `yaml:"access,omitempty" json:"access,omitempty"`
@@ -289,6 +291,7 @@ type CreatePlaygroundRequest struct {
 	InitTasks      map[string]InitTask `yaml:"initTasks" json:"initTasks"`
 	InitConditions InitConditions      `yaml:"initConditions" json:"initConditions"`
 	RegistryAuth   string              `yaml:"registryAuth" json:"registryAuth"`
+	PortForwards   []PortForward       `yaml:"portForwards,omitempty" json:"portForwards,omitempty"`
 
 	AccessControl PlaygroundAccessControl `yaml:"accessControl" json:"accessControl"`
 }
@@ -315,6 +318,7 @@ type UpdatePlaygroundRequest struct {
 	InitTasks      map[string]InitTask `yaml:"initTasks" json:"initTasks"`
 	InitConditions InitConditions      `yaml:"initConditions" json:"initConditions"`
 	RegistryAuth   string              `yaml:"registryAuth" json:"registryAuth"`
+	PortForwards   []PortForward       `yaml:"portForwards,omitempty" json:"portForwards,omitempty"`
 
 	AccessControl PlaygroundAccessControl `yaml:"accessControl" json:"accessControl"`
 }

api/port_forwards.go

@@ -0,0 +1,35 @@
+package api
+
+import (
+	"context"
+	"strconv"
+)
+
+type PortForward struct {
+	Kind       string `json:"kind"`
+	Machine    string `json:"machine"`
+	LocalHost  string `json:"localHost,omitempty"`
+	LocalPort  int    `json:"localPort,omitempty"`
+	RemotePort int    `json:"remotePort,omitempty"`
+	RemoteHost string `json:"remoteHost,omitempty"`
+}
+
+func (c *Client) ListPortForwards(ctx context.Context, playID string) ([]*PortForward, error) {
+	var resp []*PortForward
+	return resp, c.GetInto(ctx, "/plays/"+playID+"/port-forwards", nil, nil, &resp)
+}
+
+func (c *Client) AddPortForward(ctx context.Context, playID string, pf PortForward) (*PortForward, error) {
+	body, err := toJSONBody(pf)
+	if err != nil {
+		return nil, err
+	}
+
+	var resp PortForward
+	return &resp, c.PostInto(ctx, "/plays/"+playID+"/port-forwards", nil, nil, body, &resp)
+}
+
+func (c *Client) RemovePortForward(ctx context.Context, playID string, index int) error {
+	_, err := c.Delete(ctx, "/plays/"+playID+"/port-forwards/"+strconv.Itoa(index), nil, nil)
+	return err
+}

cmd/playground/manifest.go

@@ -55,6 +55,7 @@ func runManifest(ctx context.Context, cli labcli.CLI, opts *manifestOptions) err
 			InitTasks:      playground.InitTasks,
 			InitConditions: playground.InitConditions,
 			RegistryAuth:   playground.RegistryAuth,
+			PortForwards:   playground.PortForwards,
 			AccessControl:  playground.AccessControl,
 		},
 	}

cmd/playground/restart.go

@@ -14,6 +14,7 @@ import (
 	"github.com/iximiuz/labctl/cmd/ssh"
 	"github.com/iximiuz/labctl/cmd/sshproxy"
 	"github.com/iximiuz/labctl/internal/labcli"
+	"github.com/iximiuz/labctl/internal/portforward"
 )
 
 const restartCommandTimeout = 5 * time.Minute
@@ -30,6 +31,8 @@ type restartOptions struct {
 
 	forwardAgent bool
 
+	restorePortForwards bool
+
 	quiet bool
 }
 
@@ -105,6 +108,12 @@ func newRestartCommand(cli labcli.CLI) *cobra.Command {
 		false,
 		`Do not print any diagnostic messages`,
 	)
+	flags.BoolVar(
+		&opts.restorePortForwards,
+		"restore-port-forwards",
+		false,
+		`Automatically restore all forwarded ports from previous session`,
+	)
 
 	return cmd
 }
@@ -168,6 +177,17 @@ func runRestartPlayground(ctx context.Context, cli labcli.CLI, opts *restartOpti
 		}
 	}
 
+	// Start port forwarding if requested.
+	// If combined with --ide or --ssh, run in background; otherwise block.
+	var portForwardErrCh <-chan error
+	if opts.restorePortForwards {
+		var err error
+		portForwardErrCh, err = portforward.RestoreSavedForwards(ctx, cli.Client(), play.ID, cli)
+		if err != nil {
+			return err
+		}
+	}
+
 	if opts.ide != "" {
 		return sshproxy.RunSSHProxy(ctx, cli, &sshproxy.Options{
 			PlayID:  play.ID,
@@ -197,5 +217,11 @@ func runRestartPlayground(ctx context.Context, cli labcli.CLI, opts *restartOpti
 	}
 
 	cli.PrintOut("%s\n", play.ID)
+
+	// If only --port-forward was provided (no --ide or --ssh), wait for it
+	if portForwardErrCh != nil {
+		return <-portForwardErrCh
+	}
+
 	return nil
 }

cmd/portforward/portforward.go

@@ -3,6 +3,7 @@ package portforward
 import (
 	"context"
 	"fmt"
+	"strconv"
 
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
@@ -11,6 +12,16 @@ import (
 	"github.com/iximiuz/labctl/internal/portforward"
 )
 
+// runRestorePortForwards restores saved port forwards and blocks until done.
+func runRestorePortForwards(ctx context.Context, cli labcli.CLI, opts *options) error {
+	resultCh, err := portforward.RestoreSavedForwards(ctx, cli.Client(), opts.playID, cli)
+	if err != nil {
+		return err
+	}
+
+	return <-resultCh
+}
+
 type options struct {
 	playID  string
 	machine string
@@ -21,6 +32,11 @@ type options struct {
 	remotes []string
 
 	quiet bool
+
+	// New flags
+	list    bool
+	restore bool
+	remove  int
 }
 
 // Local port forwarding's possible modes (kinda sorta as in ssh -L):
@@ -32,18 +48,48 @@ type options struct {
 //   - LOCAL_HOST:LOCAL_PORT:REMOTE_HOST:REMOTE_PORT  # the most explicit form
 
 func NewCommand(cli labcli.CLI) *cobra.Command {
-	var opts options
+	opts := options{
+		remove: -1, // -1 means not set
+	}
 
 	cmd := &cobra.Command{
-		Use:   "port-forward <playground> [-m machine] -L [LOCAL:]REMOTE [-L ...] | -R [REMOTE:]:LOCAL [-R ...]",
+		Use:   "port-forward <playground> [-m machine] -L [LOCAL:]REMOTE [-L ...] | --list | --restore | --remove <index>",
 		Short: `Forward one or more local or remote ports to a running playground`,
-		Long: `While the implementation for sure differs, the behavior and semantic of the command
+		Long: `Forward one or more local or remote ports to a running playground.
+
+While the implementation differs significantly, the behavior and semantic of the command
 are meant to be similar to SSH local (-L) and remote (-R) port forwarding. The word "local" always
-refers to the labctl side. The word "remote" always refers to the target playground side.`,
+refers to the labctl side. The word "remote" always refers to the target playground side.
+
+The command also supports managing "saved" port forwards:
+  --list     List all "should be forwarded" ports for the playground
+  --restore  Forward all "should be forwarded" ports (handy after a persistent playground restart)
+  --remove   Remove a "should be forwarded" port from the playground's config by its index (0-based)
+
+When using -L|-R flags, port forwards are automatically saved to the playground's config for later restoration.`,
 		Args: cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
+			opts.playID = args[0]
+			cli.SetQuiet(opts.quiet)
+
+			// Handle list mode
+			if opts.list {
+				return labcli.WrapStatusError(runListPortForwards(cmd.Context(), cli, &opts))
+			}
+
+			// Handle remove mode
+			if opts.remove >= 0 {
+				return labcli.WrapStatusError(runRemovePortForward(cmd.Context(), cli, &opts))
+			}
+
+			// Handle restore mode
+			if opts.restore {
+				return labcli.WrapStatusError(runRestorePortForwards(cmd.Context(), cli, &opts))
+			}
+
+			// Regular port forwarding mode
 			if len(opts.locals)+len(opts.remotes) == 0 {
-				return labcli.NewStatusError(1, "at least one -L or -R flag must be provided")
+				return labcli.NewStatusError(1, "at least one -L or -R flag must be provided (or use --list, --restore, --remove)")
 			}
 			if len(opts.remotes) > 0 {
 				// TODO: Implement me!
@@ -58,10 +104,6 @@ refers to the labctl side. The word "remote" always refers to the target playgro
 				opts.localsParsed = append(opts.localsParsed, parsed)
 			}
 
-			cli.SetQuiet(opts.quiet)
-
-			opts.playID = args[0]
-
 			return labcli.WrapStatusError(runPortForward(cmd.Context(), cli, &opts))
 		},
 	}
@@ -97,9 +139,68 @@ refers to the labctl side. The word "remote" always refers to the target playgro
 		`Suppress verbose output`,
 	)
 
+	flags.BoolVar(&opts.list, "list", false, `List saved port forwards ("saved" means "should be forwarded")`)
+	flags.BoolVar(&opts.restore, "restore", false, `Forward all "should be forwarded" ports for the playground`)
+	flags.IntVar(&opts.remove, "remove", -1, `Remove a "should be forwarded" port from the playground's config by index (0-based)`)
+
 	return cmd
 }
 
+func runListPortForwards(ctx context.Context, cli labcli.CLI, opts *options) error {
+	forwards, err := cli.Client().ListPortForwards(ctx, opts.playID)
+	if err != nil {
+		return fmt.Errorf("couldn't list port forwards: %w", err)
+	}
+
+	if len(forwards) == 0 {
+		cli.PrintAux("No saved port forwards found.\n")
+		return nil
+	}
+
+	cli.PrintAux("Saved port forwards:\n")
+	for i, pf := range forwards {
+		localPart := ""
+		if pf.LocalHost != "" || pf.LocalPort > 0 {
+			if pf.LocalHost != "" {
+				localPart = pf.LocalHost
+			}
+			if pf.LocalPort > 0 {
+				if localPart != "" {
+					localPart += ":"
+				}
+				localPart += strconv.Itoa(pf.LocalPort)
+			}
+			localPart += " -> "
+		}
+
+		remotePart := ""
+		if pf.RemoteHost != "" {
+			remotePart = pf.RemoteHost + ":"
+		}
+		if pf.RemotePort > 0 {
+			remotePart += strconv.Itoa(pf.RemotePort)
+		}
+
+		kindLabel := pf.Kind
+		if kindLabel == "" {
+			kindLabel = "local"
+		}
+
+		cli.PrintAux("  [%d] %s (%s): %s%s\n", i, pf.Machine, kindLabel, localPart, remotePart)
+	}
+
+	return nil
+}
+
+func runRemovePortForward(ctx context.Context, cli labcli.CLI, opts *options) error {
+	if err := cli.Client().RemovePortForward(ctx, opts.playID, opts.remove); err != nil {
+		return fmt.Errorf("couldn't remove port forward: %w", err)
+	}
+
+	cli.PrintAux("Port forward at index %d removed.\n", opts.remove)
+	return nil
+}
+
 func runPortForward(ctx context.Context, cli labcli.CLI, opts *options) error {
 	p, err := cli.Client().GetPlay(ctx, opts.playID)
 	if err != nil {
@@ -114,11 +215,20 @@ func runPortForward(ctx context.Context, cli labcli.CLI, opts *options) error {
 		}
 	}
 
+	// Save port forwards to play's config
+	for _, spec := range opts.localsParsed {
+		pf, err := spec.ToPortForward(opts.machine)
+		if err != nil {
+			return fmt.Errorf("couldn't convert port forwarding spec to API port forward model: %w", err)
+		}
+		if _, err := cli.Client().AddPortForward(ctx, p.ID, *pf); err != nil {
+			cli.PrintErr("Warning: couldn't save port forward: %v\n", err)
+		}
+	}
+
 	tunnel, err := portforward.StartTunnel(ctx, cli.Client(), portforward.TunnelOptions{
-		PlayID:    p.ID,
-		FactoryID: p.FactoryID(),
-		Machine:   opts.machine,
-		PlaysDir:  cli.Config().PlaysDir,
+		PlayID:  p.ID,
+		Machine: opts.machine,
 	})
 	if err != nil {
 		return fmt.Errorf("couldn't start tunnel: %w", err)

cmd/ssh/ssh.go

@@ -130,9 +130,7 @@ func StartSSHSession(
 ) (*ssh.Session, <-chan error, error) {
 	tunnel, err := portforward.StartTunnel(ctx, cli.Client(), portforward.TunnelOptions{
 		PlayID:          play.ID,
-		FactoryID:       play.FactoryID(),
 		Machine:         machine,
-		PlaysDir:        cli.Config().PlaysDir,
 		SSHUser:         user,
 		SSHIdentityFile: cli.Config().SSHIdentityFile,
 	})

cmd/sshproxy/sshproxy.go

@@ -128,9 +128,7 @@ func RunSSHProxy(ctx context.Context, cli labcli.CLI, opts *Options) error {
 
 	tunnel, err := portforward.StartTunnel(ctx, cli.Client(), portforward.TunnelOptions{
 		PlayID:          p.ID,
-		FactoryID:       p.FactoryID(),
 		Machine:         opts.Machine,
-		PlaysDir:        cli.Config().PlaysDir,
 		SSHUser:         opts.User,
 		SSHIdentityFile: cli.Config().SSHIdentityFile,
 	})

internal/labcli/cliutil.go

@@ -20,6 +20,12 @@ type Streams interface {
 	ErrorStream() io.Writer
 }
 
+type Outputer interface {
+	PrintOut(string, ...any)
+	PrintErr(string, ...any)
+	PrintAux(string, ...any)
+}
+
 type CLI interface {
 	Streams
 
@@ -61,7 +67,10 @@ type cli struct {
 	version string
 }
 
-var _ CLI = &cli{}
+var (
+	_ CLI      = &cli{}
+	_ Outputer = &cli{}
+)
 
 func NewCLI(cin io.ReadCloser, cout io.Writer, cerr io.Writer, version string) CLI {
 	return &cli{