testutil/compose: add dkg lock (#575)

Adds support for DKG key generation to compose lock command.

category: feature
ticket: #568

testutil/compose/compose/main.go

@@ -87,9 +87,10 @@ func newDefineCmd() *cobra.Command {
 	dir := addDirFlag(cmd.Flags())
 	clean := cmd.Flags().Bool("clean", true, "Clean compose dir before defining a new cluster")
 	seed := cmd.Flags().Int("seed", int(time.Now().UnixNano()), "Randomness seed")
+	keygen := cmd.Flags().String("keygen", "", "Key generation process: create, split, dkg")
 
 	cmd.RunE = func(cmd *cobra.Command, _ []string) error {
-		return compose.Define(cmd.Context(), *dir, *clean, *seed)
+		return compose.Define(cmd.Context(), *dir, *clean, *seed, *keygen)
 	}
 
 	return cmd

testutil/compose/define.go

@@ -35,7 +35,7 @@ import (
 )
 
 // Define defines a compose cluster; including both keygen and running definitions.
-func Define(ctx context.Context, dir string, clean bool, seed int) error {
+func Define(ctx context.Context, dir string, clean bool, seed int, keygen string) error {
 	ctx = log.WithTopic(ctx, "define")
 
 	if clean {
@@ -53,6 +53,10 @@ func Define(ctx context.Context, dir string, clean bool, seed int) error {
 
 	conf, p2pkeys := newDefaultConfig(seed)
 
+	if keygen != "" {
+		conf.KeyGen = keyGen(keygen)
+	}
+
 	// TODO(corver): Serve a web UI to allow configuration of default values.
 
 	log.Info(ctx, "Using default config")

testutil/compose/define_test.go

@@ -31,7 +31,7 @@ func TestDefine(t *testing.T) {
 	dir, err := os.MkdirTemp("", "")
 	require.NoError(t, err)
 
-	err = compose.Define(context.Background(), dir, false, 1)
+	err = compose.Define(context.Background(), dir, false, 1, "")
 	require.NoError(t, err)
 
 	conf, err := os.ReadFile(path.Join(dir, "charon-compose.yml"))

testutil/compose/lock.go

@@ -17,6 +17,7 @@ package compose
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 	"path"
@@ -25,6 +26,7 @@ import (
 
 	"github.com/obolnetwork/charon/app/errors"
 	"github.com/obolnetwork/charon/app/log"
+	"github.com/obolnetwork/charon/cluster"
 )
 
 // Lock creates a docker-compose.yml from a charon-compose.yml for generating keys and a cluster lock file.
@@ -36,24 +38,45 @@ func Lock(ctx context.Context, dir string) error {
 		return err
 	}
 
-	if conf.KeyGen != keyGenCreate {
-		return errors.New("only keygen create supported")
-	}
-
-	// Only single node to call charon create cluster generate keys
-	n := node{EnvVars: []kv{
-		{"threshold", fmt.Sprint(conf.Def.Threshold)},
-		{"nodes", fmt.Sprint(len(conf.Def.Operators))},
-		{"cluster_dir", "/compose"},
-	}}
-
-	data := tmplData{
-		NodeOnly:         true,
-		ComposeDir:       dir,
-		CharonImageTag:   conf.ImageTag,
-		CharonEntrypoint: containerBinary,
-		CharonCommand:    cmdCreateCluster,
-		Nodes:            []node{n},
+	var data tmplData
+	switch conf.KeyGen {
+	case keyGenCreate:
+		// Only single node to call charon create cluster generate keys
+		n := node{EnvVars: []kv{
+			{"threshold", fmt.Sprint(conf.Def.Threshold)},
+			{"nodes", fmt.Sprint(len(conf.Def.Operators))},
+			{"cluster_dir", "/compose"},
+		}}
+
+		data = tmplData{
+			NodeOnly:         true,
+			ComposeDir:       dir,
+			CharonImageTag:   conf.ImageTag,
+			CharonEntrypoint: containerBinary,
+			CharonCommand:    cmdCreateCluster,
+			Nodes:            []node{n},
+		}
+	case keyGenDKG:
+
+		if err := writeDefinition(dir, conf.Def); err != nil {
+			return err
+		}
+
+		var nodes []node
+		for i := 0; i < len(conf.Def.Operators); i++ {
+			n := node{EnvVars: newNodeEnvs(i, true, true)}
+			nodes = append(nodes, n)
+		}
+
+		data = tmplData{
+			ComposeDir:       dir,
+			CharonImageTag:   conf.ImageTag,
+			CharonEntrypoint: containerBinary,
+			CharonCommand:    cmdDKG,
+			Nodes:            nodes,
+		}
+	default:
+		return errors.New("supported keygen")
 	}
 
 	log.Info(ctx, "Created docker-compose.yml")
@@ -62,6 +85,20 @@ func Lock(ctx context.Context, dir string) error {
 	return writeDockerCompose(dir, data)
 }
 
+func writeDefinition(dir string, def cluster.Definition) error {
+	b, err := json.MarshalIndent(def, "", " ")
+	if err != nil {
+		return errors.Wrap(err, "marshal definition")
+	}
+
+	err = os.WriteFile(path.Join(dir, "cluster-definition.json"), b, 0o755) //nolint:gosec
+	if err != nil {
+		return errors.Wrap(err, "write definition")
+	}
+
+	return nil
+}
+
 // newNodeEnvs returns the default node environment variable to run a charon docker container.
 
 func newNodeEnvs(index int, validatorMock, beaconMock bool) []kv {
@@ -90,8 +127,13 @@ func loadConfig(dir string) (config, error) {
 		return config{}, errors.Wrap(err, "load config")
 	}
 
+	b, err = yaml.YAMLToJSON(b)
+	if err != nil {
+		return config{}, errors.Wrap(err, "yaml config")
+	}
+
 	var resp config
-	if err := yaml.Unmarshal(b, &resp); err != nil {
+	if err := json.Unmarshal(b, &resp); err != nil {
 		return config{}, errors.Wrap(err, "unmarshal config")
 	}