nwaku/waku/waku_enr/multiaddr.nim

when (NimMajor, NimMinor) < (1, 4):
  {.push raises: [Defect].}
else:
  {.push raises: [].}

import
  std/[options, sequtils],
  stew/[endians2, results],
  stew/shims/net,
  eth/keys,
  libp2p/[multiaddress, multicodec],
  libp2p/crypto/crypto
import
  ../common/enr

const
  MultiaddrEnrField* = "multiaddrs"


func encodeMultiaddrs*(multiaddrs: seq[MultiAddress]): seq[byte] =
  var buffer = newSeq[byte]()
  for multiaddr in multiaddrs:

    let
      raw = multiaddr.data.buffer # binary encoded multiaddr
      size = raw.len.uint16.toBytes(Endianness.bigEndian) # size as Big Endian unsigned 16-bit integer

    buffer.add(concat(@size, raw))

  buffer

func readBytes(rawBytes: seq[byte], numBytes: int, pos: var int = 0): Result[seq[byte], cstring] =
  ## Attempts to read `numBytes` from a sequence, from
  ## position `pos`. Returns the requested slice or
  ## an error if `rawBytes` boundary is exceeded.
  ##
  ## If successful, `pos` is advanced by `numBytes`
  if rawBytes[pos..^1].len() < numBytes:
    return err("insufficient bytes")

  let slicedSeq = rawBytes[pos..<pos+numBytes]
  pos += numBytes

  return ok(slicedSeq)

func decodeMultiaddrs(buffer: seq[byte]): EnrResult[seq[MultiAddress]] =
  ## Parses a `multiaddrs` ENR field according to
  ## https://rfc.vac.dev/spec/31/
  var multiaddrs: seq[MultiAddress]

  var pos = 0
  while pos < buffer.len():
    let addrLenRaw = ? readBytes(buffer, 2, pos)
    let addrLen = uint16.fromBytesBE(addrLenRaw)
    if addrLen == 0:
      # Ensure pos always advances and we don't get stuck in infinite loop
      return err("malformed multiaddr field: invalid length")

    let addrRaw = ? readBytes(buffer, addrLen.int, pos)
    let address = MultiAddress.init(addrRaw).get()

    multiaddrs.add(address)

  return ok(multiaddrs)


# ENR builder extension
func stripPeerId(multiaddr: MultiAddress): MultiAddress =
  if not multiaddr.contains(multiCodec("p2p")).get():
    return multiaddr

  var cleanAddr = MultiAddress.init()
  for item in multiaddr.items:
    if item.value.protoName().get() != "p2p":
      # Add all parts except p2p peerId
      discard cleanAddr.append(item.value)

  return cleanAddr

func withMultiaddrs*(builder: var EnrBuilder, multiaddrs: seq[MultiAddress]) =
  let multiaddrs = multiaddrs.map(stripPeerId)
  let value = encodeMultiaddrs(multiaddrs)
  builder.addFieldPair(MultiaddrEnrField, value)

func withMultiaddrs*(builder: var EnrBuilder, multiaddrs: varargs[MultiAddress]) =
  withMultiaddrs(builder, @multiaddrs)

# ENR record accessors (e.g., Record, TypedRecord, etc.)

func multiaddrs*(record: TypedRecord): Option[seq[MultiAddress]] =
  let field = record.tryGet(MultiaddrEnrField, seq[byte])
  if field.isNone():
    return none(seq[MultiAddress])

  let decodeRes = decodeMultiaddrs(field.get())
  if decodeRes.isErr():
    return none(seq[MultiAddress])

  some(decodeRes.value)