Fix review changes

x/wasm/client/cli/tx.go

@@ -25,6 +25,9 @@ const (
 	flagAmount = "amount"
 )
 
+// limit max bytes read to prevent gzip bombs
+const maxSize = 400 * 1024
+
 // GetTxCmd returns the transaction commands for this module
 func GetTxCmd(cdc *codec.Codec) *cobra.Command {
 	txCmd := &cobra.Command{
@@ -59,6 +62,12 @@ func StoreCodeCmd(cdc *codec.Codec) *cobra.Command {
 				return err
 			}
 
+			// limit the input size
+			if len(wasm) > maxSize {
+				return fmt.Errorf("input size exceeds the max size allowed (allowed:%d, actual: %d)",
+					maxSize, len(wasm))
+			}
+
 			// gzip the wasm file
 			if wasmUtils.IsWasm(wasm) {
 				wasm, err = wasmUtils.GzipIt(wasm)
@@ -67,7 +76,7 @@ func StoreCodeCmd(cdc *codec.Codec) *cobra.Command {
 					return err
 				}
 			} else if !wasmUtils.IsGzip(wasm) {
-				return fmt.Errorf("invalid input file. Accepts only wasm binary or gzip %s")
+				return fmt.Errorf("invalid input file. Use wasm binary or gzip")
 			}
 
 			// build and sign the transaction, then broadcast to Tendermint

x/wasm/client/utils/utils.go

@@ -3,7 +3,6 @@ package utils
 import (
 	"bytes"
 	"compress/gzip"
-	"io"
 )
 
 var (
@@ -11,43 +10,14 @@ var (
 	wasmIdent = []byte("\x00\x61\x73\x6D")
 )
 
-// limit max bytes read to prevent gzip bombs
-const maxSize = 400 * 1024
-
 // IsGzip returns checks if the file contents are gzip compressed
 func IsGzip(input []byte) bool {
-	if len(input) < 3 {
-		return false
-	}
-
-	in := io.LimitReader(bytes.NewReader(input), maxSize)
-	buf := make([]byte, 3)
-
-	if _, err := io.ReadAtLeast(in, buf, 3); err != nil {
-		return false
-	}
-
-	return bytes.Equal(gzipIdent, buf)
+	return bytes.Equal(input[:3], gzipIdent)
 }
 
 // IsWasm checks if the file contents are of wasm binary
 func IsWasm(input []byte) bool {
-	if len(input) < 3 {
-		return false
-	}
-
-	in := io.LimitReader(bytes.NewReader(input), maxSize)
-	buf := make([]byte, 4)
-
-	if _, err := io.ReadAtLeast(in, buf, 4); err != nil {
-		return false
-	}
-
-	if bytes.Equal(wasmIdent, buf) {
-		return true
-	}
-
-	return false
+	return bytes.Equal(input[:4], wasmIdent)
 }
 
 // GzipIt compresses the input ([]byte)

x/wasm/client/utils/utils_test.go

@@ -0,0 +1,61 @@
+package utils
+
+import (
+	"github.com/stretchr/testify/require"
+	"io/ioutil"
+	"testing"
+)
+
+func GetTestData() ([]byte, []byte, []byte, error){
+	wasmCode, err := ioutil.ReadFile("../../internal/keeper/testdata/contract.wasm")
+
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	gzipData, err := GzipIt(wasmCode)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
+	someRandomStr := []byte("hello world")
+
+	return wasmCode, someRandomStr, gzipData, nil
+}
+
+func TestIsWasm (t *testing.T) {
+	wasmCode, someRandomStr, gzipData, err := GetTestData()
+	require.NoError(t, err)
+
+	t.Log("should return false for some random string data")
+	require.False(t, IsWasm(someRandomStr))
+	t.Log("should return false for gzip data")
+	require.False(t, IsWasm(gzipData))
+	t.Log("should return true for exact wasm")
+	require.True(t, IsWasm(wasmCode))
+}
+
+func TestIsGzip (t *testing.T) {
+	wasmCode, someRandomStr, gzipData, err := GetTestData()
+	require.NoError(t, err)
+
+	require.False(t, IsGzip(wasmCode))
+	require.False(t, IsGzip(someRandomStr))
+	require.True(t, IsGzip(gzipData))
+}
+
+func TestGzipIt (t *testing.T) {
+	wasmCode, someRandomStr, _, err := GetTestData()
+	require.NoError(t, err)
+
+	t.Log("gzip wasm with no error")
+	_, err = GzipIt(wasmCode)
+	require.NoError(t, err)
+
+	t.Log("gzip of a string should return exact gzip data")
+	strToGzip, err := GzipIt(someRandomStr)
+	originalGzipData := []byte{31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 202, 72, 205, 201, 201, 87, 40, 207, 47, 202, 73, 1,
+		4, 0, 0, 255, 255, 133, 17, 74, 13, 11, 0, 0, 0}
+	require.NoError(t, err)
+	require.Equal(t, originalGzipData, strToGzip)
+}