github.com/annchain/OG@v0.0.9/p2p/rlpx.go (about)

     1  // Copyright 2015 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  //go:generate msgp
    18  
    19  package p2p
    20  
    21  import (
    22  	"bytes"
    23  	"crypto/aes"
    24  	"crypto/cipher"
    25  	"crypto/ecdsa"
    26  	"crypto/elliptic"
    27  	"crypto/hmac"
    28  	"crypto/rand"
    29  	"encoding/binary"
    30  	"errors"
    31  	"fmt"
    32  	"github.com/annchain/OG/arefactor/common/goroutine"
    33  	"github.com/annchain/OG/deprecated/ogcrypto"
    34  	"github.com/annchain/OG/p2p/ioperformance"
    35  	"hash"
    36  	"io"
    37  	"io/ioutil"
    38  	mrand "math/rand"
    39  	"net"
    40  	"sync"
    41  	"time"
    42  
    43  	"github.com/annchain/OG/common/bitutil"
    44  	"github.com/annchain/OG/deprecated/ogcrypto/ecies"
    45  	"github.com/golang/snappy"
    46  	"golang.org/x/crypto/sha3"
    47  )
    48  
    49  const (
    50  	maxUint24 = ^uint32(0) >> 8
    51  
    52  	sskLen = 16 // ecies.MaxSharedKeyLength(pubKey) / 2
    53  	sigLen = 65 // elliptic S256
    54  	pubLen = 64 // 512 bit pubkey in uncompressed representation without format byte
    55  	shaLen = 32 // hash length (for nonce etc)
    56  
    57  	authMsgLen  = sigLen + shaLen + pubLen + shaLen + 1
    58  	authRespLen = pubLen + shaLen + 1
    59  
    60  	eciesOverhead = 65 /* pubkey */ + 16 /* IV */ + 32 /* MAC */
    61  
    62  	encAuthMsgLen  = authMsgLen + eciesOverhead  // size of encrypted pre-EIP-8 initiator handshake
    63  	encAuthRespLen = authRespLen + eciesOverhead // size of encrypted pre-EIP-8 handshake reply
    64  
    65  	// total timeout for encryption handshake and protocol
    66  	// handshake in both directions.
    67  	handshakeTimeout = 5 * time.Second
    68  
    69  	// This is the timeout for sending the disconnect reason.
    70  	// This is shorter than the usual timeout because we don't want
    71  	// to wait if the connection is known to be bad anyway.
    72  	discWriteTimeout = 1 * time.Second
    73  )
    74  
    75  // errPlainMessageTooLarge is returned if a decompressed message length exceeds
    76  // the allowed 24 bits (i.e. length >= 16MB).
    77  var errPlainMessageTooLarge = errors.New("message length >= 16MB")
    78  
    79  // rlpx is the transport protocol used by actual (non-test) connections.
    80  // It wraps the frame encoder with locks and read/write deadlines.
    81  type rlpx struct {
    82  	fd net.Conn
    83  
    84  	rmu, wmu sync.Mutex
    85  	rw       *rlpxFrameRW
    86  }
    87  
    88  // RLPx v4 handshake auth (defined in EIP-8).
    89  //msgp:tuple AuthMsgV4
    90  type AuthMsgV4 struct {
    91  	gotPlain bool // whether read packet had plain format.
    92  
    93  	Signature       [sigLen]byte
    94  	InitiatorPubkey [pubLen]byte
    95  	Nonce           [shaLen]byte
    96  	Version         uint32
    97  }
    98  
    99  // RLPx v4 handshake response (defined in EIP-8).
   100  //msgp:tuple AuthRespV4
   101  type AuthRespV4 struct {
   102  	RandomPubkey [pubLen]byte
   103  	Nonce        [shaLen]byte
   104  	Version      uint32
   105  }
   106  
   107  func newRLPX(fd net.Conn) transport {
   108  	fd.SetDeadline(time.Now().Add(handshakeTimeout))
   109  	return &rlpx{fd: fd}
   110  }
   111  
   112  func (t *rlpx) ReadMsg() (Msg, error) {
   113  	t.rmu.Lock()
   114  	defer t.rmu.Unlock()
   115  	t.fd.SetReadDeadline(time.Now().Add(frameReadTimeout))
   116  	return t.rw.ReadMsg()
   117  }
   118  
   119  func (t *rlpx) WriteMsg(msg Msg) error {
   120  	t.wmu.Lock()
   121  	defer t.wmu.Unlock()
   122  	t.fd.SetWriteDeadline(time.Now().Add(frameWriteTimeout))
   123  	return t.rw.WriteMsg(msg)
   124  }
   125  
   126  func (t *rlpx) close(err error) {
   127  	t.wmu.Lock()
   128  	defer t.wmu.Unlock()
   129  	// Tell the remote end why we're disconnecting if possible.
   130  	if t.rw != nil {
   131  		if r, ok := err.(DiscReason); ok && r != DiscNetworkError {
   132  			// rlpx tries to send DiscReason to disconnected peer
   133  			// if the connection is net.Pipe (in-memory simulation)
   134  			// it hangs forever, since net.Pipe does not implement
   135  			// a write deadline. Because of this only try to send
   136  			// the disconnect reason message if there is no error.
   137  			if err := t.fd.SetWriteDeadline(time.Now().Add(discWriteTimeout)); err == nil {
   138  				b, _ := r.MarshalMsg(nil)
   139  				Send(t.rw, discMsg, b)
   140  			}
   141  		}
   142  	}
   143  	t.fd.Close()
   144  }
   145  
   146  func (t *rlpx) doProtoHandshake(our *ProtoHandshake) (their *ProtoHandshake, err error) {
   147  	// Writing our handshake happens concurrently, we prefer
   148  	// returning the handshake read error. If the remote side
   149  	// disconnects us early with a valid reason, we should return it
   150  	// as the error so it can be tracked elsewhere.
   151  	werr := make(chan error, 1)
   152  	b, _ := our.MarshalMsg(nil)
   153  	sendFunc := func() { werr <- Send(t.rw, handshakeMsg, b) }
   154  	goroutine.New(sendFunc)
   155  	if their, err = readProtocolHandshake(t.rw); err != nil {
   156  		<-werr // make sure the write terminates too
   157  		return nil, err
   158  	}
   159  	if err := <-werr; err != nil {
   160  		return nil, fmt.Errorf("write error: %v", err)
   161  	}
   162  	// If the protocol version supports Snappy encoding, upgrade immediately
   163  	t.rw.snappy = their.Version >= snappyProtocolVersion
   164  
   165  	return their, nil
   166  }
   167  
   168  func readProtocolHandshake(rw MsgReader) (*ProtoHandshake, error) {
   169  	msg, err := rw.ReadMsg()
   170  	if err != nil {
   171  		return nil, err
   172  	}
   173  	if msg.Size > baseProtocolMaxMsgSize {
   174  		return nil, fmt.Errorf("message too big")
   175  	}
   176  	if msg.Code == discMsg {
   177  		// Disconnect before protocol handshake is valid according to the
   178  		// spec and we send it ourself if the posthanshake checks fail.
   179  		// We can't return the reason directly, though, because it is echoed
   180  		// back otherwise. Wrap it in a string instead.
   181  		var reason DiscReason
   182  		//rlp.Decode(msg.Payload, &reason)
   183  		b, _ := msg.GetPayLoad()
   184  		reason.UnmarshalMsg(b)
   185  		return nil, reason
   186  	}
   187  	if msg.Code != handshakeMsg {
   188  		return nil, fmt.Errorf("expected handshake, got %x", msg.Code)
   189  	}
   190  	var hs ProtoHandshake
   191  	data, _ := msg.GetPayLoad()
   192  	_, err = hs.UnmarshalMsg(data)
   193  	if err != nil {
   194  		err = newPeerError(errInvalidMsg, "(code %x) (size %d) %v", msg.Code, msg.Size, err)
   195  		//if err := msg.Decode(&hs); err != nil {
   196  		return nil, err
   197  	}
   198  	if len(hs.ID) != 64 || !bitutil.TestBytes(hs.ID) {
   199  		return nil, DiscInvalidIdentity
   200  	}
   201  	ioperformance.AddRecvSize(int(msg.Size))
   202  	return &hs, nil
   203  }
   204  
   205  // doEncHandshake runs the protocol handshake using authenticated
   206  // messages. the protocol handshake is the first authenticated message
   207  // and also verifies whether the encryption handshake 'worked' and the
   208  // remote side actually provided the right public key.
   209  func (t *rlpx) doEncHandshake(prv *ecdsa.PrivateKey, dial *ecdsa.PublicKey) (*ecdsa.PublicKey, error) {
   210  	var (
   211  		sec secrets
   212  		err error
   213  	)
   214  	if dial == nil {
   215  		sec, err = receiverEncHandshake(t.fd, prv)
   216  	} else {
   217  		sec, err = initiatorEncHandshake(t.fd, prv, dial)
   218  	}
   219  	if err != nil {
   220  		return nil, err
   221  	}
   222  	t.wmu.Lock()
   223  	t.rw = newRLPXFrameRW(t.fd, sec)
   224  	t.wmu.Unlock()
   225  	return sec.Remote.ExportECDSA(), nil
   226  }
   227  
   228  // encHandshake contains the state of the encryption handshake.
   229  type encHandshake struct {
   230  	initiator            bool
   231  	remote               *ecies.PublicKey  // remote-pubk
   232  	initNonce, respNonce []byte            // nonce
   233  	randomPrivKey        *ecies.PrivateKey // ecdhe-random
   234  	remoteRandomPub      *ecies.PublicKey  // ecdhe-random-pubk
   235  }
   236  
   237  // secrets represents the connection secrets
   238  // which are negotiated during the encryption handshake.
   239  type secrets struct {
   240  	Remote                *ecies.PublicKey
   241  	AES, MAC              []byte
   242  	EgressMAC, IngressMAC hash.Hash
   243  	Token                 []byte
   244  }
   245  
   246  // secrets is called after the handshake is completed.
   247  // It extracts the connection secrets from the handshake values.
   248  func (h *encHandshake) secrets(auth, authResp []byte) (secrets, error) {
   249  	ecdheSecret, err := h.randomPrivKey.GenerateShared(h.remoteRandomPub, sskLen, sskLen)
   250  	if err != nil {
   251  		return secrets{}, err
   252  	}
   253  
   254  	// derive base secrets from ephemeral key agreement
   255  	sharedSecret := ogcrypto.Keccak256(ecdheSecret, ogcrypto.Keccak256(h.respNonce, h.initNonce))
   256  	aesSecret := ogcrypto.Keccak256(ecdheSecret, sharedSecret)
   257  	s := secrets{
   258  		Remote: h.remote,
   259  		AES:    aesSecret,
   260  		MAC:    ogcrypto.Keccak256(ecdheSecret, aesSecret),
   261  	}
   262  
   263  	// setup sha3 instances for the MACs
   264  	mac1 := sha3.NewLegacyKeccak256()
   265  	mac1.Write(xor(s.MAC, h.respNonce))
   266  	mac1.Write(auth)
   267  	mac2 := sha3.NewLegacyKeccak256()
   268  	mac2.Write(xor(s.MAC, h.initNonce))
   269  	mac2.Write(authResp)
   270  	if h.initiator {
   271  		s.EgressMAC, s.IngressMAC = mac1, mac2
   272  	} else {
   273  		s.EgressMAC, s.IngressMAC = mac2, mac1
   274  	}
   275  
   276  	return s, nil
   277  }
   278  
   279  // staticSharedSecret returns the static shared secret, the result
   280  // of key agreement between the local and remote static node key.
   281  func (h *encHandshake) staticSharedSecret(prv *ecdsa.PrivateKey) ([]byte, error) {
   282  	return ecies.ImportECDSA(prv).GenerateShared(h.remote, sskLen, sskLen)
   283  }
   284  
   285  // initiatorEncHandshake negotiates a session token on conn.
   286  // it should be called on the dialing side of the connection.
   287  //
   288  // prv is the local client's private key.
   289  func initiatorEncHandshake(conn io.ReadWriter, prv *ecdsa.PrivateKey, remote *ecdsa.PublicKey) (s secrets, err error) {
   290  	h := &encHandshake{initiator: true, remote: ecies.ImportECDSAPublic(remote)}
   291  	authMsg, err := h.makeAuthMsg(prv)
   292  	if err != nil {
   293  		return s, err
   294  	}
   295  	authPacket, err := sealEIP8(authMsg, h)
   296  	if err != nil {
   297  		return s, err
   298  	}
   299  	if _, err = conn.Write(authPacket); err != nil {
   300  		return s, err
   301  	}
   302  	ioperformance.AddSendSize(len(authPacket))
   303  	authRespMsg := new(AuthRespV4)
   304  	authRespPacket, err := readHandshakeMsgResp(authRespMsg, encAuthRespLen, prv, conn)
   305  	if err != nil {
   306  		return s, err
   307  	}
   308  	if err := h.handleAuthResp(authRespMsg); err != nil {
   309  		return s, err
   310  	}
   311  	return h.secrets(authPacket, authRespPacket)
   312  }
   313  
   314  // makeAuthMsg creates the initiator handshake message.
   315  func (h *encHandshake) makeAuthMsg(prv *ecdsa.PrivateKey) (*AuthMsgV4, error) {
   316  	// Generate random initiator nonce.
   317  	h.initNonce = make([]byte, shaLen)
   318  	_, err := rand.Read(h.initNonce)
   319  	if err != nil {
   320  		return nil, err
   321  	}
   322  	// Generate random keypair to for ECDH.
   323  	h.randomPrivKey, err = ecies.GenerateKey(rand.Reader, ogcrypto.S256(), nil)
   324  	if err != nil {
   325  		return nil, err
   326  	}
   327  	// Sign known message: static-shared-secret ^ nonce
   328  	token, err := h.staticSharedSecret(prv)
   329  	if err != nil {
   330  		return nil, err
   331  	}
   332  	signed := xor(token, h.initNonce)
   333  	signature, err := ogcrypto.Sign(signed, h.randomPrivKey.ExportECDSA())
   334  	if err != nil {
   335  		return nil, err
   336  	}
   337  
   338  	msg := new(AuthMsgV4)
   339  	copy(msg.Signature[:], signature)
   340  	copy(msg.InitiatorPubkey[:], ogcrypto.FromECDSAPub(&prv.PublicKey)[1:])
   341  	copy(msg.Nonce[:], h.initNonce)
   342  	msg.Version = 4
   343  	return msg, nil
   344  }
   345  
   346  func (h *encHandshake) handleAuthResp(msg *AuthRespV4) (err error) {
   347  	h.respNonce = msg.Nonce[:]
   348  	h.remoteRandomPub, err = importPublicKey(msg.RandomPubkey[:])
   349  	return err
   350  }
   351  
   352  // receiverEncHandshake negotiates a session token on conn.
   353  // it should be called on the listening side of the connection.
   354  //
   355  // prv is the local client's private key.
   356  func receiverEncHandshake(conn io.ReadWriter, prv *ecdsa.PrivateKey) (s secrets, err error) {
   357  	authMsg := new(AuthMsgV4)
   358  	authPacket, err := readHandshakeMsg(authMsg, encAuthMsgLen, prv, conn)
   359  	if err != nil {
   360  		return s, err
   361  	}
   362  	h := new(encHandshake)
   363  	if err := h.handleAuthMsg(authMsg, prv); err != nil {
   364  		return s, err
   365  	}
   366  
   367  	authRespMsg, err := h.makeAuthResp()
   368  	if err != nil {
   369  		return s, err
   370  	}
   371  	var authRespPacket []byte
   372  	if authMsg.gotPlain {
   373  		authRespPacket, err = authRespMsg.sealPlain(h)
   374  	} else {
   375  		authRespPacket, err = sealEIP8Resp(authRespMsg, h)
   376  	}
   377  	if err != nil {
   378  		return s, err
   379  	}
   380  	if _, err = conn.Write(authRespPacket); err != nil {
   381  		return s, err
   382  	}
   383  	ioperformance.AddSendSize(len(authRespPacket))
   384  	return h.secrets(authPacket, authRespPacket)
   385  }
   386  
   387  func (h *encHandshake) handleAuthMsg(msg *AuthMsgV4, prv *ecdsa.PrivateKey) error {
   388  	// Import the remote identity.
   389  	rpub, err := importPublicKey(msg.InitiatorPubkey[:])
   390  	if err != nil {
   391  		return err
   392  	}
   393  	h.initNonce = msg.Nonce[:]
   394  	h.remote = rpub
   395  
   396  	// Generate random keypair for ECDH.
   397  	// If a private key is already set, use it instead of generating one (for testing).
   398  	if h.randomPrivKey == nil {
   399  		h.randomPrivKey, err = ecies.GenerateKey(rand.Reader, ogcrypto.S256(), nil)
   400  		if err != nil {
   401  			return err
   402  		}
   403  	}
   404  
   405  	// Check the signature.
   406  	token, err := h.staticSharedSecret(prv)
   407  	if err != nil {
   408  		return err
   409  	}
   410  	signedMsg := xor(token, h.initNonce)
   411  	remoteRandomPub, err := ogcrypto.Ecrecover(signedMsg, msg.Signature[:])
   412  	if err != nil {
   413  		return err
   414  	}
   415  	h.remoteRandomPub, _ = importPublicKey(remoteRandomPub)
   416  	return nil
   417  }
   418  
   419  func (h *encHandshake) makeAuthResp() (msg *AuthRespV4, err error) {
   420  	// Generate random nonce.
   421  	h.respNonce = make([]byte, shaLen)
   422  	if _, err = rand.Read(h.respNonce); err != nil {
   423  		return nil, err
   424  	}
   425  
   426  	msg = new(AuthRespV4)
   427  	copy(msg.Nonce[:], h.respNonce)
   428  	copy(msg.RandomPubkey[:], exportPubkey(&h.randomPrivKey.PublicKey))
   429  	msg.Version = 4
   430  	return msg, nil
   431  }
   432  
   433  func (msg *AuthMsgV4) sealPlain(h *encHandshake) ([]byte, error) {
   434  	buf := make([]byte, authMsgLen)
   435  	n := copy(buf, msg.Signature[:])
   436  	n += copy(buf[n:], ogcrypto.Keccak256(exportPubkey(&h.randomPrivKey.PublicKey)))
   437  	n += copy(buf[n:], msg.InitiatorPubkey[:])
   438  	n += copy(buf[n:], msg.Nonce[:])
   439  	buf[n] = 0 // token-flag
   440  	return ecies.Encrypt(rand.Reader, h.remote, buf, nil, nil)
   441  }
   442  
   443  func (msg *AuthMsgV4) decodePlain(input []byte) {
   444  	n := copy(msg.Signature[:], input)
   445  	n += shaLen // skip sha3(initiator-ephemeral-pubk)
   446  	n += copy(msg.InitiatorPubkey[:], input[n:])
   447  	copy(msg.Nonce[:], input[n:])
   448  	msg.Version = 4
   449  	msg.gotPlain = true
   450  }
   451  
   452  func (msg *AuthRespV4) sealPlain(hs *encHandshake) ([]byte, error) {
   453  	buf := make([]byte, authRespLen)
   454  	n := copy(buf, msg.RandomPubkey[:])
   455  	copy(buf[n:], msg.Nonce[:])
   456  	return ecies.Encrypt(rand.Reader, hs.remote, buf, nil, nil)
   457  }
   458  
   459  func (msg *AuthRespV4) decodePlain(input []byte) {
   460  	n := copy(msg.RandomPubkey[:], input)
   461  	copy(msg.Nonce[:], input[n:])
   462  	msg.Version = 4
   463  }
   464  
   465  var padSpace = make([]byte, 300)
   466  
   467  func sealEIP8(msg *AuthMsgV4, h *encHandshake) ([]byte, error) {
   468  	buf := new(bytes.Buffer)
   469  	b, err := msg.MarshalMsg(nil)
   470  	if err != nil {
   471  		return nil, err
   472  	}
   473  	buf.Write(b)
   474  
   475  	/*
   476  		if err := rlp.Encode(buf, msg); err != nil {
   477  			return nil, err
   478  		}
   479  	*/
   480  	// pad with random amount of data. the amount needs to be at least 100 bytes to make
   481  	// the message distinguishable from pre-EIP-8 handshakes.
   482  	pad := padSpace[:mrand.Intn(len(padSpace)-100)+100]
   483  	buf.Write(pad)
   484  	prefix := make([]byte, 2)
   485  	binary.BigEndian.PutUint16(prefix, uint16(buf.Len()+eciesOverhead))
   486  
   487  	enc, err := ecies.Encrypt(rand.Reader, h.remote, buf.Bytes(), nil, prefix)
   488  	return append(prefix, enc...), err
   489  }
   490  
   491  func sealEIP8Resp(msg *AuthRespV4, h *encHandshake) ([]byte, error) {
   492  	buf := new(bytes.Buffer)
   493  	b, err := msg.MarshalMsg(nil)
   494  	if err != nil {
   495  		return nil, err
   496  	}
   497  	buf.Write(b)
   498  	/*
   499  		if err := rlp.Encode(buf, msg); err != nil {
   500  			return nil, err
   501  		}
   502  	*/
   503  
   504  	// pad with random amount of data. the amount needs to be at least 100 bytes to make
   505  	// the message distinguishable from pre-EIP-8 handshakes.
   506  	pad := padSpace[:mrand.Intn(len(padSpace)-100)+100]
   507  	buf.Write(pad)
   508  	prefix := make([]byte, 2)
   509  	binary.BigEndian.PutUint16(prefix, uint16(buf.Len()+eciesOverhead))
   510  
   511  	enc, err := ecies.Encrypt(rand.Reader, h.remote, buf.Bytes(), nil, prefix)
   512  	return append(prefix, enc...), err
   513  }
   514  
   515  //test function ,just for test
   516  func preEip8test(msg *AuthMsgV4, remotePubkey *ecies.PublicKey) ([]byte, error) {
   517  	buf := new(bytes.Buffer)
   518  	b, err := msg.MarshalMsg(nil)
   519  	if err != nil {
   520  		return nil, err
   521  	}
   522  	buf.Write(b)
   523  
   524  	pad := padSpace[:mrand.Intn(len(padSpace)-100)+100]
   525  	buf.Write(pad)
   526  	prefix := make([]byte, 2)
   527  	binary.BigEndian.PutUint16(prefix, uint16(buf.Len()+eciesOverhead))
   528  
   529  	enc, err := ecies.Encrypt(rand.Reader, remotePubkey, buf.Bytes(), nil, prefix)
   530  	return append(prefix, enc...), err
   531  
   532  }
   533  
   534  ////test function ,just for test
   535  func preEip8Resptest(msg *AuthRespV4, remotePubkey *ecies.PublicKey) ([]byte, error) {
   536  	buf := new(bytes.Buffer)
   537  	b, err := msg.MarshalMsg(nil)
   538  	if err != nil {
   539  		return nil, err
   540  	}
   541  	buf.Write(b)
   542  
   543  	pad := padSpace[:mrand.Intn(len(padSpace)-100)+100]
   544  	buf.Write(pad)
   545  	prefix := make([]byte, 2)
   546  	binary.BigEndian.PutUint16(prefix, uint16(buf.Len()+eciesOverhead))
   547  
   548  	enc, err := ecies.Encrypt(rand.Reader, remotePubkey, buf.Bytes(), nil, prefix)
   549  	return append(prefix, enc...), err
   550  
   551  }
   552  
   553  type plainDecoder interface {
   554  	decodePlain([]byte)
   555  }
   556  
   557  func readHandshakeMsgResp(msg *AuthRespV4, plainSize int, prv *ecdsa.PrivateKey, r io.Reader) ([]byte, error) {
   558  	buf := make([]byte, plainSize)
   559  	if _, err := io.ReadFull(r, buf); err != nil {
   560  		return buf, err
   561  	}
   562  	ioperformance.AddRecvSize(len(buf))
   563  	// Attempt decoding pre-EIP-8 "plain" format.
   564  	key := ecies.ImportECDSA(prv)
   565  	if dec, err := key.Decrypt(buf, nil, nil); err == nil {
   566  		msg.decodePlain(dec)
   567  		//	fmt.Println("is plain",msg)
   568  		return buf, nil
   569  	}
   570  	// Could be EIP-8 format, try that.
   571  	prefix := buf[:2]
   572  	size := binary.BigEndian.Uint16(prefix)
   573  	if size < uint16(plainSize) {
   574  		return buf, fmt.Errorf("size underflow, need at least %d bytes", plainSize)
   575  	}
   576  	buf = append(buf, make([]byte, size-uint16(plainSize)+2)...)
   577  	if _, err := io.ReadFull(r, buf[plainSize:]); err != nil {
   578  		return buf, err
   579  	}
   580  	dec, err := key.Decrypt(buf[2:], nil, prefix)
   581  	if err != nil {
   582  		return buf, err
   583  	}
   584  	// Can't use rlp.DecodeBytes here because it rejects
   585  	// trailing data (forward-compatibility).
   586  	//s := rlp.NewStream(bytes.NewReader(dec), 0)
   587  	_, err = msg.UnmarshalMsg(dec)
   588  	//return buf, s.Decode(msg)
   589  	return buf, err
   590  }
   591  
   592  func readHandshakeMsg(msg *AuthMsgV4, plainSize int, prv *ecdsa.PrivateKey, r io.Reader) ([]byte, error) {
   593  	buf := make([]byte, plainSize)
   594  	if _, err := io.ReadFull(r, buf); err != nil {
   595  		return buf, err
   596  	}
   597  	// Attempt decoding pre-EIP-8 "plain" format.
   598  	key := ecies.ImportECDSA(prv)
   599  	if dec, err := key.Decrypt(buf, nil, nil); err == nil {
   600  		msg.decodePlain(dec)
   601  		//fmt.Println("is plain",msg)
   602  		return buf, nil
   603  	}
   604  
   605  	// Could be EIP-8 format, try that.
   606  	prefix := buf[:2]
   607  	size := binary.BigEndian.Uint16(prefix)
   608  	if size < uint16(plainSize) {
   609  		return buf, fmt.Errorf("size underflow, need at least %d bytes", plainSize)
   610  	}
   611  	buf = append(buf, make([]byte, size-uint16(plainSize)+2)...)
   612  	if _, err := io.ReadFull(r, buf[plainSize:]); err != nil {
   613  		return buf, err
   614  	}
   615  	dec, err := key.Decrypt(buf[2:], nil, prefix)
   616  	if err != nil {
   617  		return buf, err
   618  	}
   619  	// Can't use rlp.DecodeBytes here because it rejects
   620  	// trailing data (forward-compatibility).
   621  	//s := rlp.NewStream(bytes.NewReader(dec), 0)
   622  	_, err = msg.UnmarshalMsg(dec)
   623  	//return buf, s.Decode(msg)
   624  	return buf, err
   625  }
   626  
   627  // importPublicKey unmarshals 512 bit public keys.
   628  func importPublicKey(pubKey []byte) (*ecies.PublicKey, error) {
   629  	var pubKey65 []byte
   630  	switch len(pubKey) {
   631  	case 64:
   632  		// add 'uncompressed key' flag
   633  		pubKey65 = append([]byte{0x04}, pubKey...)
   634  	case 65:
   635  		pubKey65 = pubKey
   636  	default:
   637  		return nil, fmt.Errorf("invalid public key length %v (expect 64/65)", len(pubKey))
   638  	}
   639  	// TODO: fewer pointless conversions
   640  	pub, err := ogcrypto.UnmarshalPubkey(pubKey65)
   641  	if err != nil {
   642  		return nil, err
   643  	}
   644  	return ecies.ImportECDSAPublic(pub), nil
   645  }
   646  
   647  func exportPubkey(pub *ecies.PublicKey) []byte {
   648  	if pub == nil {
   649  		panic("nil pubkey")
   650  	}
   651  	return elliptic.Marshal(pub.Curve, pub.X, pub.Y)[1:]
   652  }
   653  
   654  func xor(one, other []byte) (xor []byte) {
   655  	xor = make([]byte, len(one))
   656  	for i := 0; i < len(one); i++ {
   657  		xor[i] = one[i] ^ other[i]
   658  	}
   659  	return xor
   660  }
   661  
   662  var (
   663  	// this is used in place of actual frame header data.
   664  	// TODO: replace this when Msg contains the protocol type code.
   665  	zeroHeader = []byte{0xC2, 0x80, 0x80}
   666  	// sixteen zero bytes
   667  	zero16 = make([]byte, 16)
   668  )
   669  
   670  // rlpxFrameRW implements a simplified version of RLPx framing.
   671  // chunked messages are not supported and all headers are equal to
   672  // zeroHeader.
   673  //
   674  // rlpxFrameRW is not safe for concurrent use from multiple goroutines.
   675  type rlpxFrameRW struct {
   676  	conn io.ReadWriter
   677  	enc  cipher.Stream
   678  	dec  cipher.Stream
   679  
   680  	macCipher  cipher.Block
   681  	egressMAC  hash.Hash
   682  	ingressMAC hash.Hash
   683  
   684  	snappy bool
   685  }
   686  
   687  func newRLPXFrameRW(conn io.ReadWriter, s secrets) *rlpxFrameRW {
   688  	macc, err := aes.NewCipher(s.MAC)
   689  	if err != nil {
   690  		panic("invalid MAC secret: " + err.Error())
   691  	}
   692  	encc, err := aes.NewCipher(s.AES)
   693  	if err != nil {
   694  		panic("invalid AES secret: " + err.Error())
   695  	}
   696  	// we use an all-zeroes IV for AES because the key used
   697  	// for encryption is ephemeral.
   698  	iv := make([]byte, encc.BlockSize())
   699  	return &rlpxFrameRW{
   700  		conn:       conn,
   701  		enc:        cipher.NewCTR(encc, iv),
   702  		dec:        cipher.NewCTR(encc, iv),
   703  		macCipher:  macc,
   704  		egressMAC:  s.EgressMAC,
   705  		ingressMAC: s.IngressMAC,
   706  	}
   707  }
   708  
   709  func (rw *rlpxFrameRW) WriteMsg(msg Msg) error {
   710  	//ptype, _ := rlp.EncodeToBytes(msg.Code)
   711  	ptype, _ := msg.Code.MarshalMsg(nil)
   712  	// if snappy is enabled, compress message now
   713  	if rw.snappy {
   714  		if msg.Size > maxUint24 {
   715  			return errPlainMessageTooLarge
   716  		}
   717  		payload, _ := msg.GetPayLoad()
   718  		payload = snappy.Encode(nil, payload)
   719  
   720  		msg.Payload = bytes.NewReader(payload)
   721  		msg.Size = uint32(len(payload))
   722  	}
   723  	// write header
   724  	headbuf := make([]byte, 32)
   725  	fsize := uint32(len(ptype)) + msg.Size
   726  	if fsize > maxUint24 {
   727  		return errors.New("message size overflows uint24")
   728  	}
   729  	putInt24(fsize, headbuf) // TODO: check overflow
   730  	copy(headbuf[3:], zeroHeader)
   731  	rw.enc.XORKeyStream(headbuf[:16], headbuf[:16]) // first half is now encrypted
   732  
   733  	// write header MAC
   734  	copy(headbuf[16:], updateMAC(rw.egressMAC, rw.macCipher, headbuf[:16]))
   735  	if _, err := rw.conn.Write(headbuf); err != nil {
   736  		return err
   737  	}
   738  
   739  	// write encrypted frame, updating the egress MAC hash with
   740  	// the data written to conn.
   741  	tee := cipher.StreamWriter{S: rw.enc, W: io.MultiWriter(rw.conn, rw.egressMAC)}
   742  	if _, err := tee.Write(ptype); err != nil {
   743  		return err
   744  	}
   745  
   746  	if _, err := io.Copy(tee, msg.Payload); err != nil {
   747  		return err
   748  	}
   749  
   750  	if padding := fsize % 16; padding > 0 {
   751  		if _, err := tee.Write(zero16[:16-padding]); err != nil {
   752  			return err
   753  		}
   754  	}
   755  
   756  	// write frame MAC. egress MAC hash is up to date because
   757  	// frame content was written to it as well.
   758  	fmacseed := rw.egressMAC.Sum(nil)
   759  	mac := updateMAC(rw.egressMAC, rw.macCipher, fmacseed)
   760  	_, err := rw.conn.Write(mac)
   761  
   762  	totalSize := len(headbuf) + len(ptype) + len(mac) + int(msg.Size) + len(zero16)
   763  	log.Trace("write len , ", totalSize)
   764  	ioperformance.AddSendSize(totalSize)
   765  	return err
   766  
   767  }
   768  
   769  func (rw *rlpxFrameRW) ReadMsg() (msg Msg, err error) {
   770  	// read the header
   771  	headbuf := make([]byte, 32)
   772  	if _, err := io.ReadFull(rw.conn, headbuf); err != nil {
   773  		return msg, err
   774  	}
   775  	// verify header mac
   776  	shouldMAC := updateMAC(rw.ingressMAC, rw.macCipher, headbuf[:16])
   777  	if !hmac.Equal(shouldMAC, headbuf[16:]) {
   778  		return msg, errors.New("bad header MAC")
   779  	}
   780  	rw.dec.XORKeyStream(headbuf[:16], headbuf[:16]) // first half is now decrypted
   781  	fsize := readInt24(headbuf)
   782  	// ignore protocol type for now
   783  
   784  	// read the frame content
   785  	var rsize = fsize // frame size rounded up to 16 byte boundary
   786  	if padding := fsize % 16; padding > 0 {
   787  		rsize += 16 - padding
   788  	}
   789  	framebuf := make([]byte, rsize)
   790  	if _, err := io.ReadFull(rw.conn, framebuf); err != nil {
   791  		return msg, err
   792  	}
   793  
   794  	// read and validate frame MAC. we can re-use headbuf for that.
   795  	rw.ingressMAC.Write(framebuf)
   796  	fmacseed := rw.ingressMAC.Sum(nil)
   797  	if _, err := io.ReadFull(rw.conn, headbuf[:16]); err != nil {
   798  		return msg, err
   799  	}
   800  	shouldMAC = updateMAC(rw.ingressMAC, rw.macCipher, fmacseed)
   801  	if !hmac.Equal(shouldMAC, headbuf[:16]) {
   802  		return msg, errors.New("bad frame MAC")
   803  	}
   804  
   805  	// decrypt frame content
   806  	rw.dec.XORKeyStream(framebuf, framebuf)
   807  
   808  	// decode message code
   809  	//content := bytes.NewReader(framebuf[:fsize])
   810  	/*
   811  		if err := rlp.Decode(content, &msg.Code); err != nil {
   812  			return msg, err
   813  		}
   814  	*/
   815  	out, err := msg.Code.UnmarshalMsg(framebuf[:fsize])
   816  	if err != nil {
   817  		return msg, err
   818  	}
   819  	content := bytes.NewReader(out)
   820  	msg.Size = uint32(content.Len())
   821  	msg.Payload = content
   822  	// if snappy is enabled, verify and decompress message
   823  	if rw.snappy {
   824  		payload, err := ioutil.ReadAll(msg.Payload)
   825  		if err != nil {
   826  			return msg, err
   827  		}
   828  		size, err := snappy.DecodedLen(payload)
   829  		if err != nil {
   830  			return msg, err
   831  		}
   832  		if size > int(maxUint24) {
   833  			return msg, errPlainMessageTooLarge
   834  		}
   835  		payload, err = snappy.Decode(nil, payload)
   836  		if err != nil {
   837  			return msg, err
   838  		}
   839  		msg.Size, msg.Payload = uint32(size), bytes.NewReader(payload)
   840  	}
   841  	totalSize := len(headbuf) + 3 + len(shouldMAC) + int(msg.Size) + len(zero16)
   842  	log.Trace("write len , ", totalSize)
   843  	ioperformance.AddRecvSize(totalSize)
   844  	return msg, nil
   845  }
   846  
   847  // updateMAC reseeds the given hash with encrypted seed.
   848  // it returns the first 16 bytes of the hash sum after seeding.
   849  func updateMAC(mac hash.Hash, block cipher.Block, seed []byte) []byte {
   850  	aesbuf := make([]byte, aes.BlockSize)
   851  	block.Encrypt(aesbuf, mac.Sum(nil))
   852  	for i := range aesbuf {
   853  		aesbuf[i] ^= seed[i]
   854  	}
   855  	mac.Write(aesbuf)
   856  	return mac.Sum(nil)[:16]
   857  }
   858  
   859  func readInt24(b []byte) uint32 {
   860  	return uint32(b[2]) | uint32(b[1])<<8 | uint32(b[0])<<16
   861  }
   862  
   863  func putInt24(v uint32, b []byte) {
   864  	b[0] = byte(v >> 16)
   865  	b[1] = byte(v >> 8)
   866  	b[2] = byte(v)
   867  }