Reconnect ssh to reload changes after updating environment file (#649)

* bug: reconnect ssh upon environment change

This is a workaround. UpdateEnvironment on rig's os/linux.go writes
the environment to /etc/environment and then exports the same variables
using 'export' command.

Unfortunately this is not enough for the environment to be preserved
across multiple ssh sessions. We need to write the environment and then
reopen the ssh connection. Go's ssh client.Setenv() depends on ssh server
configuration (sshd only accepts LC_* variables by default).

This problem happens, for example, when we need to deploy k0s behing a
proxy. We set the necessary environment variables in k0sctl.yaml but the
command fails when attempting to use `apt` to install `curl`.

As `HTTP_PROXY`, and any other variable, are not set when we run the apt
command it fail as it can't reach the remote servers.

Signed-off-by: Ricardo Maraschini <[email protected]>

* chore: add retry logic to reconnection

* chore: logging before the disconnection

---------

Signed-off-by: Ricardo Maraschini <[email protected]>

phase/prepare_hosts.go

@@ -1,9 +1,15 @@
 package phase
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"strings"
+	"time"
 
 	"github.com/k0sproject/k0sctl/pkg/apis/k0sctl.k0sproject.io/v1beta1/cluster"
+	"github.com/k0sproject/k0sctl/pkg/retry"
+	"github.com/k0sproject/rig"
 	"github.com/k0sproject/rig/os"
 	"github.com/k0sproject/version"
 	log "github.com/sirupsen/logrus"
@@ -30,6 +36,34 @@ type prepare interface {
 	Prepare(os.Host) error
 }
 
+// updateEnvironment updates the environment variables on the host and reconnects to
+// it if necessary.
+func (p *PrepareHosts) updateEnvironment(h *cluster.Host) error {
+	if err := h.Configurer.UpdateEnvironment(h, h.Environment); err != nil {
+		return err
+	}
+	if h.Connection.Protocol() != "SSH" {
+		return nil
+	}
+	// XXX: this is a workaround. UpdateEnvironment on rig's os/linux.go writes
+	// the environment to /etc/environment and then exports the same variables
+	// using 'export' command. This is not enough for the environment to be
+	// preserved across multiple ssh sessions. We need to write the environment
+	// and then reopen the ssh session. Go's ssh client.Setenv() depends on ssh
+	// server configuration (sshd only accepts LC_* variables by default).
+	log.Infof("%s: reconnecting to apply new environment", h)
+	h.Disconnect()
+	return retry.Timeout(context.TODO(), 10*time.Minute, func(_ context.Context) error {
+		if err := h.Connect(); err != nil {
+			if errors.Is(err, rig.ErrCantConnect) || strings.Contains(err.Error(), "host key mismatch") {
+				return errors.Join(retry.ErrAbort, err)
+			}
+			return fmt.Errorf("failed to reconnect to %s: %w", h, err)
+		}
+		return nil
+	})
+}
+
 func (p *PrepareHosts) prepareHost(h *cluster.Host) error {
 	if c, ok := h.Configurer.(prepare); ok {
 		if err := c.Prepare(h); err != nil {
@@ -39,8 +73,8 @@ func (p *PrepareHosts) prepareHost(h *cluster.Host) error {
 
 	if len(h.Environment) > 0 {
 		log.Infof("%s: updating environment", h)
-		if err := h.Configurer.UpdateEnvironment(h, h.Environment); err != nil {
-			return err
+		if err := p.updateEnvironment(h); err != nil {
+			return fmt.Errorf("failed to updated environment: %w", err)
 		}
 	}