ceremonyclient/go-libp2p-blossomsub/midgen.go
2023-07-07 01:07:10 -05:00

53 lines
1.2 KiB
Go

package blossomsub
import (
"sync"
pb "source.quilibrium.com/quilibrium/monorepo/go-libp2p-blossomsub/pb"
)
// msgIDGenerator handles computing IDs for msgs
// It allows setting custom generators(MsgIdFunction) per bitmask
type msgIDGenerator struct {
Default MsgIdFunction
bitmaskGensLk sync.RWMutex
bitmaskGens map[string]MsgIdFunction
}
func newMsgIdGenerator() *msgIDGenerator {
return &msgIDGenerator{
Default: DefaultMsgIdFn,
bitmaskGens: make(map[string]MsgIdFunction),
}
}
// Set sets custom id generator(MsgIdFunction) for bitmask.
func (m *msgIDGenerator) Set(bitmask []byte, gen MsgIdFunction) {
m.bitmaskGensLk.Lock()
m.bitmaskGens[string(bitmask)] = gen
m.bitmaskGensLk.Unlock()
}
// ID computes ID for the msg or short-circuits with the cached value.
func (m *msgIDGenerator) ID(msg *Message) string {
if msg.ID != "" {
return msg.ID
}
msg.ID = m.RawID(msg.Message)
return msg.ID
}
// RawID computes ID for the proto 'msg'.
func (m *msgIDGenerator) RawID(msg *pb.Message) string {
m.bitmaskGensLk.RLock()
gen, ok := m.bitmaskGens[string(msg.GetBitmask())]
m.bitmaskGensLk.RUnlock()
if !ok {
gen = m.Default
}
return gen(msg)
}