Load generator docker

Dockerfile

@@ -89,6 +89,13 @@ RUN --mount=type=cache,target=/go/pkg/mod \
     --mount=type=cache,target=/root/.cache/go-build \
     go build -o ./bin/relay ./cmd
 
+# Traffic Generator V2 build stage
+FROM common-builder AS generator2-builder
+WORKDIR /app/test/v2
+RUN --mount=type=cache,target=/go/pkg/mod \
+    --mount=type=cache,target=/root/.cache/go-build \
+    make build
+
 # Final stages for each component
 FROM alpine:3.18 AS churner
 COPY --from=churner-builder /app/operators/bin/churner /usr/local/bin
@@ -129,3 +136,7 @@ ENTRYPOINT ["controller"]
 FROM alpine:3.18 AS relay
 COPY --from=relay-builder /app/relay/bin/relay /usr/local/bin
 ENTRYPOINT ["relay"]
+
+FROM alpine:3.18 AS generator2
+COPY --from=generator2-builder /app/test/v2/bin/load /usr/local/bin
+ENTRYPOINT ["load", "-", "-"]

docker-bake.hcl

@@ -205,7 +205,7 @@ target "traffic-generator-internal" {
 
 target "traffic-generator-v2" {
   context    = "."
-  dockerfile = "./trafficgenerator-v2.Dockerfile"
+  dockerfile = "./Dockerfile"
   target     = "generator2"
   tags       = ["${REGISTRY}/${REPO}/traffic-generator-v2:${BUILD_TAG}"]
 }

test/v2/client/test_client.go

@@ -75,20 +75,12 @@ func NewTestClient(
 
 	// Construct the disperser client
 
-	privateKeyFile, err := ResolveTildeInPath(config.KeyPath)
+	privateKey, err := loadPrivateKey(config.KeyPath, config.KeyVar)
 	if err != nil {
-		return nil, fmt.Errorf("failed to resolve tilde in path: %w", err)
-	}
-	privateKey, err := os.ReadFile(privateKeyFile)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read private key file: %w", err)
+		return nil, fmt.Errorf("failed to load private key: %w", err)
 	}
 
-	privateKeyString := string(privateKey)
-	privateKeyString = strings.Trim(privateKeyString, "\n \t")
-	privateKeyString, _ = strings.CutPrefix(privateKeyString, "0x")
-
-	signer, err := auth.NewLocalBlobRequestSigner(privateKeyString)
+	signer, err := auth.NewLocalBlobRequestSigner(privateKey)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create signer: %w", err)
 	}
@@ -149,7 +141,7 @@ func NewTestClient(
 
 	ethClientConfig := geth.EthClientConfig{
 		RPCURLs:          config.EthRPCURLs,
-		PrivateKeyString: privateKeyString,
+		PrivateKeyString: privateKey,
 		NumConfirmations: 0,
 		NumRetries:       3,
 	}
@@ -277,13 +269,45 @@ func NewTestClient(
 		retrievalClient:           retrievalClient,
 		validatorPayloadRetriever: validatorPayloadRetriever,
 		certVerifier:              certVerifier,
-		privateKey:                privateKeyString,
+		privateKey:                privateKey,
 		metricsRegistry:           metrics.registry,
 		metrics:                   metrics,
 		blobCodec:                 blobCodec,
 	}, nil
 }
 
+// loadPrivateKey loads the private key from the file/env var specified in the config.
+func loadPrivateKey(keyPath string, keyVar string) (string, error) {
+	if keyPath != "" {
+		privateKeyFile, err := ResolveTildeInPath(keyPath)
+		if err != nil {
+			return "", fmt.Errorf("failed to resolve tilde in path: %w", err)
+		}
+		privateKey, err := os.ReadFile(privateKeyFile)
+		if err != nil {
+			return "", fmt.Errorf("failed to read private key file: %w", err)
+		}
+
+		return formatPrivateKey(string(privateKey)), nil
+	}
+
+	if keyVar == "" {
+		return "", fmt.Errorf("either KeyPath or KeyVar must be set")
+	}
+	privateKey := os.Getenv(keyVar)
+	if privateKey == "" {
+		return "", fmt.Errorf("key not found in environment variable %s", keyVar)
+	}
+	return formatPrivateKey(privateKey), nil
+}
+
+// formatPrivateKey formats the private key by removing leading/trailing whitespace and "0x" prefix.
+func formatPrivateKey(privateKey string) string {
+	privateKey = strings.Trim(privateKey, "\n \t")
+	privateKey, _ = strings.CutPrefix(privateKey, "0x")
+	return privateKey
+}
+
 // GetConfig returns the test client's configuration.
 func (c *TestClient) GetConfig() *TestClientConfig {
 	return c.config

test/v2/client/test_client_config.go

@@ -9,9 +9,15 @@ import (
 type TestClientConfig struct {
 	// The location where persistent test data is stored (e.g. SRS files). Often private keys are stored here too.
 	TestDataPath string
-	// The location where the test client's private key is stored.
-	// This is the key for the account that is paying for dispersals.
+	// The location where the test client's private key is stored. This is the key for the account that is
+	// paying for dispersals.
+	//
+	// Either this or KeyVar must be set. If both are set, KeyPath is used.
 	KeyPath string
+	// The environment variable that contains the private key for the account that is paying for dispersals.
+	//
+	// This is used if KeyPath is not set.
+	KeyVar string
 	// The disperser's hostname (url or IP address)
 	DisperserHostname string
 	// The disperser's port

test/v2/client/test_client_setup.go

@@ -183,17 +183,6 @@ func setupFilesystem(logger logging.Logger, config *TestClientConfig) error {
 		return fmt.Errorf("failed to create SRS tables directory: %w", err)
 	}
 
-	// Check to see if the private key file exists. If not, stop the test.
-	filePath, err := ResolveTildeInPath(config.KeyPath)
-	if err != nil {
-		return fmt.Errorf("failed to resolve tilde in path: %w", err)
-	}
-	_, err = os.Stat(filePath)
-	if err != nil {
-		return fmt.Errorf("private key file %s does not exist. This file should "+
-			"contain the private key for the account used in the test, in hex: %w", filePath, err)
-	}
-
 	// If any of the srs files do not exist, download them.
 	err = ensureFileIsPresent(config, SRSPathG1, G1URL)
 	if err != nil {

test/v2/load/main/load_main.go

@@ -11,12 +11,27 @@ import (
 
 func main() {
 	if len(os.Args) != 3 {
-		panic(fmt.Sprintf("Expected 3 args, got %d. Usage: %s <env_file> <load_file>\n",
+		panic(fmt.Sprintf("Expected 3 args, got %d. Usage: %s <env_file> <load_file>.\n"+
+			"If '-' is passed in lieu of a config file, the config file path is read from the environment variable "+
+			"$GENERATOR_ENV or $GENERATOR_LOAD, respectively.\n",
 			len(os.Args), os.Args[0]))
 	}
 
 	envFile := os.Args[1]
+	if envFile == "-" {
+		envFile = os.Getenv("GENERATOR_ENV")
+		if envFile == "" {
+			panic("$GENERATOR_ENV not set")
+		}
+	}
+
 	loadFile := os.Args[2]
+	if loadFile == "-" {
+		loadFile = os.Getenv("GENERATOR_LOAD")
+		if loadFile == "" {
+			panic("$GENERATOR_LOAD not set")
+		}
+	}
 
 	c, err := client.GetClient(envFile)
 	if err != nil {

trafficgenerator-v2.Dockerfile

@@ -1,20 +1,24 @@
-FROM golang:1.21.13-alpine3.20 as builder
+FROM golang:1.21.13-alpine3.20 AS builder
 
 RUN apk add --no-cache make musl-dev linux-headers gcc git jq bash
 
 WORKDIR /app
 
+RUN apk add --no-cache make tree
+
 # Copy Entire Repo here in order to not copy individual dependencies
 COPY . .
 
-WORKDIR /app/tools/traffic
+WORKDIR /app/test/v2
 
+RUN pwd
+RUN tree -L 2
 RUN --mount=type=cache,target=/go/pkg/mod \
     --mount=type=cache,target=/root/.cache/go-build \
-    go build -o ./bin/generator ./cmd2
+    make build
 
 FROM alpine:3.18 AS generator2
 
-COPY --from=builder /app/tools/traffic/bin/generator /usr/local/bin
+COPY --from=builder /app/test/v2/bin/load /usr/local/bin
 
 ENTRYPOINT ["generator"]