From 45d82f97aeaea2f24ca822eee526854930ec2ed5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9phane=20Graber?= <stgraber@stgraber.org>
Date: Mon, 11 Dec 2023 00:01:57 -0500
Subject: [PATCH] shared/cliconfig: Add keepalive proxy support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Stéphane Graber <stgraber@stgraber.org>
---
 shared/cliconfig/keepalive.go         | 60 +++++++++++++++++++++++++++
 shared/cliconfig/keepalive_windows.go | 13 ++++++
 shared/cliconfig/remote.go            | 18 ++++++--
 3 files changed, 88 insertions(+), 3 deletions(-)
 create mode 100644 shared/cliconfig/keepalive.go
 create mode 100644 shared/cliconfig/keepalive_windows.go

diff --git a/shared/cliconfig/keepalive.go b/shared/cliconfig/keepalive.go
new file mode 100644
index 00000000000..46d6be3e04e
--- /dev/null
+++ b/shared/cliconfig/keepalive.go
@@ -0,0 +1,60 @@
+//go:build !windows
+
+package cliconfig
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/lxc/incus/client"
+	"github.com/lxc/incus/shared/subprocess"
+	"github.com/lxc/incus/shared/util"
+)
+
+func (c *Config) handleKeepAlive(remote Remote, name string, args *incus.ConnectionArgs) (incus.InstanceServer, error) {
+	// Create the socker directory if missing.
+	socketDir := filepath.Join(c.ConfigDir, "keepalive")
+	err := os.Mkdir(socketDir, 0700)
+	if err != nil && !os.IsExist(err) {
+		return nil, err
+	}
+
+	// Attempt to use the existing socket.
+	socketPath := filepath.Join(socketDir, fmt.Sprintf("%s.socket", name))
+	d, err := incus.ConnectIncusUnix(socketPath, args)
+	if err != nil {
+		// Delete any existing sockets.
+		_ = os.Remove(socketPath)
+
+		// Spawn the proxy.
+		proc, err := subprocess.NewProcess("incus", []string{"remote", "proxy", name, socketPath, fmt.Sprintf("--timeout=%d", remote.KeepAlive)}, "", "")
+		if err != nil {
+			return nil, err
+		}
+
+		err = proc.Start(context.Background())
+		if err != nil {
+			return nil, err
+		}
+
+		// Try up to 10 times over 5s.
+		for i := 0; i < 10; i++ {
+			if util.PathExists(socketPath) {
+				break
+			}
+
+			time.Sleep(500 * time.Millisecond)
+		}
+
+		// Connect to the proxy.
+		d, err = incus.ConnectIncusUnix(socketPath, args)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return d, nil
+}
diff --git a/shared/cliconfig/keepalive_windows.go b/shared/cliconfig/keepalive_windows.go
new file mode 100644
index 00000000000..1d3b1398ce6
--- /dev/null
+++ b/shared/cliconfig/keepalive_windows.go
@@ -0,0 +1,13 @@
+//go:build windows
+
+package cliconfig
+
+import (
+	"fmt"
+
+	"github.com/lxc/incus/client"
+)
+
+func (c *Config) handleKeepAlive(remote Remote, name string, args *incus.ConnectionArgs) (incus.InstanceServer, error) {
+	return nil, fmt.Errorf("Keepalive isn't supported on Windows")
+}
diff --git a/shared/cliconfig/remote.go b/shared/cliconfig/remote.go
index 1c7350d0148..6b18ee0af4a 100644
--- a/shared/cliconfig/remote.go
+++ b/shared/cliconfig/remote.go
@@ -103,9 +103,21 @@ func (c *Config) GetInstanceServer(name string) (incus.InstanceServer, error) {
 		return nil, fmt.Errorf("Missing TLS client certificate and key")
 	}
 
-	d, err := incus.ConnectIncus(remote.Addr, args)
-	if err != nil {
-		return nil, err
+	var d incus.InstanceServer
+	if remote.KeepAlive > 0 {
+		d, err = c.handleKeepAlive(remote, name, args)
+		if err != nil {
+			// On proxy failure, just fallback to regular client.
+			d, err = incus.ConnectIncus(remote.Addr, args)
+			if err != nil {
+				return nil, err
+			}
+		}
+	} else {
+		d, err = incus.ConnectIncus(remote.Addr, args)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	if remote.Project != "" && remote.Project != "default" {