Attempt to add upgradeConfig support

client/client.go

@@ -226,10 +226,11 @@ func (c *client) AddNode(ctx context.Context, name string, execPath string, opts
 	ret.applyOpts(opts)
 
 	req := &rpcpb.AddNodeRequest{
-		Name:         name,
-		ExecPath:     execPath,
-		NodeConfig:   &ret.globalNodeConfig,
-		ChainConfigs: ret.chainConfigs,
+		Name:           name,
+		ExecPath:       execPath,
+		NodeConfig:     &ret.globalNodeConfig,
+		ChainConfigs:   ret.chainConfigs,
+		UpgradeConfigs: ret.upgradeConfigs,
 	}
 
 	c.log.Info("add node", zap.String("name", name))
@@ -253,6 +254,7 @@ func (c *client) RestartNode(ctx context.Context, name string, opts ...OpOption)
 		req.WhitelistedSubnets = &ret.whitelistedSubnets
 	}
 	req.ChainConfigs = ret.chainConfigs
+	req.UpgradeConfigs = ret.upgradeConfigs
 
 	c.log.Info("restart node", zap.String("name", name))
 	return c.controlc.RestartNode(ctx, req)
@@ -283,8 +285,9 @@ func (c *client) LoadSnapshot(ctx context.Context, snapshotName string, opts ...
 	ret := &Op{}
 	ret.applyOpts(opts)
 	req := rpcpb.LoadSnapshotRequest{
-		SnapshotName: snapshotName,
-		ChainConfigs: ret.chainConfigs,
+		SnapshotName:   snapshotName,
+		ChainConfigs:   ret.chainConfigs,
+		UpgradeConfigs: ret.upgradeConfigs,
 	}
 	if ret.execPath != "" {
 		req.ExecPath = &ret.execPath
@@ -333,6 +336,7 @@ type Op struct {
 	customNodeConfigs  map[string]string
 	numSubnets         uint32
 	chainConfigs       map[string]string
+	upgradeConfigs     map[string]string
 }
 
 type OpOption func(*Op)
@@ -393,6 +397,13 @@ func WithChainConfigs(chainConfigs map[string]string) OpOption {
 	}
 }
 
+// Map from chain name to its upgrade json contents.
+func WithUpgradeConfigs(upgradeConfigs map[string]string) OpOption {
+	return func(op *Op) {
+		op.upgradeConfigs = upgradeConfigs
+	}
+}
+
 // Map from node name to its custom node config
 func WithCustomNodeConfigs(customNodeConfigs map[string]string) OpOption {
 	return func(op *Op) {

cmd/control/control.go

@@ -100,6 +100,7 @@ var (
 	rootDataDir        string
 	numSubnets         uint32
 	chainConfigs       string
+	upgradeConfigs     string
 )
 
 func newStartCommand() *cobra.Command {
@@ -163,6 +164,12 @@ func newStartCommand() *cobra.Command {
 		"",
 		"[optional] JSON string of map from chain id to its config file contents",
 	)
+	cmd.PersistentFlags().StringVar(
+		&upgradeConfigs,
+		"upgrade-configs",
+		"",
+		"[optional] JSON string of map from chain id to its upgrade file contents",
+	)
 	if err := cmd.MarkPersistentFlagRequired("avalanchego-path"); err != nil {
 		panic(err)
 	}
@@ -216,6 +223,13 @@ func startFunc(cmd *cobra.Command, args []string) error {
 		}
 		opts = append(opts, client.WithChainConfigs(chainConfigsMap))
 	}
+	if upgradeConfigs != "" {
+		upgradeConfigsMap := make(map[string]string)
+		if err := json.Unmarshal([]byte(upgradeConfigs), &upgradeConfigsMap); err != nil {
+			return err
+		}
+		opts = append(opts, client.WithUpgradeConfigs(upgradeConfigsMap))
+	}
 
 	ctx := getAsyncContext()
 
@@ -501,7 +515,13 @@ func newAddNodeCommand() *cobra.Command {
 		&chainConfigs,
 		"chain-configs",
 		"",
-		"[optional] JSON string of map that maps from chain id to its config file contents",
+		"[optional] JSON string of map from chain id to its config file contents",
+	)
+	cmd.PersistentFlags().StringVar(
+		&upgradeConfigs,
+		"upgrade-configs",
+		"",
+		"[optional] JSON string of map from chain id to its upgrade file contents",
 	)
 	return cmd
 }
@@ -534,6 +554,13 @@ func addNodeFunc(cmd *cobra.Command, args []string) error {
 		}
 		opts = append(opts, client.WithChainConfigs(chainConfigsMap))
 	}
+	if upgradeConfigs != "" {
+		upgradeConfigsMap := make(map[string]string)
+		if err := json.Unmarshal([]byte(upgradeConfigs), &upgradeConfigsMap); err != nil {
+			return err
+		}
+		opts = append(opts, client.WithUpgradeConfigs(upgradeConfigsMap))
+	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
 	info, err := cli.AddNode(
@@ -574,7 +601,13 @@ func newRestartNodeCommand() *cobra.Command {
 		&chainConfigs,
 		"chain-configs",
 		"",
-		"[optional] JSON string of map that maps from chain id to its config file contents",
+		"[optional] JSON string of map from chain id to its config file contents",
+	)
+	cmd.PersistentFlags().StringVar(
+		&upgradeConfigs,
+		"upgrade-configs",
+		"",
+		"[optional] JSON string of map from chain id to its upgrade file contents",
 	)
 	return cmd
 }
@@ -600,6 +633,13 @@ func restartNodeFunc(cmd *cobra.Command, args []string) error {
 		}
 		opts = append(opts, client.WithChainConfigs(chainConfigsMap))
 	}
+	if upgradeConfigs != "" {
+		upgradeConfigsMap := make(map[string]string)
+		if err := json.Unmarshal([]byte(upgradeConfigs), &upgradeConfigsMap); err != nil {
+			return err
+		}
+		opts = append(opts, client.WithUpgradeConfigs(upgradeConfigsMap))
+	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
 	info, err := cli.RestartNode(
@@ -800,7 +840,13 @@ func newLoadSnapshotCommand() *cobra.Command {
 		&chainConfigs,
 		"chain-configs",
 		"",
-		"[optional] JSON string of map that maps from chain id to its config file contents",
+		"[optional] JSON string of map from chain id to its config file contents",
+	)
+	cmd.PersistentFlags().StringVar(
+		&upgradeConfigs,
+		"upgrade-configs",
+		"",
+		"[optional] JSON string of map from chain id to its upgrade file contents",
 	)
 	cmd.PersistentFlags().StringVar(
 		&globalNodeConfig,
@@ -831,6 +877,13 @@ func loadSnapshotFunc(cmd *cobra.Command, args []string) error {
 		}
 		opts = append(opts, client.WithChainConfigs(chainConfigsMap))
 	}
+	if upgradeConfigs != "" {
+		upgradeConfigsMap := make(map[string]string)
+		if err := json.Unmarshal([]byte(upgradeConfigs), &upgradeConfigsMap); err != nil {
+			return err
+		}
+		opts = append(opts, client.WithUpgradeConfigs(upgradeConfigsMap))
+	}
 
 	if globalNodeConfig != "" {
 		ux.Print(log, logging.Yellow.Wrap("global node config provided, will be applied to all nodes: %s"), globalNodeConfig)

local/helpers.go

@@ -55,7 +55,7 @@ func writeFiles(genesis []byte, nodeRootDir string, nodeConfig *node.Config) ([]
 			return nil, fmt.Errorf("couldn't write file at %q: %w", f.path, err)
 		}
 	}
-	if nodeConfig.ChainConfigFiles != nil {
+	if nodeConfig.ChainConfigFiles != nil || nodeConfig.UpgradeConfigFiles != nil {
 		// only one flag and multiple files
 		chainConfigDir := filepath.Join(nodeRootDir, chainConfigSubDir)
 		flags = append(flags, fmt.Sprintf("--%s=%s", config.ChainConfigDirKey, chainConfigDir))
@@ -65,6 +65,12 @@ func writeFiles(genesis []byte, nodeRootDir string, nodeConfig *node.Config) ([]
 				return nil, fmt.Errorf("couldn't write file at %q: %w", chainConfigPath, err)
 			}
 		}
+		for chainAlias, chainUpgradeFile := range nodeConfig.UpgradeConfigFiles {
+			chainUpgradePath := filepath.Join(chainConfigDir, chainAlias, upgradeConfigFileName)
+			if err := createFileAndWrite(chainUpgradePath, []byte(chainUpgradeFile)); err != nil {
+				return nil, fmt.Errorf("couldn't write file at %q: %w", chainUpgradePath, err)
+			}
+		}
 	}
 	return flags, nil
 }

local/network.go

@@ -33,6 +33,7 @@ import (
 const (
 	defaultNodeNamePrefix = "node"
 	configFileName        = "config.json"
+	upgradeConfigFileName = "upgrade.json"
 	stakingKeyFileName    = "staking.key"
 	stakingCertFileName   = "staking.crt"
 	genesisFileName       = "genesis.json"
@@ -97,6 +98,8 @@ type localNetwork struct {
 	binaryPath string
 	// chain config files to use per default
 	chainConfigFiles map[string]string
+	// upgrade config files to use per default
+	upgradeConfigFiles map[string]string
 }
 
 var (
@@ -181,6 +184,7 @@ func init() {
 		ChainConfigFiles: map[string]string{
 			"C": string(cChainConfig),
 		},
+		UpgradeConfigFiles: map[string]string{},
 	}
 
 	for i := 0; i < len(defaultNetworkConfig.NodeConfigs); i++ {
@@ -392,6 +396,7 @@ func (ln *localNetwork) loadConfig(ctx context.Context, networkConfig network.Co
 	ln.flags = networkConfig.Flags
 	ln.binaryPath = networkConfig.BinaryPath
 	ln.chainConfigFiles = networkConfig.ChainConfigFiles
+	ln.upgradeConfigFiles = networkConfig.UpgradeConfigFiles
 
 	// Sort node configs so beacons start first
 	var nodeConfigs []node.Config
@@ -450,6 +455,12 @@ func (ln *localNetwork) addNode(nodeConfig node.Config) (node.Node, error) {
 			nodeConfig.ChainConfigFiles[k] = v
 		}
 	}
+	for k, v := range ln.upgradeConfigFiles {
+		_, ok := nodeConfig.UpgradeConfigFiles[k]
+		if !ok {
+			nodeConfig.UpgradeConfigFiles[k] = v
+		}
+	}
 	addNetworkFlags(ln.log, ln.flags, nodeConfig.Flags)
 
 	// it shouldn't happen that just one is empty, most probably both,
@@ -725,6 +736,7 @@ func (ln *localNetwork) RestartNode(
 	binaryPath string,
 	whitelistedSubnets string,
 	chainConfigs map[string]string,
+	upgradeConfigs map[string]string,
 ) error {
 	ln.lock.Lock()
 	defer ln.lock.Unlock()
@@ -753,6 +765,10 @@ func (ln *localNetwork) RestartNode(
 	for k, v := range chainConfigs {
 		nodeConfig.ChainConfigFiles[k] = v
 	}
+	// apply upgrade configs
+	for k, v := range upgradeConfigs {
+		nodeConfig.UpgradeConfigFiles[k] = v
+	}
 
 	if err := ln.removeNode(ctx, nodeName); err != nil {
 		return err

local/snapshot.go

@@ -28,6 +28,7 @@ func NewNetworkFromSnapshot(
 	binaryPath string,
 	buildDir string,
 	chainConfigs map[string]string,
+	upgradeConfigs map[string]string,
 	flags map[string]interface{},
 ) (network.Network, error) {
 	net, err := newNetwork(
@@ -45,7 +46,7 @@ func NewNetworkFromSnapshot(
 	if err != nil {
 		return net, err
 	}
-	err = net.loadSnapshot(context.Background(), snapshotName, binaryPath, buildDir, chainConfigs, flags)
+	err = net.loadSnapshot(context.Background(), snapshotName, binaryPath, buildDir, chainConfigs, upgradeConfigs, flags)
 	return net, err
 }
 
@@ -128,11 +129,12 @@ func (ln *localNetwork) SaveSnapshot(ctx context.Context, snapshotName string) (
 	}
 	// save network conf
 	networkConfig := network.Config{
-		Genesis:          string(ln.genesis),
-		Flags:            networkConfigFlags,
-		NodeConfigs:      []node.Config{},
-		BinaryPath:       ln.binaryPath,
-		ChainConfigFiles: ln.chainConfigFiles,
+		Genesis:            string(ln.genesis),
+		Flags:              networkConfigFlags,
+		NodeConfigs:        []node.Config{},
+		BinaryPath:         ln.binaryPath,
+		ChainConfigFiles:   ln.chainConfigFiles,
+		UpgradeConfigFiles: ln.upgradeConfigFiles,
 	}
 
 	for _, nodeConfig := range nodesConfig {
@@ -157,6 +159,7 @@ func (ln *localNetwork) loadSnapshot(
 	binaryPath string,
 	buildDir string,
 	chainConfigs map[string]string,
+	upgradeConfigs map[string]string,
 	flags map[string]interface{},
 ) error {
 	ln.lock.Lock()
@@ -208,14 +211,20 @@ func (ln *localNetwork) loadSnapshot(
 			networkConfig.NodeConfigs[i].Flags[config.BuildDirKey] = buildDir
 		}
 	}
-	// add chain configs
+	// add chain configs and upgrade configs
 	for i := range networkConfig.NodeConfigs {
 		if networkConfig.NodeConfigs[i].ChainConfigFiles == nil {
 			networkConfig.NodeConfigs[i].ChainConfigFiles = map[string]string{}
 		}
+		if networkConfig.NodeConfigs[i].UpgradeConfigFiles == nil {
+			networkConfig.NodeConfigs[i].UpgradeConfigFiles = map[string]string{}
+		}
 		for k, v := range chainConfigs {
 			networkConfig.NodeConfigs[i].ChainConfigFiles[k] = v
 		}
+		for k, v := range upgradeConfigs {
+			networkConfig.NodeConfigs[i].UpgradeConfigFiles[k] = v
+		}
 	}
 	return ln.loadConfig(ctx, networkConfig)
 }

network/config.go

@@ -59,6 +59,8 @@ type Config struct {
 	BinaryPath string `json:"binaryPath"`
 	// Chain config files to use per default, if not specified in node config
 	ChainConfigFiles map[string]string `json:"chainConfigFiles"`
+	// Upgrade config files to use per default, if not specified in node config
+	UpgradeConfigFiles map[string]string `json:"upgradeConfigFiles"`
 }
 
 // Validate returns an error if this config is invalid

network/network.go

@@ -53,8 +53,8 @@ type Network interface {
 	// Get name of available snapshots
 	GetSnapshotNames() ([]string, error)
 	// Restart a given node using the same config, optionally changing binary path,
-	// whitelisted subnets, and a map of chain configs
-	RestartNode(context.Context, string, string, string, map[string]string) error
+	// whitelisted subnets, a map of chain configs, and a map of upgrade configs
+	RestartNode(context.Context, string, string, string, map[string]string, map[string]string) error
 	// Create the specified blockchains
 	CreateBlockchains(context.Context, []BlockchainSpec) error
 	// Create the given numbers of subnets

network/node/node.go

@@ -72,6 +72,8 @@ type Config struct {
 	ConfigFile string `json:"configFile"`
 	// May be nil.
 	ChainConfigFiles map[string]string `json:"chainConfigFiles"`
+	// May be nil.
+	UpgradeConfigFiles map[string]string `json:"upgradeConfigFiles"`
 	// Flags can hold additional flags for the node.
 	// It can be empty.
 	// The precedence of flags handling is: