github.com/ethereum/go-ethereum@v1.16.1/p2p/discover/v5wire/session.go (about)

     1  // Copyright 2020 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  package v5wire
    18  
    19  import (
    20  	"crypto/ecdsa"
    21  	crand "crypto/rand"
    22  	"encoding/binary"
    23  	"time"
    24  
    25  	"github.com/ethereum/go-ethereum/common/lru"
    26  	"github.com/ethereum/go-ethereum/common/mclock"
    27  	"github.com/ethereum/go-ethereum/crypto"
    28  	"github.com/ethereum/go-ethereum/p2p/enode"
    29  )
    30  
    31  const handshakeTimeout = time.Second
    32  
    33  // The SessionCache keeps negotiated encryption keys and
    34  // state for in-progress handshakes in the Discovery v5 wire protocol.
    35  type SessionCache struct {
    36  	sessions   lru.BasicLRU[sessionID, *session]
    37  	handshakes map[sessionID]*Whoareyou
    38  	clock      mclock.Clock
    39  
    40  	// hooks for overriding randomness.
    41  	nonceGen        func(uint32) (Nonce, error)
    42  	maskingIVGen    func([]byte) error
    43  	ephemeralKeyGen func() (*ecdsa.PrivateKey, error)
    44  }
    45  
    46  // sessionID identifies a session or handshake.
    47  type sessionID struct {
    48  	id   enode.ID
    49  	addr string
    50  }
    51  
    52  // session contains session information
    53  type session struct {
    54  	writeKey     []byte
    55  	readKey      []byte
    56  	nonceCounter uint32
    57  	node         *enode.Node
    58  }
    59  
    60  // keysFlipped returns a copy of s with the read and write keys flipped.
    61  func (s *session) keysFlipped() *session {
    62  	return &session{s.readKey, s.writeKey, s.nonceCounter, s.node}
    63  }
    64  
    65  func NewSessionCache(maxItems int, clock mclock.Clock) *SessionCache {
    66  	return &SessionCache{
    67  		sessions:        lru.NewBasicLRU[sessionID, *session](maxItems),
    68  		handshakes:      make(map[sessionID]*Whoareyou),
    69  		clock:           clock,
    70  		nonceGen:        generateNonce,
    71  		maskingIVGen:    generateMaskingIV,
    72  		ephemeralKeyGen: crypto.GenerateKey,
    73  	}
    74  }
    75  
    76  func generateNonce(counter uint32) (n Nonce, err error) {
    77  	binary.BigEndian.PutUint32(n[:4], counter)
    78  	_, err = crand.Read(n[4:])
    79  	return n, err
    80  }
    81  
    82  func generateMaskingIV(buf []byte) error {
    83  	_, err := crand.Read(buf)
    84  	return err
    85  }
    86  
    87  // nextNonce creates a nonce for encrypting a message to the given session.
    88  func (sc *SessionCache) nextNonce(s *session) (Nonce, error) {
    89  	s.nonceCounter++
    90  	return sc.nonceGen(s.nonceCounter)
    91  }
    92  
    93  // session returns the current session for the given node, if any.
    94  func (sc *SessionCache) session(id enode.ID, addr string) *session {
    95  	item, _ := sc.sessions.Get(sessionID{id, addr})
    96  	return item
    97  }
    98  
    99  // readKey returns the current read key for the given node.
   100  func (sc *SessionCache) readKey(id enode.ID, addr string) []byte {
   101  	if s := sc.session(id, addr); s != nil {
   102  		return s.readKey
   103  	}
   104  	return nil
   105  }
   106  
   107  func (sc *SessionCache) readNode(id enode.ID, addr string) *enode.Node {
   108  	if s := sc.session(id, addr); s != nil {
   109  		return s.node
   110  	}
   111  	return nil
   112  }
   113  
   114  // storeNewSession stores new encryption keys in the cache.
   115  func (sc *SessionCache) storeNewSession(id enode.ID, addr string, s *session, n *enode.Node) {
   116  	if n == nil {
   117  		panic("nil node in storeNewSession")
   118  	}
   119  	s.node = n
   120  	sc.sessions.Add(sessionID{id, addr}, s)
   121  }
   122  
   123  // getHandshake gets the handshake challenge we previously sent to the given remote node.
   124  func (sc *SessionCache) getHandshake(id enode.ID, addr string) *Whoareyou {
   125  	return sc.handshakes[sessionID{id, addr}]
   126  }
   127  
   128  // storeSentHandshake stores the handshake challenge sent to the given remote node.
   129  func (sc *SessionCache) storeSentHandshake(id enode.ID, addr string, challenge *Whoareyou) {
   130  	challenge.sent = sc.clock.Now()
   131  	sc.handshakes[sessionID{id, addr}] = challenge
   132  }
   133  
   134  // deleteHandshake deletes handshake data for the given node.
   135  func (sc *SessionCache) deleteHandshake(id enode.ID, addr string) {
   136  	delete(sc.handshakes, sessionID{id, addr})
   137  }
   138  
   139  // handshakeGC deletes timed-out handshakes.
   140  func (sc *SessionCache) handshakeGC() {
   141  	deadline := sc.clock.Now().Add(-handshakeTimeout)
   142  	for key, challenge := range sc.handshakes {
   143  		if challenge.sent < deadline {
   144  			delete(sc.handshakes, key)
   145  		}
   146  	}
   147  }