1 files changed, 160 insertions, 0 deletions
diff --git a/p2p/enode/idscheme.go b/p2p/enode/idscheme.go
new file mode 100644
index 000000000..9b495fd4f
--- /dev/null
+++ b/p2p/enode/idscheme.go
@@ -0,0 +1,160 @@
+// Copyright 2018 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package enode
+
+import (
+	"crypto/ecdsa"
+	"fmt"
+	"io"
+
+	"github.com/ethereum/go-ethereum/common/math"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/crypto/sha3"
+	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/ethereum/go-ethereum/rlp"
+)
+
+// List of known secure identity schemes.
+var ValidSchemes = enr.SchemeMap{
+	"v4": V4ID{},
+}
+
+var ValidSchemesForTesting = enr.SchemeMap{
+	"v4":   V4ID{},
+	"null": NullID{},
+}
+
+// v4ID is the "v4" identity scheme.
+type V4ID struct{}
+
+// SignV4 signs a record using the v4 scheme.
+func SignV4(r *enr.Record, privkey *ecdsa.PrivateKey) error {
+	// Copy r to avoid modifying it if signing fails.
+	cpy := *r
+	cpy.Set(enr.ID("v4"))
+	cpy.Set(Secp256k1(privkey.PublicKey))
+
+	h := sha3.NewKeccak256()
+	rlp.Encode(h, cpy.AppendElements(nil))
+	sig, err := crypto.Sign(h.Sum(nil), privkey)
+	if err != nil {
+		return err
+	}
+	sig = sig[:len(sig)-1] // remove v
+	if err = cpy.SetSig(V4ID{}, sig); err == nil {
+		*r = cpy
+	}
+	return err
+}
+
+func (V4ID) Verify(r *enr.Record, sig []byte) error {
+	var entry s256raw
+	if err := r.Load(&entry); err != nil {
+		return err
+	} else if len(entry) != 33 {
+		return fmt.Errorf("invalid public key")
+	}
+
+	h := sha3.NewKeccak256()
+	rlp.Encode(h, r.AppendElements(nil))
+	if !crypto.VerifySignature(entry, h.Sum(nil), sig) {
+		return enr.ErrInvalidSig
+	}
+	return nil
+}
+
+func (V4ID) NodeAddr(r *enr.Record) []byte {
+	var pubkey Secp256k1
+	err := r.Load(&pubkey)
+	if err != nil {
+		return nil
+	}
+	buf := make([]byte, 64)
+	math.ReadBits(pubkey.X, buf[:32])
+	math.ReadBits(pubkey.Y, buf[32:])
+	return crypto.Keccak256(buf)
+}
+
+// Secp256k1 is the "secp256k1" key, which holds a public key.
+type Secp256k1 ecdsa.PublicKey
+
+func (v Secp256k1) ENRKey() string { return "secp256k1" }
+
+// EncodeRLP implements rlp.Encoder.
+func (v Secp256k1) EncodeRLP(w io.Writer) error {
+	return rlp.Encode(w, crypto.CompressPubkey((*ecdsa.PublicKey)(&v)))
+}
+
+// DecodeRLP implements rlp.Decoder.
+func (v *Secp256k1) DecodeRLP(s *rlp.Stream) error {
+	buf, err := s.Bytes()
+	if err != nil {
+		return err
+	}
+	pk, err := crypto.DecompressPubkey(buf)
+	if err != nil {
+		return err
+	}
+	*v = (Secp256k1)(*pk)
+	return nil
+}
+
+// s256raw is an unparsed secp256k1 public key entry.
+type s256raw []byte
+
+func (s256raw) ENRKey() string { return "secp256k1" }
+
+// v4CompatID is a weaker and insecure version of the "v4" scheme which only checks for the
+// presence of a secp256k1 public key, but doesn't verify the signature.
+type v4CompatID struct {
+	V4ID
+}
+
+func (v4CompatID) Verify(r *enr.Record, sig []byte) error {
+	var pubkey Secp256k1
+	return r.Load(&pubkey)
+}
+
+func signV4Compat(r *enr.Record, pubkey *ecdsa.PublicKey) {
+	r.Set((*Secp256k1)(pubkey))
+	if err := r.SetSig(v4CompatID{}, []byte{}); err != nil {
+		panic(err)
+	}
+}
+
+// NullID is the "null" ENR identity scheme. This scheme stores the node
+// ID in the record without any signature.
+type NullID struct{}
+
+func (NullID) Verify(r *enr.Record, sig []byte) error {
+	return nil
+}
+
+func (NullID) NodeAddr(r *enr.Record) []byte {
+	var id ID
+	r.Load(enr.WithEntry("nulladdr", &id))
+	return id[:]
+}
+
+func SignNull(r *enr.Record, id ID) *Node {
+	r.Set(enr.ID("null"))
+	r.Set(enr.WithEntry("nulladdr", id))
+	if err := r.SetSig(NullID{}, []byte{}); err != nil {
+		panic(err)
+	}
+	return &Node{r: *r, id: id}
+}