From 301ac8b8fdd931a86c888e162a7b2a11a2346989 Mon Sep 17 00:00:00 2001
From: Anton Evangelatov <anton.evangelatov@gmail.com>
Date: Wed, 6 Dec 2017 12:43:08 +0100
Subject: [PATCH] p2p/enr: check for size limit in Sign; WithKey(k, v) for
 adding random key/values

---
 p2p/enr/enr.go      | 11 ++++++++++-
 p2p/enr/enr_test.go | 37 +++++++++++++++++++++++++++++++++++++
 p2p/enr/keys.go     | 21 +++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/p2p/enr/enr.go b/p2p/enr/enr.go
index 79ea3bf536..df6ce8c85d 100644
--- a/p2p/enr/enr.go
+++ b/p2p/enr/enr.go
@@ -40,6 +40,7 @@ var (
 	errNotSorted      = errors.New("record key/value pairs are not sorted by key")
 	errDuplicateKey   = errors.New("record contains duplicate key")
 	errIncompletePair = errors.New("record contains incomplete k/v pair")
+	errTooBig         = errors.New("record bigger than 300 bytes")
 )
 
 // Key is implemented by known node record key types.
@@ -221,7 +222,15 @@ func (r *Record) signAndEncode(privkey *ecdsa.PrivateKey) error {
 	r.signature = encodeCompactSignature(sig)
 	list[0] = r.signature
 	r.raw, err = rlp.EncodeToBytes(list)
-	return err
+	if err != nil {
+		return err
+	}
+
+	if len(r.raw) > 300 {
+		return errTooBig
+	}
+
+	return nil
 }
 
 func (r *Record) verifySignature() error {
diff --git a/p2p/enr/enr_test.go b/p2p/enr/enr_test.go
index 6dd8ae3abd..52ccac7391 100644
--- a/p2p/enr/enr_test.go
+++ b/p2p/enr/enr_test.go
@@ -19,9 +19,11 @@ package enr
 import (
 	"bytes"
 	"encoding/hex"
+	"math/rand"
 	"net"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/ethereum/go-ethereum/crypto"
@@ -268,3 +270,38 @@ func TestPythonInterop(t *testing.T) {
 		}
 	}
 }
+
+func TestRecordTooBig(t *testing.T) {
+	privkey, err := crypto.HexToECDSA(privkeyHex)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var r Record
+
+	key := randomString(10)
+
+	// set a big value for random key, expect error
+	r.Set(WithKey(key, randomString(300)))
+	err = r.Sign(privkey)
+	if err != errTooBig {
+		t.Fatalf("expected to get errTooBig, got %#v", err)
+	}
+
+	// set an acceptable value for random key, expect no error
+	r.Set(WithKey(key, randomString(100)))
+	err = r.Sign(privkey)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func randomString(strlen int) string {
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	const chars = "abcdefghijklmnopqrstuvwxyz0123456789"
+	result := make([]byte, strlen)
+	for i := range result {
+		result[i] = chars[r.Intn(len(chars))]
+	}
+	return string(result)
+}
diff --git a/p2p/enr/keys.go b/p2p/enr/keys.go
index 96f729997c..ef94126803 100644
--- a/p2p/enr/keys.go
+++ b/p2p/enr/keys.go
@@ -26,6 +26,27 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
+type generic struct {
+	key   string
+	value interface{}
+}
+
+func (g generic) ENRKey() string {
+	return g.key
+}
+
+func (g generic) EncodeRLP(w io.Writer) error {
+	return rlp.Encode(w, g.value)
+}
+
+func (g *generic) DecodeRLP(s *rlp.Stream) error {
+	return s.Decode(&g.value)
+}
+
+func WithKey(k string, v interface{}) Key {
+	return &generic{key: k, value: v}
+}
+
 // DiscPort represents an UDP port for discovery v5.
 type DiscPort uint16