diff --git a/handshake/ephermal_cache.go b/handshake/ephermal_cache.go index 8dcd7fcc6..794bcbdd7 100644 --- a/handshake/ephermal_cache.go +++ b/handshake/ephermal_cache.go @@ -1,7 +1,6 @@ package handshake import ( - "fmt" "sync" "time" @@ -11,10 +10,10 @@ import ( ) var ( - kexLifetime = protocol.EphermalKeyLifetime - kexCurrent crypto.KeyExchange - kexSetOnce sync.Once - kexMutex sync.RWMutex + kexLifetime = protocol.EphermalKeyLifetime + kexCurrent crypto.KeyExchange + kexCurrentTime time.Time + kexMutex sync.RWMutex ) // getEphermalKEX returns the currently active KEX, which changes every protocol.EphermalKeyLifetime @@ -25,31 +24,27 @@ var ( // used for all connections for 60 seconds is negligible. Thus we can amortise // the Diffie-Hellman key generation at the server over all the connections in a // small time span. -func getEphermalKEX() crypto.KeyExchange { - kexSetOnce.Do(func() { go setKexRoutine() }) +func getEphermalKEX() (res crypto.KeyExchange) { kexMutex.RLock() - defer kexMutex.RUnlock() - return kexCurrent -} + res = kexCurrent + t := kexCurrentTime + kexMutex.RUnlock() + if res != nil && time.Now().Sub(t) < kexLifetime { + return res + } -func setKexRoutine() { - for { - time.Sleep(kexLifetime) + kexMutex.Lock() + defer kexMutex.Unlock() + // Check if still unfulfilled + if kexCurrent == nil || time.Now().Sub(kexCurrentTime) > kexLifetime { kex, err := crypto.NewCurve25519KEX() if err != nil { utils.Errorf("could not set KEX: %s", err.Error()) - continue + return kexCurrent } - kexMutex.Lock() kexCurrent = kex - kexMutex.Unlock() + kexCurrentTime = time.Now() + return kexCurrent } -} - -func init() { - kex, err := crypto.NewCurve25519KEX() - if err != nil { - panic(fmt.Sprintf("Could not set KEX: %s", err.Error())) - } - kexCurrent = kex + return kexCurrent }