From 36914c5c2c9a91fc7c02223f63ac95eaed0a7fb4 Mon Sep 17 00:00:00 2001
From: Emmanuel T Odeke <emmanuel@orijtech.com>
Date: Wed, 17 Feb 2021 02:13:00 -0800
Subject: [PATCH] codec/types: avoid unnecessary allocations for
 NewAnyWithCustomTypeURL on error (#8605)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Avoids a bleed out attack in which a node can be made to allocate
memory slowly or very fast in small strides, by sending bad data
to code that invokes NewAnyWithCustomTypeURL, in which we
unconditionally returned a new Any object. On a 64-bit machine,
this would waste 96 bytes per invocation even on error.

Added a test to ensure zero allocations with a fixed error returned.
Also added a benchmark which shows reduction in wasted allocations and
wasted CPU time:

```shell
$ benchstat before.txt after.txt
name                                            old time/op    new time/op    delta
NewAnyWithCustomTypeURLWithErrorReturned-8      142ns ± 6%      55ns ±12%   -61.65%  (p=0.000 n=9+10)

name                                            old alloc/op   new alloc/op   delta
NewAnyWithCustomTypeURLWithErrorReturned-8      96.0B ± 0%      0.0B       -100.00%  (p=0.000 n=10+10)

name                                            old allocs/op  new allocs/op  delta
NewAnyWithCustomTypeURLWithErrorReturned-8      1.00 ± 0%      0.00       -100.00%  (p=0.000 n=10+10)
```

Fixes #8537

(cherry picked from commit 56fc3fc5726a38e4fd379cc1dc245df6e71ba9f6)
---
 codec/types/any.go      |  5 ++-
 codec/types/any_test.go | 68 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 codec/types/any_test.go

diff --git a/codec/types/any.go b/codec/types/any.go
index eebf13ca6cf0..faa8693d731a 100644
--- a/codec/types/any.go
+++ b/codec/types/any.go
@@ -71,11 +71,14 @@ func NewAnyWithValue(v proto.Message) (*Any, error) {
 // into the protobuf Any serialization. For simple marshaling you should use NewAnyWithValue.
 func NewAnyWithCustomTypeURL(v proto.Message, typeURL string) (*Any, error) {
 	bz, err := proto.Marshal(v)
+	if err != nil {
+		return nil, err
+	}
 	return &Any{
 		TypeUrl:     typeURL,
 		Value:       bz,
 		cachedValue: v,
-	}, err
+	}, nil
 }
 
 // UnsafePackAny packs the value x in the Any and instead of returning the error
diff --git a/codec/types/any_test.go b/codec/types/any_test.go
new file mode 100644
index 000000000000..cea3e0444efb
--- /dev/null
+++ b/codec/types/any_test.go
@@ -0,0 +1,68 @@
+package types_test
+
+import (
+	"fmt"
+	"runtime"
+	"testing"
+
+	"github.com/gogo/protobuf/proto"
+
+	"github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/testutil/testdata"
+)
+
+type errOnMarshal struct {
+	testdata.Dog
+}
+
+var _ proto.Message = (*errOnMarshal)(nil)
+
+var errAlways = fmt.Errorf("always erroring")
+
+func (eom *errOnMarshal) XXX_Marshal(b []byte, deterministic bool) ([]byte, error) {
+	return nil, errAlways
+}
+
+const fauxURL = "/anyhere"
+
+var eom = &errOnMarshal{}
+
+// Ensure that returning an error doesn't suddenly allocate and waste bytes.
+// See https://github.com/cosmos/cosmos-sdk/issues/8537
+func TestNewAnyWithCustomTypeURLWithErrorNoAllocation(t *testing.T) {
+	var ms1, ms2 runtime.MemStats
+	runtime.ReadMemStats(&ms1)
+	any, err := types.NewAnyWithCustomTypeURL(eom, fauxURL)
+	runtime.ReadMemStats(&ms2)
+	// Ensure that no fresh allocation was made.
+	if diff := ms2.HeapAlloc - ms1.HeapAlloc; diff > 0 {
+		t.Errorf("Unexpected allocation of %d bytes", diff)
+	}
+	if err == nil {
+		t.Fatal("err wasn't returned")
+	}
+	if any != nil {
+		t.Fatalf("Unexpectedly got a non-nil Any value: %v", any)
+	}
+}
+
+var sink interface{}
+
+func BenchmarkNewAnyWithCustomTypeURLWithErrorReturned(b *testing.B) {
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		any, err := types.NewAnyWithCustomTypeURL(eom, fauxURL)
+		if err == nil {
+			b.Fatal("err wasn't returned")
+		}
+		if any != nil {
+			b.Fatalf("Unexpectedly got a non-nil Any value: %v", any)
+		}
+		sink = any
+	}
+	if sink == nil {
+		b.Fatal("benchmark didn't run")
+	}
+	sink = (interface{})(nil)
+}