From 472306f5ce8b5fd9ab515f953b8ac623f5128d7e Mon Sep 17 00:00:00 2001
From: Ludovic Chenut <ludovic@status.im>
Date: Fri, 21 Apr 2023 17:41:42 +0200
Subject: [PATCH] Added Xor-Mapped-Address attribute

---
 webrtc/stun.nim           |  8 +++++--
 webrtc/stunattributes.nim | 45 +++++++++++++++++++++++++++++++++++++--
 webrtc/webrtc.nim         | 40 ++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+), 4 deletions(-)
 create mode 100644 webrtc/webrtc.nim

diff --git a/webrtc/stun.nim b/webrtc/stun.nim
index 9ef11da..9e575ff 100644
--- a/webrtc/stun.nim
+++ b/webrtc/stun.nim
@@ -49,6 +49,7 @@ type
 
   RawStunMessage = object
     msgType: uint16
+    # it.conten.len() + 8 Because the Fingerprint is added after the encoding
     length* {.bin_value: it.content.len() + 8.}: uint16
     magicCookie: uint32
     transactionId: array[12, byte]
@@ -89,7 +90,9 @@ proc encode*(msg: StunMessage): seq[byte] =
   result.add(Binary.encode(Fingerprint.encode(result)))
 
 proc getResponse*(T: typedesc[Stun], msg: seq[byte],
-    address: TransportAddress): Option[StunMessage] =
+    ta: TransportAddress): Option[StunMessage] =
+  if ta.family != AddressFamily.IPv4 and ta.family != AddressFamily.IPv6:
+    return none(StunMessage)
   let sm =
     try:
       StunMessage.decode(msg)
@@ -112,7 +115,8 @@ proc getResponse*(T: typedesc[Stun], msg: seq[byte],
     res.attributes.add(UnknownAttribute.encode(unknownAttr))
     return some(res)
 
-  #if sm.attributes.getAttribute())
+  res.attributes.add(XorMappedAddress.encode(ta, sm.transactionId))
+  return some(res)
 
 proc new*(T: typedesc[Stun]): T =
   result = T()
diff --git a/webrtc/stunattributes.nim b/webrtc/stunattributes.nim
index f908568..e57a829 100644
--- a/webrtc/stunattributes.nim
+++ b/webrtc/stunattributes.nim
@@ -1,8 +1,11 @@
+import sequtils, typetraits
 import binary_serialization,
-       stew/byteutils
+       stew/byteutils,
+       chronos
 import utils
 
 type
+  StunAttributeEncodingError* = object of CatchableError
 # Stun Attribute
 # 0                   1                   2                   3
 # 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
@@ -100,12 +103,50 @@ proc encode*(T: typedesc[UnknownAttribute], unknownAttr: seq[uint16]): RawStunAt
                             length: value.len().uint16,
                             value: value)
 
+# Fingerprint
+
 type
   Fingerprint* = object
     crc32: uint32
 
 proc encode*(T: typedesc[Fingerprint], msg: seq[byte]): RawStunAttribute =
-  let value = Binary.encode(Fingerprint(crc32: crc32(msg) xor 0x5354554e'u32))
+  let value = Binary.encode(T(crc32: crc32(msg) xor 0x5354554e'u32))
   result = RawStunAttribute(attributeType: AttrFingerprint.uint16,
                             length: value.len().uint16,
                             value: value)
+
+# Xor Mapped Address
+
+type
+  MappedAddressFamily {.size: 1.} = enum
+    MAFIPv4 = 0x01
+    MAFIPv6 = 0x02
+
+  XorMappedAddress* = object
+    reserved: uint8 # should be 0
+    family: MappedAddressFamily
+    port: uint16
+    address: seq[byte]
+
+proc encode*(T: typedesc[XorMappedAddress], ta: TransportAddress,
+             tid: array[12, byte]): RawStunAttribute =
+  const magicCookie = @[ 0x21'u8, 0x12, 0xa4, 0x42 ]
+  let
+    address =
+      if ta.family == AddressFamily.IPv4:
+        var s = newSeq[uint8](4)
+        for i in 0..3:
+          s[i] = ta.address_v4[i] xor magicCookie[i]
+        s
+      else:
+        let magicCookieTid = magicCookie.concat(@tid)
+        var s = newSeq[uint8](16)
+        for i in 0..15:
+          s[i] = ta.address_v6[i] xor magicCookieTid[i]
+        s
+    xma = T(family: if ta.family == AddressFamily.IPv4: MAFIPv4 else: MAFIPv6,
+            port: ta.port.distinctBase xor 0x2112'u16, address: address)
+    value = Binary.encode(xma)
+  result = RawStunAttribute(attributeType: AttrXORMappedAddress.uint16,
+                            length: value.len().uint16,
+                            value: value)
diff --git a/webrtc/webrtc.nim b/webrtc/webrtc.nim
new file mode 100644
index 0000000..05f451b
--- /dev/null
+++ b/webrtc/webrtc.nim
@@ -0,0 +1,40 @@
+import chronos, chronicles
+import stun
+
+logScope:
+  topics = "webrtc"
+
+let fut = newFuture[void]()
+type
+  WebRTC* = object
+    udp: DatagramTransport
+
+proc new*(T: typedesc[WebRTC], port: uint16 = 42657): T =
+  logScope: topics = "webrtc"
+  var webrtc = T()
+  proc onReceive(udp: DatagramTransport, address: TransportAddress) {.async, gcsafe.} =
+    let
+      msg = udp.getMessage()
+    if Stun.isMessage(msg):
+      let res = Stun.getResponse(msg, address)
+      echo res
+      if res.isSome():
+        await udp.sendTo(address, res.get().encode())
+
+    trace "onReceive", isStun = Stun.isMessage(msg)
+    if not fut.completed(): fut.complete()
+
+  let
+    laddr = initTAddress("127.0.0.1:" & $port)
+    udp = newDatagramTransport(onReceive, local = laddr)
+  trace "local address", laddr
+  webrtc.udp = udp
+  return webrtc
+#
+#proc main {.async.} =
+#  echo "/ip4/127.0.0.1/udp/42657/webrtc/certhash/uEiDKBGpmOW3zQhiCHagHZ8igwfKNIp8rQCJWd5E5mIhGHw/p2p/12D3KooWFjMiMZLaCKEZRvMqKp5qUGduS6iBZ9RWQgYZXYtAAaPC"
+#  discard WebRTC.new()
+#  await fut
+#  await sleepAsync(10.seconds)
+#
+#waitFor(main())