Attempt to add upgradeConfig support

client/client.go

@@ -246,6 +246,7 @@ func (c *client) AddNode(ctx context.Context, name string, execPath string, opts
 		req.StartRequest.PluginDir = &ret.pluginDir
 	}
 	req.StartRequest.ChainConfigs = ret.chainConfigs
+	req.StartRequest.UpgradeConfigs = ret.upgradeConfigs
 
 	zap.L().Info("add node", zap.String("name", name))
 	return c.controlc.AddNode(ctx, req)
@@ -270,6 +271,8 @@ func (c *client) RestartNode(ctx context.Context, name string, opts ...OpOption)
 	if ret.rootDataDir != "" {
 		req.RootDataDir = &ret.rootDataDir
 	}
+	req.ChainConfigs = ret.chainConfigs
+	req.UpgradeConfigs = ret.upgradeConfigs
 
 	zap.L().Info("restart node", zap.String("name", name))
 	return c.controlc.RestartNode(ctx, req)
@@ -350,6 +353,7 @@ type Op struct {
 	customNodeConfigs  map[string]string
 	numSubnets         uint32
 	chainConfigs       map[string]string
+	upgradeConfigs     map[string]string
 }
 
 type OpOption func(*Op)
@@ -410,6 +414,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

@@ -82,6 +82,7 @@ var (
 	rootDataDir               string
 	numSubnets                uint32
 	chainConfigs              string
+	upgradeConfigs            string
 )
 
 func newStartCommand() *cobra.Command {
@@ -587,6 +588,18 @@ func newRestartNodeCommand() *cobra.Command {
 		"",
 		"whitelisted subnets (comma-separated)",
 	)
+	cmd.PersistentFlags().StringVar(
+		&chainConfigs,
+		"chain-configs",
+		"",
+		"[optional] JSON string of map that maps from chain id to its config file contents",
+	)
+	cmd.PersistentFlags().StringVar(
+		&upgradeConfigs,
+		"upgrade-configs",
+		"",
+		"[optional] JSON string of map that maps from chain id to its upgrade file contents",
+	)
 	return cmd
 }
 
@@ -597,12 +610,32 @@ func restartNodeFunc(cmd *cobra.Command, args []string) error {
 	}
 	defer cli.Close()
 
+	opts := []client.OpOption{
+		client.WithExecPath(avalancheGoBinPath),
+		client.WithWhitelistedSubnets(whitelistedSubnets),
+	}
+
+	if chainConfigs != "" {
+		chainConfigsMap := make(map[string]string)
+		if err := json.Unmarshal([]byte(chainConfigs), &chainConfigsMap); err != nil {
+			return err
+		}
+		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.WithChainConfigs(upgradeConfigsMap))
+	}
+
 	ctx, cancel := context.WithTimeout(context.Background(), requestTimeout)
 	info, err := cli.RestartNode(
 		ctx,
 		nodeName,
-		client.WithExecPath(avalancheGoBinPath),
-		client.WithWhitelistedSubnets(whitelistedSubnets),
+		opts...,
 	)
 	cancel()
 	if err != nil {

local/network.go

@@ -37,6 +37,7 @@ import (
 const (
 	defaultNodeNamePrefix = "node"
 	configFileName        = "config.json"
+	upgradeConfigFileName = "upgrade.json"
 	stakingKeyFileName    = "staking.key"
 	stakingCertFileName   = "staking.crt"
 	genesisFileName       = "genesis.json"
@@ -322,6 +323,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(
@@ -338,7 +340,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
 }
 
@@ -809,6 +811,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()
@@ -865,9 +868,15 @@ func (ln *localNetwork) loadSnapshot(
 		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)
 }
@@ -1209,5 +1218,17 @@ func writeFiles(genesis []byte, nodeRootDir string, nodeConfig *node.Config) ([]
 			}
 		}
 	}
+	if nodeConfig.UpgradeConfigFiles != nil {
+		// only one flag and multiple files
+		chainConfigDir := filepath.Join(nodeRootDir, chainConfigSubDir)
+		// TODO: only need to specify this flag once
+		flags = append(flags, fmt.Sprintf("--%s=%s", config.ChainConfigDirKey, chainConfigDir))
+		for chainAlias, upgradeChainFile := range nodeConfig.UpgradeConfigFiles {
+			upgradeConfigPath := filepath.Join(chainConfigDir, chainAlias, upgradeConfigFileName)
+			if err := createFileAndWrite(upgradeConfigPath, []byte(upgradeChainFile)); err != nil {
+				return nil, fmt.Errorf("couldn't write file at %q: %w", upgradeConfigPath, err)
+			}
+		}
+	}
 	return flags, nil
 }

network/node/node.go

@@ -68,6 +68,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: