diff --git a/pkg/ltsm/embind.go b/pkg/ltsm/embind.go index f8fb76a6..2ad4b2ed 100644 --- a/pkg/ltsm/embind.go +++ b/pkg/ltsm/embind.go @@ -960,7 +960,7 @@ func (imp *Imports) noteChannelState(className string) { } switch className { case "Curve25519Key", "E2EEKey", "E2EEChannel", "E2EEKeychain": - imp.owner.cryptoDirty = true + imp.owner.markChanged() } } diff --git a/pkg/ltsm/runtime.go b/pkg/ltsm/runtime.go index 114d602b..55cbaf33 100644 --- a/pkg/ltsm/runtime.go +++ b/pkg/ltsm/runtime.go @@ -6,17 +6,18 @@ import ( "crypto/sha256" "encoding/base64" "fmt" + "sync" + "sync/atomic" ) // Runtime wraps the transpiled Module and embind Imports to provide // the same high-level API as wasm.Runtime. type Runtime struct { - mod *Module - imp *Imports - crypto *Runtime - cryptoCalls int - cryptoDirty bool - isCrypto bool + mod *Module + imp *Imports + generation uint64 + calls int + isCrypto bool } // NewRuntime creates and initializes a new transpiled LTSM runtime. @@ -31,6 +32,7 @@ func NewRuntime() (*Runtime, error) { mod.fT() rt := &Runtime{mod: mod, imp: imp} + rt.markChanged() imp.owner = rt return rt, nil } @@ -47,49 +49,52 @@ func (rt *Runtime) Close() {} // disposable copy so their allocator cannot exhaust the authoritative module. const cryptoRefreshCalls = 128 -func (rt *Runtime) refreshCrypto() error { - if rt.crypto == nil { - crypto, err := NewRuntime() - if err != nil { - return err +var ( + scratchRuntimes sync.Pool + runtimeGenerations atomic.Uint64 +) + +func (rt *Runtime) markChanged() { + rt.generation = runtimeGenerations.Add(1) +} + +func (rt *Runtime) scratch() (*Runtime, error) { + crypto, _ := scratchRuntimes.Get().(*Runtime) + if crypto == nil { + var err error + if crypto, err = NewRuntime(); err != nil { + return nil, err } crypto.isCrypto = true - rt.crypto = crypto } - copy(rt.crypto.mod.mem, rt.mod.mem) - rt.crypto.mod.g0 = rt.mod.g0 - values, refCnts, free := rt.imp.emval.snapshot() - rt.crypto.imp.emval.restore(values, refCnts, free) - rt.crypto.imp.urandomFD, rt.crypto.imp.nextFD = rt.imp.urandomFD, rt.imp.nextFD - rt.cryptoCalls = 0 - rt.cryptoDirty = false - return nil + if crypto.generation != rt.generation || crypto.calls >= cryptoRefreshCalls { + copy(crypto.mod.mem, rt.mod.mem) + crypto.mod.g0 = rt.mod.g0 + values, refCnts, free := rt.imp.emval.snapshot() + crypto.imp.emval.restore(values, refCnts, free) + crypto.imp.urandomFD, crypto.imp.nextFD = rt.imp.urandomFD, rt.imp.nextFD + crypto.generation, crypto.calls = rt.generation, 0 + } + return crypto, nil } -func (rt *Runtime) channelCrypto(call func(*Runtime) ([]byte, error)) (data []byte, err error) { +func (rt *Runtime) channelCrypto(call func(*Runtime) ([]byte, error)) ([]byte, error) { if rt.isCrypto { return call(rt) } - if rt.crypto == nil || rt.cryptoDirty || rt.cryptoCalls >= cryptoRefreshCalls { - if err := rt.refreshCrypto(); err != nil { - return nil, err - } + crypto, err := rt.scratch() + if err != nil { + return nil, err } - defer func() { - if recovered := recover(); recovered != nil { - rt.cryptoDirty = true - panic(recovered) - } - if err != nil { - rt.cryptoDirty = true - } - }() - data, err = call(rt.crypto) + data, err := call(crypto) if err != nil { + crypto.markChanged() + scratchRuntimes.Put(crypto) return nil, err } data = append([]byte(nil), data...) - rt.cryptoCalls++ + crypto.calls++ + scratchRuntimes.Put(crypto) return data, nil } @@ -381,7 +386,7 @@ func (rt *Runtime) E2EEChannelUnwrapGroupSharedKey(chanPtr uint32, encKey []byte func (rt *Runtime) WrapRegistrationGroupKey(chanPtrs []uint32) ([][]byte, error) { var wrapped [][]byte - defer func() { rt.cryptoDirty = true }() + defer rt.markChanged() _, err := rt.channelCrypto(func(crypto *Runtime) ([]byte, error) { keyPtr, err := crypto.Curve25519KeyGenerate() if err != nil {