From 10b69064ee0f2343f5917fd4ee4cf9c990a2c681 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Mon, 28 Sep 2020 14:03:52 +0200
Subject: [PATCH] Write payload of short packets

---
 quic/packets.nim           |  1 +
 quic/packets/packet.nim    |  1 +
 quic/packets/write.nim     | 14 +++++++++-----
 tests/quic/testPackets.nim |  6 ++++++
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/quic/packets.nim b/quic/packets.nim
index f370824..d8767e4 100644
--- a/quic/packets.nim
+++ b/quic/packets.nim
@@ -46,6 +46,7 @@ proc write*(datagram: var Datagram, packet: Packet) =
     writer.writeKeyPhase(datagram)
     writer.writeShortDestination(datagram)
     writer.writePacketNumber(datagram)
+    writer.writePayload(datagram)
   of formLong:
     writer.writeKind(datagram)
     writer.writeVersion(datagram)
diff --git a/quic/packets/packet.nim b/quic/packets/packet.nim
index adbe4dc..fefa32a 100644
--- a/quic/packets/packet.nim
+++ b/quic/packets/packet.nim
@@ -19,6 +19,7 @@ type
     spinBit*: bool
     keyPhase*: bool
     packetnumber*: PacketNumber
+    payload*: seq[byte]
   PacketInitial* = object
     version*: uint32
     token*: seq[byte]
diff --git a/quic/packets/write.nim b/quic/packets/write.nim
index 44c89ee..0e677d3 100644
--- a/quic/packets/write.nim
+++ b/quic/packets/write.nim
@@ -55,11 +55,15 @@ proc writeIntegrity*(writer: var PacketWriter, datagram: var Datagram) =
   writer.write(datagram, writer.packet.retry.integrity)
 
 proc payload(packet: Packet): seq[byte] =
-  case packet.kind
-  of packetHandshake: packet.handshake.payload
-  of packet0RTT: packet.rtt.payload
-  of packetInitial: packet.initial.payload
-  else: @[]
+  case packet.form
+  of formShort:
+    packet.short.payload
+  of formLong:
+    case packet.kind
+    of packetHandshake: packet.handshake.payload
+    of packet0RTT: packet.rtt.payload
+    of packetInitial: packet.initial.payload
+    else: @[]
 
 proc packetNumber(packet: Packet): PacketNumber =
   case packet.form
diff --git a/tests/quic/testPackets.nim b/tests/quic/testPackets.nim
index 5613208..e1ad4c1 100644
--- a/tests/quic/testPackets.nim
+++ b/tests/quic/testPackets.nim
@@ -158,6 +158,12 @@ suite "packet writing":
     check int(datagram[0] and 0b11'u8) + 1 == sizeof(packetnumber)
     check datagram[1..2] == @[0xAA'u8, 0xBB'u8]
 
+  test "writes payload for short packet":
+    const payload = repeat(0xAB'u8, 1024)
+    var packet = Packet(form: formShort, short: PacketShort(payload: payload))
+    datagram.write(packet)
+    check datagram[2..1025] == payload
+
 suite "packet reading":
 
   var datagram: seq[byte]