waku-org/go-waku

package enr

import (
    "errors"

    "github.com/ethereum/go-ethereum/p2p/enode"
    "github.com/ethereum/go-ethereum/p2p/enr"
    "github.com/waku-org/go-waku/waku/v2/protocol"
)

func deleteShardingENREntries(localnode *enode.LocalNode) {
    localnode.Delete(enr.WithEntry(ShardingBitVectorEnrField, struct{}{}))
    localnode.Delete(enr.WithEntry(ShardingIndicesListEnrField, struct{}{}))
}

func WithWakuRelayShardList(rs protocol.RelayShards) ENROption {
    return func(localnode *enode.LocalNode) error {
        value, err := rs.ShardList()
        if err != nil {
            return err
        }
        deleteShardingENREntries(localnode)
        localnode.Set(enr.WithEntry(ShardingIndicesListEnrField, value))
        return nil
    }
}

func WithWakuRelayShardingBitVector(rs protocol.RelayShards) ENROption {
    return func(localnode *enode.LocalNode) error {
        deleteShardingENREntries(localnode)
        localnode.Set(enr.WithEntry(ShardingBitVectorEnrField, rs.BitVector()))
        return nil
    }
}

func WithWakuRelaySharding(rs protocol.RelayShards) ENROption {
    return func(localnode *enode.LocalNode) error {
        if len(rs.ShardIDs) >= 64 {
            return WithWakuRelayShardingBitVector(rs)(localnode)
        }

        return WithWakuRelayShardList(rs)(localnode)
    }
}

func WithWakuRelayShardingTopics(topics ...string) ENROption {
    return func(localnode *enode.LocalNode) error {
        rs, err := protocol.TopicsToRelayShards(topics...)
        if err != nil {
            return err
        }

        if len(rs) != 1 {
            return errors.New("expected a single RelayShards")
        }

        return WithWakuRelaySharding(rs[0])(localnode)
    }
}

// ENR record accessors

func RelayShardList(record *enr.Record) (*protocol.RelayShards, error) {
    var field []byte
    if err := record.Load(enr.WithEntry(ShardingIndicesListEnrField, &field)); err != nil {
        if enr.IsNotFound(err) {
            return nil, nil
        }
        return nil, err
    }

    res, err := protocol.FromShardList(field)
    if err != nil {
        return nil, err
    }

    return &res, nil
}

func RelayShardingBitVector(record *enr.Record) (*protocol.RelayShards, error) {
    var field []byte
    if err := record.Load(enr.WithEntry(ShardingBitVectorEnrField, &field)); err != nil {
        if enr.IsNotFound(err) {
            return nil, nil
        }
        return nil, err
    }

    res, err := protocol.FromBitVector(field)
    if err != nil {
        return nil, err
    }

    return &res, nil
}

func RelaySharding(record *enr.Record) (*protocol.RelayShards, error) {
    res, err := RelayShardList(record)
    if err != nil {
        return nil, err
    }

    if res != nil {
        return res, nil
    }

    return RelayShardingBitVector(record)
}

// Utils

func ContainsShard(record *enr.Record, cluster uint16, index uint16) bool {
    if index > protocol.MaxShardIndex {
        return false
    }

    rs, err := RelaySharding(record)
    if err != nil {
        return false
    }

    return rs.Contains(cluster, index)
}

func ContainsShardWithWakuTopic(record *enr.Record, topic protocol.WakuPubSubTopic) bool {
    if shardTopic, err := protocol.ToShardPubsubTopic(topic); err != nil {
        return false
    } else {
        return ContainsShard(record, shardTopic.Cluster(), shardTopic.Shard())
    }
}

func ContainsRelayShard(record *enr.Record, topic protocol.StaticShardingPubsubTopic) bool {
    return ContainsShardWithWakuTopic(record, topic)
}

func ContainsShardTopic(record *enr.Record, topic string) bool {
    shardTopic, err := protocol.ToWakuPubsubTopic(topic)
    if err != nil {
        return false
    }
    return ContainsShardWithWakuTopic(record, shardTopic)
}