1 files changed, 95 insertions, 73 deletions
diff --git a/p2p/enr/enr.go b/p2p/enr/enr.go
index 48683471d..251caf458 100644
--- a/p2p/enr/enr.go
+++ b/p2p/enr/enr.go
@@ -15,14 +15,20 @@
 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
 
 // Package enr implements Ethereum Node Records as defined in EIP-778. A node record holds
-// arbitrary information about a node on the peer-to-peer network.
-//
-// Records contain named keys. To store and retrieve key/values in a record, use the Entry
+// arbitrary information about a node on the peer-to-peer network. Node information is
+// stored in key/value pairs. To store and retrieve key/values in a record, use the Entry
 // interface.
 //
-// Records must be signed before transmitting them to another node. Decoding a record verifies
-// its signature. When creating a record, set the entries you want, then call Sign to add the
-// signature. Modifying a record invalidates the signature.
+// Signature Handling
+//
+// Records must be signed before transmitting them to another node.
+//
+// Decoding a record doesn't check its signature. Code working with records from an
+// untrusted source must always verify two things: that the record uses an identity scheme
+// deemed secure, and that the signature is valid according to the declared scheme.
+//
+// When creating a record, set the entries you want and use a signing function provided by
+// the identity scheme to add the signature. Modifying a record invalidates the signature.
 //
 // Package enr supports the "secp256k1-keccak" identity scheme.
 package enr
@@ -40,8 +46,7 @@ import (
 const SizeLimit = 300 // maximum encoded size of a node record in bytes
 
 var (
-	errNoID           = errors.New("unknown or unspecified identity scheme")
-	errInvalidSig     = errors.New("invalid signature")
+	ErrInvalidSig     = errors.New("invalid signature on node record")
 	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")
@@ -50,6 +55,32 @@ var (
 	errNotFound       = errors.New("no such key in record")
 )
 
+// An IdentityScheme is capable of verifying record signatures and
+// deriving node addresses.
+type IdentityScheme interface {
+	Verify(r *Record, sig []byte) error
+	NodeAddr(r *Record) []byte
+}
+
+// SchemeMap is a registry of named identity schemes.
+type SchemeMap map[string]IdentityScheme
+
+func (m SchemeMap) Verify(r *Record, sig []byte) error {
+	s := m[r.IdentityScheme()]
+	if s == nil {
+		return ErrInvalidSig
+	}
+	return s.Verify(r, sig)
+}
+
+func (m SchemeMap) NodeAddr(r *Record) []byte {
+	s := m[r.IdentityScheme()]
+	if s == nil {
+		return nil
+	}
+	return s.NodeAddr(r)
+}
+
 // Record represents a node record. The zero value is an empty record.
 type Record struct {
 	seq       uint64 // sequence number
@@ -64,11 +95,6 @@ type pair struct {
 	v rlp.RawValue
 }
 
-// Signed reports whether the record has a valid signature.
-func (r *Record) Signed() bool {
-	return r.signature != nil
-}
-
 // Seq returns the sequence number.
 func (r *Record) Seq() uint64 {
 	return r.seq
@@ -140,7 +166,7 @@ func (r *Record) invalidate() {
 // EncodeRLP implements rlp.Encoder. Encoding fails if
 // the record is unsigned.
 func (r Record) EncodeRLP(w io.Writer) error {
-	if !r.Signed() {
+	if r.signature == nil {
 		return errEncodeUnsigned
 	}
 	_, err := w.Write(r.raw)
@@ -149,25 +175,34 @@ func (r Record) EncodeRLP(w io.Writer) error {
 
 // DecodeRLP implements rlp.Decoder. Decoding verifies the signature.
 func (r *Record) DecodeRLP(s *rlp.Stream) error {
-	raw, err := s.Raw()
+	dec, raw, err := decodeRecord(s)
 	if err != nil {
 		return err
 	}
+	*r = dec
+	r.raw = raw
+	return nil
+}
+
+func decodeRecord(s *rlp.Stream) (dec Record, raw []byte, err error) {
+	raw, err = s.Raw()
+	if err != nil {
+		return dec, raw, err
+	}
 	if len(raw) > SizeLimit {
-		return errTooBig
+		return dec, raw, errTooBig
 	}
 
 	// Decode the RLP container.
-	dec := Record{raw: raw}
 	s = rlp.NewStream(bytes.NewReader(raw), 0)
 	if _, err := s.List(); err != nil {
-		return err
+		return dec, raw, err
 	}
 	if err = s.Decode(&dec.signature); err != nil {
-		return err
+		return dec, raw, err
 	}
 	if err = s.Decode(&dec.seq); err != nil {
-		return err
+		return dec, raw, err
 	}
 	// The rest of the record contains sorted k/v pairs.
 	var prevkey string
@@ -177,73 +212,68 @@ func (r *Record) DecodeRLP(s *rlp.Stream) error {
 			if err == rlp.EOL {
 				break
 			}
-			return err
+			return dec, raw, err
 		}
 		if err := s.Decode(&kv.v); err != nil {
 			if err == rlp.EOL {
-				return errIncompletePair
+				return dec, raw, errIncompletePair
 			}
-			return err
+			return dec, raw, err
 		}
 		if i > 0 {
 			if kv.k == prevkey {
-				return errDuplicateKey
+				return dec, raw, errDuplicateKey
 			}
 			if kv.k < prevkey {
-				return errNotSorted
+				return dec, raw, errNotSorted
 			}
 		}
 		dec.pairs = append(dec.pairs, kv)
 		prevkey = kv.k
 	}
-	if err := s.ListEnd(); err != nil {
-		return err
-	}
+	return dec, raw, s.ListEnd()
+}
 
-	_, scheme := dec.idScheme()
-	if scheme == nil {
-		return errNoID
-	}
-	if err := scheme.Verify(&dec, dec.signature); err != nil {
-		return err
-	}
-	*r = dec
-	return nil
+// IdentityScheme returns the name of the identity scheme in the record.
+func (r *Record) IdentityScheme() string {
+	var id ID
+	r.Load(&id)
+	return string(id)
 }
 
-// NodeAddr returns the node address. The return value will be nil if the record is
-// unsigned or uses an unknown identity scheme.
-func (r *Record) NodeAddr() []byte {
-	_, scheme := r.idScheme()
-	if scheme == nil {
-		return nil
-	}
-	return scheme.NodeAddr(r)
+// VerifySignature checks whether the record is signed using the given identity scheme.
+func (r *Record) VerifySignature(s IdentityScheme) error {
+	return s.Verify(r, r.signature)
 }
 
 // SetSig sets the record signature. It returns an error if the encoded record is larger
 // than the size limit or if the signature is invalid according to the passed scheme.
-func (r *Record) SetSig(idscheme string, sig []byte) error {
-	// Check that "id" is set and matches the given scheme. This panics because
-	// inconsitencies here are always implementation bugs in the signing function calling
-	// this method.
-	id, s := r.idScheme()
-	if s == nil {
-		panic(errNoID)
-	}
-	if id != idscheme {
-		panic(fmt.Errorf("identity scheme mismatch in Sign: record has %s, want %s", id, idscheme))
-	}
-
-	// Verify against the scheme.
-	if err := s.Verify(r, sig); err != nil {
-		return err
-	}
-	raw, err := r.encode(sig)
-	if err != nil {
-		return err
+//
+// You can also use SetSig to remove the signature explicitly by passing a nil scheme
+// and signature.
+//
+// SetSig panics when either the scheme or the signature (but not both) are nil.
+func (r *Record) SetSig(s IdentityScheme, sig []byte) error {
+	switch {
+	// Prevent storing invalid data.
+	case s == nil && sig != nil:
+		panic("enr: invalid call to SetSig with non-nil signature but nil scheme")
+	case s != nil && sig == nil:
+		panic("enr: invalid call to SetSig with nil signature but non-nil scheme")
+	// Verify if we have a scheme.
+	case s != nil:
+		if err := s.Verify(r, sig); err != nil {
+			return err
+		}
+		raw, err := r.encode(sig)
+		if err != nil {
+			return err
+		}
+		r.signature, r.raw = sig, raw
+	// Reset otherwise.
+	default:
+		r.signature, r.raw = nil, nil
 	}
-	r.signature, r.raw = sig, raw
 	return nil
 }
 
@@ -268,11 +298,3 @@ func (r *Record) encode(sig []byte) (raw []byte, err error) {
 	}
 	return raw, nil
 }
-
-func (r *Record) idScheme() (string, IdentityScheme) {
-	var id ID
-	if err := r.Load(&id); err != nil {
-		return "", nil
-	}
-	return string(id), FindIdentityScheme(string(id))
-}