Skip to content

Commit b9b7d45

Browse files
acmoreacmore
committed
feat(cli): improve attach flow and add ssh sidecar mode
1 parent 2973163 commit b9b7d45

11 files changed

Lines changed: 417 additions & 32 deletions

File tree

‎internal/cli/down.go‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@ package cli
22

33
import (
44
"fmt"
5+
"os"
6+
"path/filepath"
7+
"strings"
58

69
"github.com/spf13/cobra"
710
)
@@ -47,6 +50,7 @@ func newDownCmd(opts *Options) *cobra.Command {
4750
return fmt.Errorf("delete workspace pvc: %w", err)
4851
}
4952
}
53+
_ = removeSSHConfigEntry(sshHostAlias(sn))
5054
fmt.Fprintf(cmd.OutOrStdout(), "Session stopped: %s\n", sn)
5155
if !deletePVC && cfg.Spec.Workspace.PVC.ClaimName == "" {
5256
fmt.Fprintf(cmd.OutOrStdout(), "Workspace PVC retained: %s (use --delete-pvc to remove)\n", pvcName(cfg, sn))
@@ -58,3 +62,22 @@ func newDownCmd(opts *Options) *cobra.Command {
5862
cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Preview actions without deleting resources")
5963
return cmd
6064
}
65+
66+
func removeSSHConfigEntry(hostAlias string) error {
67+
home, err := os.UserHomeDir()
68+
if err != nil {
69+
return err
70+
}
71+
configPath := filepath.Join(home, ".ssh", "config")
72+
existing, err := os.ReadFile(configPath)
73+
if err != nil {
74+
return nil
75+
}
76+
begin := "# BEGIN OKDEV " + hostAlias
77+
end := "# END OKDEV " + hostAlias
78+
updated := strings.TrimSpace(stripManagedSSHBlock(string(existing), begin, end))
79+
if updated == "" {
80+
return os.WriteFile(configPath, []byte(""), 0o600)
81+
}
82+
return os.WriteFile(configPath, []byte(updated+"\n"), 0o600)
83+
}

‎internal/cli/portforward_helper.go‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ func startManagedPortForward(opts *Options, namespace, pod string, forwards []st
1616
return startManagedPortForwardWithClient(newKubeClient(opts), namespace, pod, forwards)
1717
}
1818

19+
func startManagedPortForwardNoProbe(opts *Options, namespace, pod string, forwards []string) (context.CancelFunc, error) {
20+
return startManagedPortForwardNoProbeWithClient(newKubeClient(opts), namespace, pod, forwards)
21+
}
22+
1923
func startManagedPortForwardWithClient(k interface {
2024
PortForward(context.Context, string, string, []string, io.Writer, io.Writer) error
2125
}, namespace, pod string, forwards []string) (context.CancelFunc, error) {
@@ -58,6 +62,37 @@ func startManagedPortForwardWithClient(k interface {
5862
}
5963
}
6064

65+
func startManagedPortForwardNoProbeWithClient(k interface {
66+
PortForward(context.Context, string, string, []string, io.Writer, io.Writer) error
67+
}, namespace, pod string, forwards []string) (context.CancelFunc, error) {
68+
ctx, cancel := context.WithCancel(context.Background())
69+
errCh := make(chan error, 1)
70+
doneCh := make(chan struct{})
71+
go func() {
72+
defer close(doneCh)
73+
errCh <- k.PortForward(ctx, namespace, pod, forwards, io.Discard, io.Discard)
74+
}()
75+
cancelAndWait := func() {
76+
cancel()
77+
select {
78+
case <-doneCh:
79+
case <-time.After(portForwardShutdownWait):
80+
}
81+
}
82+
timer := time.NewTimer(500 * time.Millisecond)
83+
defer timer.Stop()
84+
select {
85+
case err := <-errCh:
86+
cancelAndWait()
87+
if err != nil {
88+
return nil, err
89+
}
90+
return nil, fmt.Errorf("port-forward exited before ready")
91+
case <-timer.C:
92+
return cancelAndWait, nil
93+
}
94+
}
95+
6196
func localPortsFromForwards(forwards []string) []int {
6297
ports := make([]int, 0, len(forwards))
6398
for _, f := range forwards {

‎internal/cli/ssh.go‎

Lines changed: 86 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ func newSSHCmd(opts *Options) *cobra.Command {
2525

2626
cmd := &cobra.Command{
2727
Use: "ssh",
28-
Short: "Connect to session pod over SSH (requires sshd in container)",
28+
Short: "Connect to session pod over SSH (dev container or okdev SSH sidecar)",
2929
RunE: func(cmd *cobra.Command, args []string) error {
3030
cfg, ns, err := loadConfigAndNamespace(opts)
3131
if err != nil {
@@ -59,7 +59,7 @@ func newSSHCmd(opts *Options) *cobra.Command {
5959
}
6060

6161
if setupKey {
62-
if err := ensureSSHKeyOnPod(opts, ns, podName(sn), keyPath); err != nil {
62+
if err := ensureSSHKeyOnPod(opts, cfg, ns, podName(sn), keyPath); err != nil {
6363
return err
6464
}
6565
}
@@ -70,16 +70,21 @@ func newSSHCmd(opts *Options) *cobra.Command {
7070
}
7171
defer cancelPF()
7272

73+
sshHost := sshHostAlias(sn)
74+
if cfgErr := ensureSSHConfigEntry(sshHost, user, localPort, keyPath); cfgErr != nil {
75+
fmt.Fprintf(cmd.ErrOrStderr(), "warning: failed to update ~/.ssh/config: %v\n", cfgErr)
76+
}
77+
7378
sshArgs := []string{
7479
"-p", fmt.Sprintf("%d", localPort),
7580
"-o", "StrictHostKeyChecking=no",
7681
"-o", "UserKnownHostsFile=/dev/null",
7782
"-i", keyPath,
7883
}
7984
if cmdStr != "" {
80-
sshArgs = append(sshArgs, fmt.Sprintf("%s@127.0.0.1", user), cmdStr)
85+
sshArgs = append(sshArgs, sshHost, cmdStr)
8186
} else {
82-
sshArgs = append(sshArgs, fmt.Sprintf("%s@127.0.0.1", user))
87+
sshArgs = append(sshArgs, sshHost)
8388
}
8489
sshCmd := exec.Command("ssh", sshArgs...)
8590
sshCmd.Stdin = os.Stdin
@@ -101,7 +106,7 @@ func newSSHCmd(opts *Options) *cobra.Command {
101106
return cmd
102107
}
103108

104-
func ensureSSHKeyOnPod(opts *Options, namespace, pod, keyPath string) error {
109+
func ensureSSHKeyOnPod(opts *Options, cfg *config.DevEnvironment, namespace, pod, keyPath string) error {
105110
if err := ensureCommand("ssh-keygen"); err != nil {
106111
return err
107112
}
@@ -127,7 +132,13 @@ func ensureSSHKeyOnPod(opts *Options, namespace, pod, keyPath string) error {
127132
script := fmt.Sprintf("mkdir -p ~/.ssh && chmod 700 ~/.ssh && touch ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && (grep -qxF %s ~/.ssh/authorized_keys || echo %s >> ~/.ssh/authorized_keys)", syncengine.ShellEscape(pubKey), syncengine.ShellEscape(pubKey))
128133
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
129134
defer cancel()
130-
_, err = newKubeClient(opts).ExecSh(ctx, namespace, pod, script)
135+
k := newKubeClient(opts)
136+
container := sshTargetContainer(cfg)
137+
if container == "" {
138+
_, err = k.ExecSh(ctx, namespace, pod, script)
139+
} else {
140+
_, err = k.ExecShInContainer(ctx, namespace, pod, container, script)
141+
}
131142
if err != nil {
132143
return fmt.Errorf("install ssh key in pod: %w", err)
133144
}
@@ -149,3 +160,72 @@ func defaultSSHKeyPath(cfg *config.DevEnvironment) (string, error) {
149160
}
150161
return filepath.Join(home, ".ssh", "okdev_ed25519"), nil
151162
}
163+
164+
func sshTargetContainer(cfg *config.DevEnvironment) string {
165+
if cfg == nil {
166+
return ""
167+
}
168+
if strings.EqualFold(strings.TrimSpace(cfg.Spec.SSH.Mode), "sidecar") {
169+
return "okdev-ssh"
170+
}
171+
return ""
172+
}
173+
174+
func sshHostAlias(sessionName string) string {
175+
return "okdev-" + sessionName
176+
}
177+
178+
func ensureSSHConfigEntry(hostAlias, user string, localPort int, keyPath string) error {
179+
home, err := os.UserHomeDir()
180+
if err != nil {
181+
return err
182+
}
183+
sshDir := filepath.Join(home, ".ssh")
184+
if err := os.MkdirAll(sshDir, 0o700); err != nil {
185+
return err
186+
}
187+
configPath := filepath.Join(sshDir, "config")
188+
existing, _ := os.ReadFile(configPath)
189+
190+
begin := "# BEGIN OKDEV " + hostAlias
191+
end := "# END OKDEV " + hostAlias
192+
blockLines := []string{
193+
begin,
194+
"Host " + hostAlias,
195+
" HostName 127.0.0.1",
196+
fmt.Sprintf(" Port %d", localPort),
197+
" User " + user,
198+
" IdentityFile " + keyPath,
199+
" StrictHostKeyChecking no",
200+
" UserKnownHostsFile /dev/null",
201+
end,
202+
}
203+
block := strings.Join(blockLines, "\n") + "\n"
204+
updated := stripManagedSSHBlock(string(existing), begin, end) + "\n" + block
205+
updated = strings.TrimSpace(updated) + "\n"
206+
return os.WriteFile(configPath, []byte(updated), 0o600)
207+
}
208+
209+
func stripManagedSSHBlock(content, begin, end string) string {
210+
lines := strings.Split(content, "\n")
211+
out := make([]string, 0, len(lines))
212+
inBlock := false
213+
for _, line := range lines {
214+
t := strings.TrimSpace(line)
215+
if t == begin {
216+
inBlock = true
217+
continue
218+
}
219+
if inBlock {
220+
if t == end {
221+
inBlock = false
222+
}
223+
continue
224+
}
225+
out = append(out, line)
226+
}
227+
for len(out) > 0 && strings.TrimSpace(out[len(out)-1]) == "" {
228+
out = out[:len(out)-1]
229+
}
230+
return strings.Join(out, "\n")
231+
}

‎internal/cli/up.go‎

Lines changed: 54 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,14 @@ import (
77
"log/slog"
88
"os"
99
"strconv"
10+
"strings"
1011
"time"
1112

1213
"github.com/acmore/okdev/internal/config"
1314
"github.com/acmore/okdev/internal/kube"
1415
"github.com/acmore/okdev/internal/session"
1516
"github.com/spf13/cobra"
17+
apierrors "k8s.io/apimachinery/pkg/api/errors"
1618
)
1719

1820
func newUpCmd(opts *Options) *cobra.Command {
@@ -61,6 +63,16 @@ func newUpCmd(opts *Options) *cobra.Command {
6163
if warnErr := warnIfConfigNewerThanSession(opts, k, ns, sn, pod, cmd.ErrOrStderr()); warnErr != nil {
6264
slog.Debug("skip config drift warning", "error", warnErr)
6365
}
66+
existingRunningReady := false
67+
{
68+
ctxCheck, cancelCheck := context.WithTimeout(context.Background(), 10*time.Second)
69+
defer cancelCheck()
70+
if ps, perr := k.GetPodSummary(ctxCheck, ns, pod); perr == nil {
71+
existingRunningReady = strings.EqualFold(ps.Phase, "Running") && podReadyFromSummary(ps.Ready)
72+
} else if !apierrors.IsNotFound(perr) {
73+
return perr
74+
}
75+
}
6476

6577
ctx, cancel := defaultContext()
6678
defer cancel()
@@ -75,36 +87,40 @@ func newUpCmd(opts *Options) *cobra.Command {
7587
}
7688
}
7789

78-
preparedSpec, err := kube.PreparePodSpec(
90+
preparedSpec, err := kube.PreparePodSpecWithSSH(
7991
cfg.Spec.PodTemplate.Spec,
8092
pvc,
8193
cfg.Spec.Workspace.MountPath,
8294
cfg.Spec.Sync.Engine == "syncthing",
8395
cfg.Spec.Sync.Syncthing.Image,
96+
strings.EqualFold(cfg.Spec.SSH.Mode, "sidecar"),
97+
cfg.Spec.SSH.SidecarImage,
8498
)
8599
if err != nil {
86100
return err
87101
}
88102

89-
podManifest, err := kube.BuildPodManifest(ns, pod, labels, annotations, preparedSpec)
90-
if err != nil {
91-
return err
92-
}
93-
if err := k.Apply(ctx, ns, podManifest); err != nil {
94-
return err
95-
}
96-
progressPrinter := func(p kube.PodReadinessProgress) {
97-
fmt.Fprintf(cmd.OutOrStdout(), "Waiting for pod/%s: phase=%s ready=%d/%d reason=%s\n", pod, p.Phase, p.ReadyContainers, p.TotalContainers, p.Reason)
98-
}
99-
if err := k.WaitReadyWithProgress(ctx, ns, pod, waitTimeout, progressPrinter); err != nil {
100-
hints := fmt.Sprintf("next steps:\n- run `okdev status --session %s`\n- run `kubectl -n %s describe pod %s`", sn, ns, pod)
101-
diag, derr := k.DescribePod(ctx, ns, pod)
102-
if derr == nil {
103-
fmt.Fprintf(cmd.ErrOrStderr(), "pod diagnostics:\n%s\n\n%s\n", diag, hints)
103+
if !existingRunningReady {
104+
podManifest, err := kube.BuildPodManifest(ns, pod, labels, annotations, preparedSpec)
105+
if err != nil {
106+
return err
107+
}
108+
if err := k.Apply(ctx, ns, podManifest); err != nil {
109+
return err
110+
}
111+
progressPrinter := func(p kube.PodReadinessProgress) {
112+
fmt.Fprintf(cmd.OutOrStdout(), "Waiting for pod/%s: phase=%s ready=%d/%d reason=%s\n", pod, p.Phase, p.ReadyContainers, p.TotalContainers, p.Reason)
113+
}
114+
if err := k.WaitReadyWithProgress(ctx, ns, pod, waitTimeout, progressPrinter); err != nil {
115+
hints := fmt.Sprintf("next steps:\n- run `okdev status --session %s`\n- run `kubectl -n %s describe pod %s`", sn, ns, pod)
116+
diag, derr := k.DescribePod(ctx, ns, pod)
117+
if derr == nil {
118+
fmt.Fprintf(cmd.ErrOrStderr(), "pod diagnostics:\n%s\n\n%s\n", diag, hints)
119+
return fmt.Errorf("wait for pod/%s readiness failed: %w", pod, err)
120+
}
121+
fmt.Fprintln(cmd.ErrOrStderr(), hints)
104122
return fmt.Errorf("wait for pod/%s readiness failed: %w", pod, err)
105123
}
106-
fmt.Fprintln(cmd.ErrOrStderr(), hints)
107-
return fmt.Errorf("wait for pod/%s readiness failed: %w", pod, err)
108124
}
109125

110126
if err := session.SaveActiveSession(sn); err != nil {
@@ -132,29 +148,34 @@ func newUpCmd(opts *Options) *cobra.Command {
132148
forwards = append(forwards, strconv.Itoa(p.Local)+":"+strconv.Itoa(p.Remote))
133149
}
134150
if len(forwards) > 0 {
135-
cancelPF, err := startManagedPortForwardWithClient(k, ns, pod, forwards)
151+
cancelPF, err := startManagedPortForwardNoProbeWithClient(k, ns, pod, forwards)
136152
if err != nil {
137-
return fmt.Errorf("start background port-forward: %w", err)
153+
fmt.Fprintf(cmd.ErrOrStderr(), "warning: failed to start background port-forward: %v\n", err)
154+
} else {
155+
stopBackgrounds = append(stopBackgrounds, cancelPF)
156+
fmt.Fprintf(cmd.OutOrStdout(), "Background port-forward active: %v\n", forwards)
138157
}
139-
stopBackgrounds = append(stopBackgrounds, cancelPF)
140-
fmt.Fprintf(cmd.OutOrStdout(), "Background port-forward active: %v\n", forwards)
141158
}
142159
}
143160
if cfg.Spec.SSH.LocalPort > 0 && cfg.Spec.SSH.RemotePort > 0 {
144161
keyPath, keyErr := defaultSSHKeyPath(cfg)
145162
if keyErr != nil {
146163
fmt.Fprintf(cmd.ErrOrStderr(), "warning: failed to resolve SSH key path: %v\n", keyErr)
147164
} else {
148-
if err := ensureSSHKeyOnPod(opts, ns, pod, keyPath); err != nil {
165+
if err := ensureSSHKeyOnPod(opts, cfg, ns, pod, keyPath); err != nil {
149166
fmt.Fprintf(cmd.ErrOrStderr(), "warning: failed to setup SSH key in pod: %v\n", err)
150167
} else {
151168
sshForward := []string{strconv.Itoa(cfg.Spec.SSH.LocalPort) + ":" + strconv.Itoa(cfg.Spec.SSH.RemotePort)}
152-
cancelSSH, err := startManagedPortForwardWithClient(k, ns, pod, sshForward)
169+
cancelSSH, err := startManagedPortForwardNoProbeWithClient(k, ns, pod, sshForward)
153170
if err != nil {
154171
fmt.Fprintf(cmd.ErrOrStderr(), "warning: failed to start SSH port-forward %v: %v\n", sshForward, err)
155172
} else {
156173
stopBackgrounds = append(stopBackgrounds, cancelSSH)
157-
fmt.Fprintf(cmd.OutOrStdout(), "Background SSH tunnel active: ssh -p %d -i %s %s@127.0.0.1\n", cfg.Spec.SSH.LocalPort, keyPath, cfg.Spec.SSH.User)
174+
alias := sshHostAlias(sn)
175+
if cfgErr := ensureSSHConfigEntry(alias, cfg.Spec.SSH.User, cfg.Spec.SSH.LocalPort, keyPath); cfgErr != nil {
176+
fmt.Fprintf(cmd.ErrOrStderr(), "warning: failed to update ~/.ssh/config: %v\n", cfgErr)
177+
}
178+
fmt.Fprintf(cmd.OutOrStdout(), "Background SSH tunnel active: ssh %s\n", alias)
158179
}
159180
}
160181
}
@@ -206,3 +227,11 @@ func warnIfConfigNewerThanSession(opts *Options, k *kube.Client, namespace, sess
206227
}
207228
return nil
208229
}
230+
231+
func podReadyFromSummary(ready string) bool {
232+
parts := strings.Split(strings.TrimSpace(ready), "/")
233+
if len(parts) != 2 {
234+
return false
235+
}
236+
return parts[0] == parts[1] && parts[0] != "0"
237+
}

0 commit comments

Comments
 (0)