keep the staking keys and client wallet of the previous node version on plugin update

int/config/plugin_info.go

@@ -2,8 +2,11 @@ package config
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"path/filepath"
+	"strconv"
+	"strings"
 	"sync"
 
 	"github.com/massalabs/node-manager-plugin/int/utils"
@@ -64,6 +67,164 @@ func (pi *PluginInfo) RemoveOldNodeVersionsArtifacts(folderPath string) error {
 	return nil
 }
 
+// nodeWalletFolders are the folders, relative to a node version folder, that hold the keys of the node user:
+// the staking keys of the node and the wallet of the client. Both are encrypted with the node password.
+var nodeWalletFolders = []string{
+	filepath.Join("massa-node", "config", "staking_wallets"),
+	filepath.Join("massa-client", "wallets"),
+}
+
+/*
+MigrateNodeVersionsData copies the staking keys and the client wallet of the previous node version of each network
+into the current version folder (e.g. after a plugin update), so that the node keeps staking the same addresses
+with the same password. A folder is only copied when the current version does not have it yet or has it empty.
+It must run before RemoveOldNodeVersionsArtifacts, which deletes the previous version folders.
+*/
+func (pi *PluginInfo) MigrateNodeVersionsData(folderPath string) error {
+	entries, err := os.ReadDir(folderPath)
+	if err != nil {
+		return fmt.Errorf("failed to read directory %s: %v", folderPath, err)
+	}
+
+	for _, currentVersion := range []string{pi.MainnetVersion, pi.BuildnetVersion} {
+		previousVersion := latestPreviousVersion(entries, currentVersion)
+		if previousVersion == "" {
+			continue
+		}
+
+		for _, walletFolder := range nodeWalletFolders {
+			src := filepath.Join(folderPath, previousVersion, walletFolder)
+			dst := filepath.Join(folderPath, currentVersion, walletFolder)
+
+			if isEmptyOrMissingDir(src) || !isEmptyOrMissingDir(dst) {
+				continue
+			}
+
+			if err := copyDirAtomically(src, dst); err != nil {
+				return fmt.Errorf("failed to copy %s to %s: %w", src, dst, err)
+			}
+		}
+	}
+
+	return nil
+}
+
+// latestPreviousVersion returns the highest version folder of the same network as currentVersion
+// (same prefix, e.g. MAIN or DEVN), other than currentVersion. It returns an empty string if there is none.
+func latestPreviousVersion(entries []os.DirEntry, currentVersion string) string {
+	prefix, _, found := strings.Cut(currentVersion, ".")
+	if !found {
+		return ""
+	}
+
+	latest := ""
+	for _, entry := range entries {
+		name := entry.Name()
+		if !entry.IsDir() || name == currentVersion || !strings.HasPrefix(name, prefix+".") {
+			continue
+		}
+
+		if latest == "" || compareVersions(name, latest) > 0 {
+			latest = name
+		}
+	}
+
+	return latest
+}
+
+// compareVersions compares the numeric parts of two node versions of the same network (e.g. MAIN.4.10 > MAIN.4.2).
+func compareVersions(a, b string) int {
+	partsA := strings.Split(a, ".")[1:]
+	partsB := strings.Split(b, ".")[1:]
+
+	for i := 0; i < len(partsA) && i < len(partsB); i++ {
+		numA, errA := strconv.Atoi(partsA[i])
+		numB, errB := strconv.Atoi(partsB[i])
+		if errA != nil || errB != nil {
+			return strings.Compare(partsA[i], partsB[i])
+		}
+
+		if numA != numB {
+			return numA - numB
+		}
+	}
+
+	return len(partsA) - len(partsB)
+}
+
+func isEmptyOrMissingDir(path string) bool {
+	entries, err := os.ReadDir(path)
+	return err != nil || len(entries) == 0
+}
+
+// copyDirAtomically copies src to a temporary folder next to dst, then renames it to dst,
+// so that an interrupted copy does not leave a partial dst that would prevent the next attempt.
+func copyDirAtomically(src, dst string) error {
+	tmp := dst + ".migrating"
+	if err := os.RemoveAll(tmp); err != nil {
+		return err
+	}
+
+	if err := copyDir(src, tmp); err != nil {
+		return err
+	}
+
+	// dst is missing or empty
+	if err := os.RemoveAll(dst); err != nil {
+		return err
+	}
+
+	return os.Rename(tmp, dst)
+}
+
+// copyDir copies the src folder content into dst, keeping the files permissions.
+func copyDir(src, dst string) error {
+	return filepath.WalkDir(src, func(path string, entry os.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+
+		relPath, err := filepath.Rel(src, path)
+		if err != nil {
+			return err
+		}
+
+		target := filepath.Join(dst, relPath)
+
+		info, err := entry.Info()
+		if err != nil {
+			return err
+		}
+
+		if entry.IsDir() {
+			return os.MkdirAll(target, info.Mode().Perm())
+		}
+
+		return copyFile(path, target, info.Mode().Perm())
+	})
+}
+
+func copyFile(src, dst string, perm os.FileMode) error {
+	in, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+	// the file is only read: a close error does not affect the copy
+	defer func() { _ = in.Close() }()
+
+	out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm)
+	if err != nil {
+		return err
+	}
+
+	if _, err := io.Copy(out, in); err != nil {
+		_ = out.Close()
+		return err
+	}
+
+	return out.Close()
+}
+
 func (pi *PluginInfo) GetPwdByNetwork(isMainnet bool) string {
 	pi.mu.RLock()
 	defer pi.mu.RUnlock()

int/config/plugin_info_test.go

@@ -50,6 +50,94 @@ func TestRemoveOldNodeVersionsArtifacts_ReadDirError(t *testing.T) {
 	}
 }
 
+func TestMigrateNodeVersionsData(t *testing.T) {
+	root := t.TempDir()
+
+	pi := &PluginInfo{
+		MainnetVersion:  "MAIN.5.1",
+		BuildnetVersion: "DEVN.30.2",
+	}
+
+	stakingWallets := filepath.Join("massa-node", "config", "staking_wallets")
+	clientWallets := filepath.Join("massa-client", "wallets")
+
+	// previous mainnet versions: the latest one holds the keys to migrate
+	mustMkdir(t, filepath.Join(root, "MAIN.4.10", stakingWallets))
+	mustWriteFile(t, filepath.Join(root, "MAIN.4.10", stakingWallets, "wallet_AU1old.yaml"), []byte("old"))
+	mustMkdir(t, filepath.Join(root, "MAIN.5.0", stakingWallets))
+	mustWriteFile(t, filepath.Join(root, "MAIN.5.0", stakingWallets, "wallet_AU1staker.yaml"), []byte("staker"))
+	mustMkdir(t, filepath.Join(root, "MAIN.5.0", clientWallets))
+	mustWriteFile(t, filepath.Join(root, "MAIN.5.0", clientWallets, "wallet_AU1client.yaml"), []byte("client"))
+
+	// the new version is shipped with empty config folders
+	mustMkdir(t, filepath.Join(root, "MAIN.5.1", "massa-node", "config"))
+	mustMkdir(t, filepath.Join(root, "MAIN.5.1", clientWallets))
+
+	// previous buildnet version: the new version already has a client wallet, which must not be replaced
+	mustMkdir(t, filepath.Join(root, "DEVN.30.0", stakingWallets))
+	mustWriteFile(t, filepath.Join(root, "DEVN.30.0", stakingWallets, "wallet_AU1build.yaml"), []byte("build"))
+	mustMkdir(t, filepath.Join(root, "DEVN.30.0", clientWallets))
+	mustWriteFile(t, filepath.Join(root, "DEVN.30.0", clientWallets, "wallet_AU1oldclient.yaml"), []byte("old"))
+	mustMkdir(t, filepath.Join(root, "DEVN.30.2", clientWallets))
+	mustWriteFile(t, filepath.Join(root, "DEVN.30.2", clientWallets, "wallet_AU1newclient.yaml"), []byte("new"))
+
+	if err := pi.MigrateNodeVersionsData(root); err != nil {
+		t.Fatalf("MigrateNodeVersionsData returned error: %v", err)
+	}
+
+	assertContent(t, filepath.Join(root, "MAIN.5.1", stakingWallets, "wallet_AU1staker.yaml"), "staker")
+	assertExists(t, filepath.Join(root, "MAIN.5.1", stakingWallets, "wallet_AU1old.yaml"), false)
+	assertContent(t, filepath.Join(root, "MAIN.5.1", clientWallets, "wallet_AU1client.yaml"), "client")
+
+	assertContent(t, filepath.Join(root, "DEVN.30.2", stakingWallets, "wallet_AU1build.yaml"), "build")
+	assertContent(t, filepath.Join(root, "DEVN.30.2", clientWallets, "wallet_AU1newclient.yaml"), "new")
+	assertExists(t, filepath.Join(root, "DEVN.30.2", clientWallets, "wallet_AU1oldclient.yaml"), false)
+
+	// a second run (next plugin start) changes nothing
+	if err := pi.MigrateNodeVersionsData(root); err != nil {
+		t.Fatalf("second MigrateNodeVersionsData returned error: %v", err)
+	}
+	assertContent(t, filepath.Join(root, "MAIN.5.1", stakingWallets, "wallet_AU1staker.yaml"), "staker")
+
+	// then the previous versions are removed, the keys stay in the current ones
+	if err := pi.RemoveOldNodeVersionsArtifacts(root); err != nil {
+		t.Fatalf("RemoveOldNodeVersionsArtifacts returned error: %v", err)
+	}
+	assertExists(t, filepath.Join(root, "MAIN.5.0"), false)
+	assertContent(t, filepath.Join(root, "MAIN.5.1", stakingWallets, "wallet_AU1staker.yaml"), "staker")
+}
+
+func TestMigrateNodeVersionsData_NoPreviousVersion(t *testing.T) {
+	root := t.TempDir()
+	pi := &PluginInfo{MainnetVersion: "MAIN.5.1", BuildnetVersion: "DEVN.30.2"}
+
+	mustMkdir(t, filepath.Join(root, "MAIN.5.1", "massa-client"))
+
+	if err := pi.MigrateNodeVersionsData(root); err != nil {
+		t.Fatalf("MigrateNodeVersionsData returned error: %v", err)
+	}
+	assertExists(t, filepath.Join(root, "MAIN.5.1", "massa-client", "wallets"), false)
+}
+
+func TestCompareVersions(t *testing.T) {
+	cases := []struct {
+		a, b string
+		sign int
+	}{
+		{"MAIN.5.0", "MAIN.4.2", 1},
+		{"MAIN.4.10", "MAIN.4.2", 1},
+		{"DEVN.29.2", "DEVN.30.0", -1},
+		{"MAIN.5.0", "MAIN.5.0", 0},
+	}
+
+	for _, c := range cases {
+		got := compareVersions(c.a, c.b)
+		if (got > 0) != (c.sign > 0) || (got < 0) != (c.sign < 0) {
+			t.Errorf("compareVersions(%s, %s) = %d, expected sign %d", c.a, c.b, got, c.sign)
+		}
+	}
+}
+
 // helpers
 func mustMkdir(t *testing.T, path string) {
 	t.Helper()
@@ -75,3 +163,14 @@ func assertExists(t *testing.T, path string, shouldExist bool) {
 		t.Fatalf("expected %s to be removed, but it still exists", path)
 	}
 }
+
+func assertContent(t *testing.T, path, expected string) {
+	t.Helper()
+	data, err := os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("read file failed for %s: %v", path, err)
+	}
+	if string(data) != expected {
+		t.Fatalf("expected %s to contain %q, got %q", path, expected, string(data))
+	}
+}

int/node-bin-dir-manager/node_dir_manager.go

@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 
 	"github.com/massalabs/node-manager-plugin/int/config"
+	"github.com/massalabs/station/pkg/logger"
 )
 
 const (
@@ -103,7 +104,14 @@ func (ndm *nodeDirManager) init() error {
 		ndm.extension = ".exe"
 	}
 
-	// clean old node versions if any
+	// keep the staking keys and the client wallet of the previous node versions, then clean them if any.
+	// If the migration fails, the previous version folders are kept so that the keys are not lost:
+	// it is attempted again at the next start.
+	if err := config.GlobalPluginInfo.MigrateNodeVersionsData(ndm.nodeFolderPath); err != nil {
+		logger.Errorf("failed to migrate the wallets of the previous node versions, keeping their folders: %v", err)
+		return nil
+	}
+
 	if err := config.GlobalPluginInfo.RemoveOldNodeVersionsArtifacts(ndm.nodeFolderPath); err != nil {
 		return fmt.Errorf("failed to remove old node versions bin folders: %+w", err)
 	}