improve sui blockchain

Author: FelixFan1992

framework/components/blockchain/sui.go

@@ -1,15 +1,18 @@
 package blockchain
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"time"
 
 	"github.com/block-vision/sui-go-sdk/models"
 	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/pkg/stdcopy"
 	"github.com/docker/go-connections/nat"
 	"github.com/go-resty/resty/v2"
 	"github.com/testcontainers/testcontainers-go"
@@ -23,6 +26,10 @@ const (
 	DefaultFaucetPort    = "9123/tcp"
 	DefaultFaucetPortNum = "9123"
 	DefaultSuiNodePort   = "9000"
+	// DefaultSuiImage is the mysten/sui-tools image when Input.Image is empty on non-arm64 hosts.
+	DefaultSuiImage = "mysten/sui-tools:devnet-v1.69.0"
+	// DefaultSuiImageARM64 is used when Input.Image is empty on arm64 (e.g. Apple Silicon).
+	DefaultSuiImageARM64 = "mysten/sui-tools:ci-arm64"
 )
 
 // SuiWalletInfo info about Sui account/wallet
@@ -54,34 +61,95 @@ func fundAccount(url string, address string) error {
 	return nil
 }
 
+// demuxDockerExecOutput converts Docker exec attach output to plain text when it uses the
+// multiplexed stream format (first byte 1=stdout / 2=stderr). Must run before stripping 0x01,
+// which appears in stream headers and would corrupt the stream if removed globally.
+func demuxDockerExecOutput(raw string) string {
+	if len(raw) == 0 {
+		return raw
+	}
+	if raw[0] != 1 && raw[0] != 2 {
+		return raw
+	}
+	var stdout, stderr bytes.Buffer
+	if _, err := stdcopy.StdCopy(&stdout, &stderr, strings.NewReader(raw)); err != nil {
+		return raw
+	}
+	out := stdout.String() + stderr.String()
+	// Invalid or partial multiplex streams can make StdCopy succeed with empty output; keep raw so
+	// parseSuiKeytoolGenerateJSON can still find JSON after a single-byte preamble (e.g. 0x01).
+	if out == "" {
+		return raw
+	}
+
+	return out
+}
+
+// parseSuiKeytoolGenerateJSON extracts a SuiWalletInfo from `sui keytool generate --json` output.
+// The CLI may print a preamble, and v1.69+ may emit compact one-line JSON; older parsers assumed a
+// legacy layout (newline after '{') and corrupt compact output.
+func parseSuiKeytoolGenerateJSON(keyOut string) (*SuiWalletInfo, error) {
+	text := demuxDockerExecOutput(keyOut)
+	s := strings.ReplaceAll(text, "\x00", "")
+	for i := range s {
+		if s[i] != '{' {
+			continue
+		}
+		var key SuiWalletInfo
+		dec := json.NewDecoder(bytes.NewReader([]byte(s[i:])))
+		if err := dec.Decode(&key); err != nil {
+			continue
+		}
+		if key.SuiAddress != "" {
+			return &key, nil
+		}
+	}
+
+	return nil, fmt.Errorf("failed to parse SuiWalletInfo from keytool output: %.200q", keyOut)
+}
+
 // generateKeyData generates a wallet and returns all the data
 func generateKeyData(ctx context.Context, containerName string, keyCipherType string) (*SuiWalletInfo, error) {
-	cmdStr := []string{"sui", "keytool", "generate", keyCipherType, "--json"}
 	dc, err := framework.NewDockerClient()
 	if err != nil {
 		return nil, err
 	}
+
+	// Ensure a valid Sui client config exists. `sui start --force-regenesis`
+	// creates its config under /root/.sui/sui_config/ but the client.yaml it
+	// generates may not exist yet when this runs, so we use `sui client --yes`
+	// with an explicit config flag to force creation.
+	initCmd := []string{"sui", "client", "--client.config", "/root/.sui/sui_config/client.yaml", "--yes", "envs"}
+	if initOut, initErr := dc.ExecContainerWithContext(ctx, containerName, initCmd); initErr != nil {
+		framework.L.Warn().Err(initErr).Str("out", initOut).Msg("sui client init returned error (may be harmless)")
+	}
+
+	cmdStr := []string{"sui", "keytool", "generate", keyCipherType, "--json"}
 	keyOut, err := dc.ExecContainerWithContext(ctx, containerName, cmdStr)
 	if err != nil {
 		return nil, err
 	}
-	// formatted JSON with, no plain --json version, remove special symbols
-	cleanKey := strings.ReplaceAll(keyOut, "\x00", "")
-	cleanKey = strings.ReplaceAll(cleanKey, "\x01", "")
-	cleanKey = strings.ReplaceAll(cleanKey, "\x02", "")
-	cleanKey = strings.ReplaceAll(cleanKey, "\n", "")
-	cleanKey = "{" + cleanKey[2:]
-	var key *SuiWalletInfo
-	if err := json.Unmarshal([]byte(cleanKey), &key); err != nil {
-		return nil, err
+	key, err := parseSuiKeytoolGenerateJSON(keyOut)
+	if err != nil {
+		return nil, fmt.Errorf("%w (raw output: %.300q)", err, keyOut)
 	}
-	framework.L.Info().Interface("Key", key).Msg("Test key")
+
+	framework.L.Info().Str("suiAddress", key.SuiAddress).Msg("CTF test key generated")
+
 	return key, nil
 }
 
 func defaultSui(in *Input) {
 	if in.Image == "" {
-		in.Image = "mysten/sui-tools:devnet-v1.68.0"
+		if runtime.GOARCH == "arm64" {
+			in.Image = DefaultSuiImageARM64
+			if in.ImagePlatform == nil {
+				arm := "linux/arm64"
+				in.ImagePlatform = &arm
+			}
+		} else {
+			in.Image = DefaultSuiImage
+		}
 	}
 	if in.Port == "" {
 		in.Port = DefaultSuiNodePort
@@ -95,17 +163,25 @@ func newSui(ctx context.Context, in *Input) (*Output, error) {
 	defaultSui(in)
 	containerName := framework.DefaultTCName("blockchain-node")
 
-	absPath, err := filepath.Abs(in.ContractsDir)
-	if err != nil {
-		return nil, err
+	var files []testcontainers.ContainerFile
+	if in.ContractsDir != "" {
+		absPath, err := filepath.Abs(in.ContractsDir)
+		if err != nil {
+			return nil, err
+		}
+		files = []testcontainers.ContainerFile{
+			{
+				HostFilePath:      absPath,
+				ContainerFilePath: "/",
+			},
+		}
 	}
 
 	// Sui container always listens on port 9000 internally
 	containerPort := fmt.Sprintf("%s/tcp", DefaultSuiNodePort)
 
-	// default to amd64, unless otherwise specified
 	imagePlatform := "linux/amd64"
-	if in.ImagePlatform != nil {
+	if in.ImagePlatform != nil && *in.ImagePlatform != "" {
 		imagePlatform = *in.ImagePlatform
 	}
 
@@ -150,13 +226,7 @@ func newSui(ctx context.Context, in *Input) (*Output, error) {
 			"--force-regenesis",
 			"--with-faucet",
 		},
-		Files: []testcontainers.ContainerFile{
-			{
-				HostFilePath:      absPath,
-				ContainerFilePath: "/",
-			},
-		},
-		// we need faucet for funding
+		Files:      files,
 		WaitingFor: wait.ForListeningPort(DefaultFaucetPort).WithStartupTimeout(1 * time.Minute).WithPollInterval(200 * time.Millisecond),
 	}
 
@@ -183,6 +253,7 @@ func newSui(ctx context.Context, in *Input) (*Output, error) {
 		Type:                in.Type,
 		Family:              FamilySui,
 		ContainerName:       containerName,
+		Container:           c,
 		NetworkSpecificData: &NetworkSpecificData{SuiAccount: suiAccount},
 		Nodes: []*Node{
 			{

framework/components/blockchain/sui_test.go

@@ -0,0 +1,54 @@
+package blockchain
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseSuiKeytoolGenerateJSON(t *testing.T) {
+	t.Parallel()
+
+	const addr = "0xabc"
+	compact := `{"alias":null,"flag":0,"keyScheme":"ed25519","mnemonic":"a b c","peerId":"p","publicBase64Key":"k","suiAddress":"` + addr + `"}`
+
+	t.Run("compact one-line JSON", func(t *testing.T) {
+		t.Parallel()
+		got, err := parseSuiKeytoolGenerateJSON(compact)
+		require.NoError(t, err)
+		require.Equal(t, addr, got.SuiAddress)
+	})
+
+	t.Run("preamble before JSON", func(t *testing.T) {
+		t.Parallel()
+		in := "some log line\n" + compact
+		got, err := parseSuiKeytoolGenerateJSON(in)
+		require.NoError(t, err)
+		require.Equal(t, addr, got.SuiAddress)
+	})
+
+	t.Run("legacy newline after brace (old parser shape)", func(t *testing.T) {
+		t.Parallel()
+		legacy := "{\n  \"suiAddress\": \"" + addr + "\"\n}"
+		got, err := parseSuiKeytoolGenerateJSON(legacy)
+		require.NoError(t, err)
+		require.Equal(t, addr, got.SuiAddress)
+	})
+
+	t.Run("docker multiplexed stdout", func(t *testing.T) {
+		t.Parallel()
+		// stdcopy multiplex: 1 = stdout, then payload
+		mux := string([]byte{1}) + compact
+		got, err := parseSuiKeytoolGenerateJSON(mux)
+		require.NoError(t, err)
+		require.Equal(t, addr, got.SuiAddress)
+	})
+
+	t.Run("invalid", func(t *testing.T) {
+		t.Parallel()
+		_, err := parseSuiKeytoolGenerateJSON("no json here")
+		require.Error(t, err)
+		require.True(t, strings.Contains(err.Error(), "failed to parse"))
+	})
+}