From fd600615e9982cdb3a49d3104fc2082b751f0e4f Mon Sep 17 00:00:00 2001
From: Dan Forbes <dan@danforbes.dev>
Date: Fri, 21 Jan 2022 15:24:49 -0800
Subject: [PATCH] feat(lib/runtime): support Substrate WASM compression (#2213)

* Support Substrate Wasm compression

https://github.com/paritytech/substrate/blob/master/primitives/maybe-compressed-blob/src/lib.rs

* Apply suggestions from code review

Co-authored-by: noot <36753753+noot@users.noreply.github.com>

* Review comments

- slices cannot be const
- create function & write tests
- `go mod tidy`

* Apply suggestions from code review

Co-authored-by: Quentin McGaw <quentin.mcgaw@gmail.com>

Co-authored-by: noot <36753753+noot@users.noreply.github.com>
Co-authored-by: Quentin McGaw <quentin.mcgaw@gmail.com>
---
 go.mod                              |  2 +-
 lib/runtime/wasmer/instance.go      | 25 ++++++++++++++++++
 lib/runtime/wasmer/instance_test.go | 40 +++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 120ef2e983..4f192034d3 100644
--- a/go.mod
+++ b/go.mod
@@ -27,6 +27,7 @@ require (
 	github.com/ipfs/go-ds-badger2 v0.1.1
 	github.com/ipfs/go-ipns v0.1.2 //indirect
 	github.com/jpillora/ipfilter v1.2.3
+	github.com/klauspost/compress v1.12.3
 	github.com/libp2p/go-libp2p v0.15.1
 	github.com/libp2p/go-libp2p-core v0.9.0
 	github.com/libp2p/go-libp2p-discovery v0.5.1
@@ -84,7 +85,6 @@ require (
 	github.com/jbenet/goprocess v0.1.4 // indirect
 	github.com/jcelliott/lumber v0.0.0-20160324203708-dd349441af25 // indirect
 	github.com/jpillora/backoff v1.0.0 // indirect
-	github.com/klauspost/compress v1.12.3 // indirect
 	github.com/klauspost/cpuid/v2 v2.0.9 // indirect
 	github.com/koron/go-ssdp v0.0.2 // indirect
 	github.com/leodido/go-urn v1.2.1 // indirect
diff --git a/lib/runtime/wasmer/instance.go b/lib/runtime/wasmer/instance.go
index 7c87ceb287..da3b5441a9 100644
--- a/lib/runtime/wasmer/instance.go
+++ b/lib/runtime/wasmer/instance.go
@@ -4,6 +4,7 @@
 package wasmer
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"sync"
@@ -18,6 +19,8 @@ import (
 	"github.com/ChainSafe/gossamer/lib/crypto"
 
 	wasm "github.com/wasmerio/go-ext-wasm/wasmer"
+
+	"github.com/klauspost/compress/zstd"
 )
 
 // Name represents the name of the interpreter
@@ -94,6 +97,12 @@ func NewInstance(code []byte, cfg *Config) (*Instance, error) {
 		return nil, errors.New("code is empty")
 	}
 
+	var err error
+	code, err = decompressWasm(code)
+	if err != nil {
+		return nil, fmt.Errorf("cannot decompress WASM code: %w", err)
+	}
+
 	logger.Patch(log.SetLevel(cfg.LogLvl), log.SetCallerFunc(true))
 
 	imports, err := cfg.Imports()
@@ -157,6 +166,22 @@ func NewInstance(code []byte, cfg *Config) (*Instance, error) {
 	return inst, nil
 }
 
+// decompressWasm decompresses a Wasm blob that may or may not be compressed with zstd
+// ref: https://github.com/paritytech/substrate/blob/master/primitives/maybe-compressed-blob/src/lib.rs
+func decompressWasm(code []byte) ([]byte, error) {
+	compressionFlag := []byte{82, 188, 83, 118, 70, 219, 142, 5}
+	if !bytes.HasPrefix(code, compressionFlag) {
+		return code, nil
+	}
+
+	decoder, err := zstd.NewReader(nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create zstd decoder: %w", err)
+	}
+
+	return decoder.DecodeAll(code[len(compressionFlag):], nil)
+}
+
 // GetCodeHash returns the code of the instance
 func (in *Instance) GetCodeHash() common.Hash {
 	return in.codeHash
diff --git a/lib/runtime/wasmer/instance_test.go b/lib/runtime/wasmer/instance_test.go
index 08e071223e..89b95d70eb 100644
--- a/lib/runtime/wasmer/instance_test.go
+++ b/lib/runtime/wasmer/instance_test.go
@@ -10,6 +10,8 @@ import (
 
 	"github.com/ChainSafe/gossamer/lib/runtime"
 	"github.com/stretchr/testify/require"
+
+	"github.com/klauspost/compress/zstd"
 )
 
 // test used for ensuring runtime exec calls can me made concurrently
@@ -63,3 +65,41 @@ func TestInstance_CheckRuntimeVersion(t *testing.T) {
 	require.Equal(t, expected.ImplVersion(), version.ImplVersion())
 	require.Equal(t, expected.TransactionVersion(), version.TransactionVersion())
 }
+
+func TestDecompressWasm(t *testing.T) {
+	encoder, err := zstd.NewWriter(nil)
+	require.NoError(t, err)
+
+	cases := []struct {
+		in       []byte
+		expected []byte
+		msg      string
+	}{
+		{
+			[]byte{82, 188, 83, 118, 70, 219, 142},
+			[]byte{82, 188, 83, 118, 70, 219, 142},
+			"partial compression flag",
+		},
+		{
+			[]byte{82, 188, 83, 118, 70, 219, 142, 6},
+			[]byte{82, 188, 83, 118, 70, 219, 142, 6},
+			"wrong compression flag",
+		},
+		{
+			[]byte{82, 188, 83, 118, 70, 219, 142, 6, 221},
+			[]byte{82, 188, 83, 118, 70, 219, 142, 6, 221},
+			"wrong compression flag with data",
+		},
+		{
+			append([]byte{82, 188, 83, 118, 70, 219, 142, 5}, encoder.EncodeAll([]byte("compressed"), nil)...),
+			[]byte("compressed"),
+			"compressed data",
+		},
+	}
+
+	for _, test := range cases {
+		actual, err := decompressWasm(test.in)
+		require.NoError(t, err)
+		require.Equal(t, test.expected, actual)
+	}
+}