From 3abf7df73d7c7290ab49433aca6fc0622b4a4a74 Mon Sep 17 00:00:00 2001
From: cheatfate <ka@hardcore.kiev.ua>
Date: Mon, 3 Dec 2018 16:59:40 +0200
Subject: [PATCH] Fix base32 encoding and added more tests.

---
 libp2p/base32.nim    | 59 +++++++++++++++++++++-----------------------
 tests/testbase32.nim | 58 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+), 31 deletions(-)

diff --git a/libp2p/base32.nim b/libp2p/base32.nim
index 64b614333..f11751600 100644
--- a/libp2p/base32.nim
+++ b/libp2p/base32.nim
@@ -83,44 +83,54 @@ proc decodedLength*(btype: typedesc[Base32Types], length: int): int =
   result = (length div 8) * 5 + ((reminder * 5) div 8)
 
 proc convert5to8(inbytes: openarray[byte], outbytes: var openarray[char],
-                 length: int) {.inline.} =
+                 length: int): int {.inline.} =
   if length >= 1:
     outbytes[0] = chr(inbytes[0] shr 3)
     outbytes[1] = chr((inbytes[0] and 7'u8) shl 2)
+    result = 2
   if length >= 2:
     outbytes[1] = chr(cast[byte](outbytes[1]) or cast[byte](inbytes[1] shr 6))
     outbytes[2] = chr((inbytes[1] shr 1) and 31'u8)
     outbytes[3] = chr((inbytes[1] and 1'u8) shl 4)
+    result = 4
   if length >= 3:
     outbytes[3] = chr(cast[byte](outbytes[3]) or (inbytes[2] shr 4))
     outbytes[4] = chr((inbytes[2] and 15'u8) shl 1)
+    result = 5
   if length >= 4:
     outbytes[4] = chr(cast[byte](outbytes[4]) or (inbytes[3] shr 7))
     outbytes[5] = chr((inbytes[3] shr 2) and 31'u8)
     outbytes[6] = chr((inbytes[3] and 3'u8) shl 3)
+    result = 7
   if length >= 5:
     outbytes[6] = chr(cast[byte](outbytes[6]) or (inbytes[4] shr 5))
     outbytes[7] = chr(inbytes[4] and 31'u8)
+    result = 8
 
 proc convert8to5(inbytes: openarray[byte], outbytes: var openarray[byte],
-                 length: int) {.inline.} =
+                 length: int): int {.inline.} =
   if length >= 2:
     outbytes[0] = inbytes[0] shl 3
     outbytes[0] = outbytes[0] or (inbytes[1] shr 2)
+    result = 1
   if length >= 4:
     outbytes[1] = (inbytes[1] and 3'u8) shl 6
     outbytes[1] = outbytes[1] or (inbytes[2] shl 1)
     outbytes[1] = outbytes[1] or (inbytes[3] shr 4)
+    result = 2
   if length >= 5:
     outbytes[2] = (inbytes[3] and 15'u8) shl 4
     outbytes[2] = outbytes[2] or (inbytes[4] shr 1)
+    result = 3
   if length >= 7:
     outbytes[3] = (inbytes[4] and 1'u8) shl 7
     outbytes[3] = outbytes[3] or (inbytes[5] shl 2)
     outbytes[3] = outbytes[3] or (inbytes[6] shr 3)
+    result = 4
   if length >= 8:
     outbytes[4] = (inbytes[6] and 7'u8) shl 5
     outbytes[4] = outbytes[4] or (inbytes[7] and 31'u8)
+    result = 5
 
 proc encode*(btype: typedesc[Base32Types], inbytes: openarray[byte],
              outstr: var openarray[char], outlen: var int): Base32Status =
@@ -152,29 +162,24 @@ proc encode*(btype: typedesc[Base32Types], inbytes: openarray[byte],
   let limit = len(inbytes) - reminder
   var i, k: int
   while i < limit:
-    convert5to8(inbytes.toOpenArray(i, i + 4),
-                outstr.toOpenArray(k, k + 7), 5)
+    discard convert5to8(inbytes.toOpenArray(i, i + 4),
+                        outstr.toOpenArray(k, k + 7), 5)
     for j in 0..7:
       outstr[k + j] = chr(alphabet.encode[ord(outstr[k + j])])
     k += 8
     i += 5
 
   if reminder != 0:
-    convert5to8(inbytes.toOpenArray(i, i + reminder - 1),
-                outstr.toOpenArray(k, length - 1), reminder)
-    outstr[k] = chr(alphabet.encode[ord(outstr[k])])
-    inc(k)
-    while k < len(outstr):
-      if ord(outstr[k]) != 0:
-        outstr[k] = chr(alphabet.encode[ord(outstr[k])])
+    let left = convert5to8(inbytes.toOpenArray(i, i + reminder - 1),
+                           outstr.toOpenArray(k, length - 1), reminder)
+    for j in 0..(left - 1):
+      outstr[k] = chr(alphabet.encode[ord(outstr[k])])
+      inc(k)
+    when (btype is Base32UpperPad) or (btype is Base32LowerPad) or
+         (btype is HexBase32UpperPad) or (btype is HexBase32LowerPad):
+      while k < len(outstr):
+        outstr[k] = '='
         inc(k)
-      else:
-        when (btype is Base32UpperPad) or (btype is Base32LowerPad) or
-             (btype is HexBase32UpperPad) or (btype is HexBase32LowerPad):
-          outstr[k] = '='
-          inc(k)
-        else:
-          discard
   outlen = k
   result = Base32Status.Success
 
@@ -242,24 +247,16 @@ proc decode*[T: byte|char](btype: typedesc[Base32Types], instr: openarray[T],
         zeroMem(addr outbytes[0], i + 8)
         return Base32Status.Incorrect
       buffer[j] = cast[byte](ch)
-    convert8to5(buffer, outbytes.toOpenArray(k, k + 4), 8)
+    discard convert8to5(buffer, outbytes.toOpenArray(k, k + 4), 8)
     k += 5
     i += 8
 
-  var incsize = 0
+  var left = 0
   if reminder != 0:
     if reminder == 1 or reminder == 3 or reminder == 6:
       outlen = 0
       zeroMem(addr outbytes[0], i + 8)
       return Base32Status.Incorrect
-    elif reminder == 2:
-      incsize = 1
-    elif reminder == 4:
-      incsize = 2
-    elif reminder == 5:
-      incsize = 3
-    elif reminder == 7:
-      incsize = 4
     for j in 0..<reminder:
       if (cast[byte](instr[i + j]) and 0x80'u8) != 0:
         outlen = 0
@@ -273,9 +270,9 @@ proc decode*[T: byte|char](btype: typedesc[Base32Types], instr: openarray[T],
         result = Base32Status.Incorrect
         return
       buffer[j] = cast[byte](ch)
-    convert8to5(buffer.toOpenArray(0, reminder - 1),
-                outbytes.toOpenArray(k, length - 1), reminder)
-  outlen = k + incsize
+    left = convert8to5(buffer.toOpenArray(0, reminder - 1),
+                          outbytes.toOpenArray(k, length - 1), reminder)
+  outlen = k + left
   result = Base32Status.Success
 
 proc decode*[T: byte|char](btype: typedesc[Base32Types],
diff --git a/tests/testbase32.nim b/tests/testbase32.nim
index 02e01840f..9ee6133f2 100644
--- a/tests/testbase32.nim
+++ b/tests/testbase32.nim
@@ -148,6 +148,25 @@ suite "BASE32 encoding test suite":
       o7 == 0
       o8 == 0
 
+  test "Zero test":
+    var s = newString(256)
+    for i in 0..255:
+      s[i] = 'A'
+    var buffer: array[256, byte]
+    for i in 0..255:
+      var a = Base32.encode(buffer.toOpenArray(0, i))
+      var b = Base32.decode(a)
+      check b == buffer[0..i]
+
+  test "Leading zero test":
+    var buffer: array[256, byte]
+    for i in 0..255:
+      buffer[255] = byte(i)
+      var a = Base32.encode(buffer)
+      var b = Base32.decode(a)
+      check:
+        equalMem(addr buffer[0], addr b[0], 256) == true
+
   test "BASE32 uppercase padding test vectors":
     for item in TVBaseUpperPadding:
       let plain = cast[seq[byte]](item[0])
@@ -347,3 +366,42 @@ suite "BASE32 encoding test suite":
       check:
         d1 == plain
         d2 == plain
+
+  test "Buffer Overrun test":
+    var encres = ""
+    var encsize = 0
+    var decres: seq[byte] = @[]
+    var decsize = 0
+    check:
+      Base32.encode([0'u8], encres, encsize) == Base32Status.Overrun
+      encsize == Base32.encodedLength(1)
+      Base32.decode("AA", decres, decsize) == Base32Status.Overrun
+      decsize == Base32.decodedLength(2)
+
+  test "Incorrect test":
+    var decres = newSeq[byte](10)
+    var decsize = 0
+    check:
+      Base32.decode("A", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      Base32.decode("AAA", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      Base32.decode("AAAAAA", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      Base32Upper.decode("aa", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      Base32Upper.decode("11", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      Base32Lower.decode("AA", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      Base32Lower.decode("11", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      HexBase32Upper.decode("aa", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      HexBase32Upper.decode("WW", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      HexBase32Lower.decode("AA", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+      HexBase32Lower.decode("ww", decres, decsize) == Base32Status.Incorrect
+      decsize == 0
+