github.com/arieschain/arieschain@v0.0.0-20191023063405-37c074544356/p2p/rlpx.go (about)

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