9afc6816d2
It seems there is no fully typed library implementation of an LRU cache. So I wrote one. Method names are the same as github.com/hashicorp/golang-lru, and the new type can be used as a drop-in replacement. Two reasons to do this: - It's much easier to understand what a cache is for when the types are right there. - Performance: the new implementation is slightly faster and performs zero memory allocations in Add when the cache is at capacity. Overall, memory usage of the cache is much reduced because keys are values are no longer wrapped in interface.
136 lines
4.1 KiB
Go
136 lines
4.1 KiB
Go
// Copyright 2020 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 v5wire
|
|
|
|
import (
|
|
"crypto/ecdsa"
|
|
crand "crypto/rand"
|
|
"encoding/binary"
|
|
"time"
|
|
|
|
"github.com/ethereum/go-ethereum/common/lru"
|
|
"github.com/ethereum/go-ethereum/common/mclock"
|
|
"github.com/ethereum/go-ethereum/crypto"
|
|
"github.com/ethereum/go-ethereum/p2p/enode"
|
|
)
|
|
|
|
const handshakeTimeout = time.Second
|
|
|
|
// The SessionCache keeps negotiated encryption keys and
|
|
// state for in-progress handshakes in the Discovery v5 wire protocol.
|
|
type SessionCache struct {
|
|
sessions lru.BasicLRU[sessionID, *session]
|
|
handshakes map[sessionID]*Whoareyou
|
|
clock mclock.Clock
|
|
|
|
// hooks for overriding randomness.
|
|
nonceGen func(uint32) (Nonce, error)
|
|
maskingIVGen func([]byte) error
|
|
ephemeralKeyGen func() (*ecdsa.PrivateKey, error)
|
|
}
|
|
|
|
// sessionID identifies a session or handshake.
|
|
type sessionID struct {
|
|
id enode.ID
|
|
addr string
|
|
}
|
|
|
|
// session contains session information
|
|
type session struct {
|
|
writeKey []byte
|
|
readKey []byte
|
|
nonceCounter uint32
|
|
}
|
|
|
|
// keysFlipped returns a copy of s with the read and write keys flipped.
|
|
func (s *session) keysFlipped() *session {
|
|
return &session{s.readKey, s.writeKey, s.nonceCounter}
|
|
}
|
|
|
|
func NewSessionCache(maxItems int, clock mclock.Clock) *SessionCache {
|
|
return &SessionCache{
|
|
sessions: lru.NewBasicLRU[sessionID, *session](maxItems),
|
|
handshakes: make(map[sessionID]*Whoareyou),
|
|
clock: clock,
|
|
nonceGen: generateNonce,
|
|
maskingIVGen: generateMaskingIV,
|
|
ephemeralKeyGen: crypto.GenerateKey,
|
|
}
|
|
}
|
|
|
|
func generateNonce(counter uint32) (n Nonce, err error) {
|
|
binary.BigEndian.PutUint32(n[:4], counter)
|
|
_, err = crand.Read(n[4:])
|
|
return n, err
|
|
}
|
|
|
|
func generateMaskingIV(buf []byte) error {
|
|
_, err := crand.Read(buf)
|
|
return err
|
|
}
|
|
|
|
// nextNonce creates a nonce for encrypting a message to the given session.
|
|
func (sc *SessionCache) nextNonce(s *session) (Nonce, error) {
|
|
s.nonceCounter++
|
|
return sc.nonceGen(s.nonceCounter)
|
|
}
|
|
|
|
// session returns the current session for the given node, if any.
|
|
func (sc *SessionCache) session(id enode.ID, addr string) *session {
|
|
item, _ := sc.sessions.Get(sessionID{id, addr})
|
|
return item
|
|
}
|
|
|
|
// readKey returns the current read key for the given node.
|
|
func (sc *SessionCache) readKey(id enode.ID, addr string) []byte {
|
|
if s := sc.session(id, addr); s != nil {
|
|
return s.readKey
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// storeNewSession stores new encryption keys in the cache.
|
|
func (sc *SessionCache) storeNewSession(id enode.ID, addr string, s *session) {
|
|
sc.sessions.Add(sessionID{id, addr}, s)
|
|
}
|
|
|
|
// getHandshake gets the handshake challenge we previously sent to the given remote node.
|
|
func (sc *SessionCache) getHandshake(id enode.ID, addr string) *Whoareyou {
|
|
return sc.handshakes[sessionID{id, addr}]
|
|
}
|
|
|
|
// storeSentHandshake stores the handshake challenge sent to the given remote node.
|
|
func (sc *SessionCache) storeSentHandshake(id enode.ID, addr string, challenge *Whoareyou) {
|
|
challenge.sent = sc.clock.Now()
|
|
sc.handshakes[sessionID{id, addr}] = challenge
|
|
}
|
|
|
|
// deleteHandshake deletes handshake data for the given node.
|
|
func (sc *SessionCache) deleteHandshake(id enode.ID, addr string) {
|
|
delete(sc.handshakes, sessionID{id, addr})
|
|
}
|
|
|
|
// handshakeGC deletes timed-out handshakes.
|
|
func (sc *SessionCache) handshakeGC() {
|
|
deadline := sc.clock.Now().Add(-handshakeTimeout)
|
|
for key, challenge := range sc.handshakes {
|
|
if challenge.sent < deadline {
|
|
delete(sc.handshakes, key)
|
|
}
|
|
}
|
|
}
|