nwaku/waku/common/enr/builder.nim

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


import
  std/options,
  stew/results,
  stew/shims/net,
  eth/keys as eth_keys,
  eth/p2p/discoveryv5/enr,
  libp2p/crypto/crypto as libp2p_crypto


## Builder

type EnrBuilder* = object
    seqNumber: uint64
    privateKey: eth_keys.PrivateKey
    fields: seq[FieldPair]


proc init*(T: type EnrBuilder, key: eth_keys.PrivateKey, seqNum: uint64 = 1): T =
  EnrBuilder(
    seqNumber: seqNum,
    privateKey: key,
    fields: newSeq[FieldPair]()
  )

proc init*(T: type EnrBuilder, key: libp2p_crypto.PrivateKey, seqNum: uint64 = 1): T =
  # TODO: Inconvenient runtime assertion. Move this assertion to compile time
  if key.scheme != PKScheme.Secp256k1:
    raise newException(Defect, "invalid private key scheme")

  let
    bytes = key.getRawBytes().expect("Private key is valid")
    privateKey = eth_keys.PrivateKey.fromRaw(bytes).expect("Raw private key is of valid length")

  EnrBuilder.init(key=privateKey, seqNum=seqNum)

proc addFieldPair*(builder: var EnrBuilder, pair: FieldPair) =
  builder.fields.add(pair)

proc addFieldPair*[V](builder: var EnrBuilder, key: string, value: V) =
  builder.addFieldPair(toFieldPair(key, value))

proc build*(builder: EnrBuilder): EnrResult[enr.Record] =
  # Note that nim-eth's `Record.init` does not deduplicate the field pairs.
  # See: https://github.com/status-im/nim-eth/blob/4b22fcd/eth/p2p/discoveryv5/enr.nim#L143-L144
  enr.Record.init(
    seqNum = builder.seqNumber,
    pk = builder.privateKey,
    ip = none(IpAddress),
    tcpPort = none(Port),
    udpPort = none(Port),
    extraFields = builder.fields
  )


## Builder extension: IP address and TCP/UDP ports

proc addAddressAndPorts(builder: var EnrBuilder, ip: IpAddress, tcpPort, udpPort: Option[Port]) =
  # Based on: https://github.com/status-im/nim-eth/blob/4b22fcd/eth/p2p/discoveryv5/enr.nim#L166
  let isV6 = ip.family == IPv6

  let ipField = if isV6: toFieldPair("ip6", ip.address_v6)
                else: toFieldPair("ip", ip.address_v4)
  builder.addFieldPair(ipField)

  if tcpPort.isSome():
    let
      tcpPortFieldKey = if isV6: "tcp6" else: "tcp"
      tcpPortFieldValue = tcpPort.get()
    builder.addFieldPair(tcpPortFieldKey, tcpPortFieldValue.uint16)

  if udpPort.isSome():
    let
      udpPortFieldKey = if isV6: "udp6" else: "udp"
      udpPortFieldValue = udpPort.get()
    builder.addFieldPair(udpPortFieldKey, udpPortFieldValue.uint16)

proc addPorts(builder: var EnrBuilder, tcp, udp: Option[Port]) =
  # Based on: https://github.com/status-im/nim-eth/blob/4b22fcd/eth/p2p/discoveryv5/enr.nim#L166

  if tcp.isSome():
    let tcpPort = tcp.get()
    builder.addFieldPair("tcp", tcpPort.uint16)

  if udp.isSome():
    let udpPort = udp.get()
    builder.addFieldPair("udp", udpPort.uint16)


proc withIpAddressAndPorts*(builder: var EnrBuilder,
                            ipAddr = none(IpAddress),
                            tcpPort = none(Port),
                            udpPort = none(Port)) =
  if ipAddr.isSome():
    addAddressAndPorts(builder, ipAddr.get(), tcpPort, udpPort)
  else:
    addPorts(builder, tcpPort, udpPort)