Add storage initialization + read bytecode from file to wrap-contract

cmd/retest/retest.go

@@ -330,7 +330,7 @@ func EthTestDataToString(data EthTestData) string {
 	case reflect.Float64:
 		// We have few tests with numeric code, ex:
 		// "code": 16449,
-		return processStorageData(data.(float64))
+		return util.GetHexString(data.(float64))
 
 	default:
 		log.Fatal().Any("input", data).Str("kind", v.Kind().String()).Msg("Attempted to convert unknown type to raw data")
@@ -378,25 +378,6 @@ func processTestDataString(data string) string {
 	return ""
 }
 
-func processStorageData(data any) string {
-	var result string
-	if reflect.TypeOf(data).Kind() == reflect.Float64 {
-		result = fmt.Sprintf("%x", int64(data.(float64)))
-	} else if reflect.TypeOf(data).Kind() == reflect.String {
-		if strings.HasPrefix(data.(string), "0x") {
-			result = strings.TrimPrefix(data.(string), "0x")
-		} else {
-			result = data.(string)
-		}
-	} else {
-		log.Fatal().Any("data", data).Msg("unknown storage data type")
-	}
-	if len(result) % 2 != 0 {
-		result = "0" + result
-	}
-	return result
-}
-
 // isStandardSolidityString will do a rough check to see if the string looks like a typical solidity file rather than
 // the contracts that are usually in the retest code base
 func isStandardSolidityString(contract string) bool {
@@ -676,8 +657,8 @@ func storageToByteCode(storage map[string]EthTestNumeric) string {
 			log.Warn().Str("slot", slot).Msg("found a storage entry for invalid slot")
 		}
 
-		s := processStorageData(slot)
-		v := processStorageData(value)
+		s := util.GetHexString(slot)
+		v := util.GetHexString(value)
 		sLen := len(s) / 2
 		vLen := len(v) / 2
 		sPushCode := 0x5F + sLen

cmd/wrapcontract/usage.md

@@ -2,9 +2,26 @@ This command takes the runtime bytecode, the bytecode deployed on-chain, as inpu
 
 ```bash
 $ polycli wrap-contract 69602a60005260206000f3600052600a6016f3
+$ echo 69602a60005260206000f3600052600a6016f3 | polycli wrap-contract 
 
 ```
 
+You can also provide a path to a file, and the bytecode while be read from there.
+
+```bash
+$ polycli wrap-contract bytecode.txt
+$ polycli wrap-contract ../bytecode.txt
+$ polycli wrap-contract /tmp/bytecode.txt
+$ echo /tmp/bytecode.txt | polycli wrap-contract
+```
+
+Additionally, you can provide storage for the contract in JSON
+```bash
+$ polycli wrap-contract 0x4455 --storage '{"0x01":"0x0034"}'
+$ polycli wrap-contract 0x4455 --storage '{"0x01":"0x0034", "0x02": "0xFF"}'
+$ echo 69602a60005260206000f3600052600a6016f3 | polycli wrap-contract --storage '{"0x01":"0x0034", "0x02": "0xFF"}'
+```
+
 The resulting bytecode will be formatted this way:
 
 		0x??   // storage initialization code if any

cmd/wrapcontract/wrapcontract.go

@@ -1,33 +1,96 @@
 package wrapcontract
 
 import (
-	"fmt"
 	_ "embed"
-	"github.com/spf13/cobra"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+	"encoding/json"
+
 	"github.com/0xPolygon/polygon-cli/util"
+	"github.com/spf13/cobra"
 )
 
 var (
 	//go:embed usage.md
-	usage                  string
+	usage string
+	jsonStorage *string
 )
 
 var WrapContractCmd = &cobra.Command{
-	Use:   "wrap-contract bytecode",
+	Use:   "wrap-contract bytecode|file",
 	Aliases: []string{"wrapcontract", "wrapContract"},
 	Short: "Wrap deployed bytecode into create bytecode.",
 	Long:  usage,
 	RunE: func(cmd *cobra.Command, args []string) error {
-		deployed_bytecode := args[0]
-		storage_bytecode := ""
-		create_bytecode := util.WrapDeployedCode(deployed_bytecode, storage_bytecode)
-		fmt.Println(create_bytecode)
+		deployedBytecode, err := getInputData(args)
+		if err != nil {
+			cmd.PrintErrf("There was an error reading input for wrapping contract: %s", err.Error())
+			return err
+		}
+		storageBytecode, err := getStorageBytecode()
+		if err != nil {
+			cmd.PrintErrf("There was an error reading storage map: %s", err.Error())
+		}
+		createBytecode := util.WrapDeployedCode(deployedBytecode, storageBytecode)
+		fmt.Println(createBytecode)
 		return nil
 	},
 	Args: func(cmd *cobra.Command, args []string) error {
-		if len(args) != 1 {
-			return fmt.Errorf("expected exactly one argument: bytecode")
+		if len(args) > 1 {
+			return fmt.Errorf("expected at most one argument: bytecode")
 		}
 		return nil
 	},
 }
+
+func init() {
+	flagSet := WrapContractCmd.PersistentFlags()
+	jsonStorage = flagSet.String("storage", "", "Provide storage slots in json format k:v")
+}
+
+func getInputData(args []string) (string, error) {
+	var deployedBytecode string
+	if len(args) == 0 {
+		deployedBytecodeBytes, err := io.ReadAll(os.Stdin)
+		if err != nil {
+			return "", err
+		}
+		deployedBytecode = string(deployedBytecodeBytes)
+	} else {
+		deployedBytecodeOrFile := args[0]
+		// Try to open the param as a file, otherwise treat it as bytecode
+		deployedBytecodeBytes, err := os.ReadFile(deployedBytecodeOrFile)
+		if err != nil {
+			deployedBytecode = deployedBytecodeOrFile
+		} else {
+			deployedBytecode = string(deployedBytecodeBytes)
+		}
+	}
+
+	return strings.TrimSpace(deployedBytecode), nil
+}
+
+func getStorageBytecode() (string, error) {
+	var storageBytecode string = ""
+
+	if jsonStorage != nil && *jsonStorage != "" {
+		var storage map[string]string
+		err := json.Unmarshal([]byte(*jsonStorage), &storage)
+		if err != nil {
+			return storageBytecode, err
+		}
+		for k, v := range storage {
+			slot := util.GetHexString(k)
+			value := util.GetHexString(v)
+			sLen := len(slot) / 2
+			vLen := len(value) / 2
+			sPushCode := 0x5f + sLen
+			vPushCode := 0x5f + vLen
+			storageBytecode += fmt.Sprintf("%02x%s%02x%s55", vPushCode, value, sPushCode, slot)
+		}
+	}
+
+	return storageBytecode, nil
+}

doc/polycli_wrap-contract.md

@@ -14,7 +14,7 @@
 Wrap deployed bytecode into create bytecode.
 
 ```bash
-polycli wrap-contract bytecode [flags]
+polycli wrap-contract bytecode|file [flags]
 ```
 
 ## Usage
@@ -23,9 +23,26 @@ This command takes the runtime bytecode, the bytecode deployed on-chain, as inpu
 
 ```bash
 $ polycli wrap-contract 69602a60005260206000f3600052600a6016f3
+$ echo 69602a60005260206000f3600052600a6016f3 | polycli wrap-contract 
 
 ```
 
+You can also provide a path to a file, and the bytecode while be read from there.
+
+```bash
+$ polycli wrap-contract bytecode.txt
+$ polycli wrap-contract ../bytecode.txt
+$ polycli wrap-contract /tmp/bytecode.txt
+$ echo /tmp/bytecode.txt | polycli wrap-contract
+```
+
+Additionally, you can provide storage for the contract in JSON
+```bash
+$ polycli wrap-contract 0x4455 --storage '{"0x01":"0x0034"}'
+$ polycli wrap-contract 0x4455 --storage '{"0x01":"0x0034", "0x02": "0xFF"}'
+$ echo 69602a60005260206000f3600052600a6016f3 | polycli wrap-contract --storage '{"0x01":"0x0034", "0x02": "0xFF"}'
+```
+
 The resulting bytecode will be formatted this way:
 
 		0x??   // storage initialization code if any
@@ -41,7 +58,8 @@ The resulting bytecode will be formatted this way:
 ## Flags
 
 ```bash
-  -h, --help   help for wrap-contract
+  -h, --help             help for wrap-contract
+      --storage string   Provide storage slots in json format k:v
 ```
 
 The command also inherits flags from parent commands.

util/util.go

@@ -7,6 +7,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+	"reflect"
 
 	"github.com/cenkalti/backoff"
 	"github.com/ethereum/go-ethereum/common/hexutil"
@@ -337,3 +338,22 @@ func WrapDeployedCode(deployedBytecode string, storageBytecode string) string {
 		"%s",			// CODE starts here.
 		storageBytecode, codeCopySize, codeCopyOffset, codeCopySize, deployedBytecode)
 }
+
+func GetHexString(data any) string {
+	var result string
+	if reflect.TypeOf(data).Kind() == reflect.Float64 {
+		result = fmt.Sprintf("%x", int64(data.(float64)))
+	} else if reflect.TypeOf(data).Kind() == reflect.String {
+		if strings.HasPrefix(data.(string), "0x") {
+			result = strings.TrimPrefix(data.(string), "0x")
+		} else {
+			result = data.(string)
+		}
+	} else {
+		log.Fatal().Any("data", data).Msg("unknown storage data type")
+	}
+	if len(result) % 2 != 0 {
+		result = "0" + result
+	}
+	return strings.ToLower(result)
+}