From 77d128618ac60231d491fd72e01198567927354a Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Sat, 23 Sep 2017 15:08:41 +0200
Subject: [PATCH] Add pad() and unpad()

---
 Sodium/Utils.swift            | 57 ++++++++++++++++++++++++++++++-----
 SodiumTests/SodiumTests.swift |  8 +++++
 2 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/Sodium/Utils.swift b/Sodium/Utils.swift
index f7fed2c..c23dd4b 100644
--- a/Sodium/Utils.swift
+++ b/Sodium/Utils.swift
@@ -167,14 +167,14 @@ public class Utils {
         let ignore_cstr = ignore != nil ? (ignore! as NSString).utf8String : nil
         
         let result = binData.withUnsafeMutableBytes { binPtr in
-            return b64Data.withUnsafeBytes { b64Ptr in
-                return sodium_base642bin(binPtr,
-                                         binDataCapacity,
-                                         b64Ptr,
-                                         b64DataLen,
-                                         ignore_cstr,
-                                         &binDataLen,
-                                         nil, variant.rawValue)
+            b64Data.withUnsafeBytes { b64Ptr in
+                sodium_base642bin(binPtr,
+                                  binDataCapacity,
+                                  b64Ptr,
+                                  b64DataLen,
+                                  ignore_cstr,
+                                  &binDataLen,
+                                  nil, variant.rawValue)
             }
         }
         
@@ -185,4 +185,45 @@ public class Utils {
         binData.count = Int(binDataLen)
         return binData
     }
+    
+    /*
+     Add padding to `data` so that its length becomes a multiple of `blockSize`
+     
+     - Parameter data: input/output buffer, will be modified in-place
+     - Parameter blocksize: the block size
+     */
+    public func pad(data: inout Data, blockSize: Int) -> ()? {
+        let dataCount = data.count
+        data.reserveCapacity(dataCount + blockSize)
+        data.count = dataCount + blockSize
+        var paddedLen: size_t = 0
+        let result = data.withUnsafeMutableBytes { dataPtr in
+            sodium_pad(&paddedLen, dataPtr, dataCount, blockSize, dataCount + blockSize)
+        }
+        if result != 0 {
+            return nil
+        }
+        data.count = Int(paddedLen)
+        
+        return ()
+    }
+    
+    /*
+     Remove padding from `data` to restore its original size
+     
+     - Parameter data: input/output buffer, will be modified in-place
+     - Parameter blocksize: the block size
+     */
+    public func unpad(data: inout Data, blockSize: Int) -> ()? {
+        var unpaddedLen: size_t = 0
+        let result = data.withUnsafeMutableBytes { dataPtr in
+            sodium_unpad(&unpaddedLen, dataPtr, data.count, blockSize)
+        }
+        if result != 0 {
+            return nil
+        }
+        data.count = Int(unpaddedLen)
+        
+        return ()
+    }
 }
diff --git a/SodiumTests/SodiumTests.swift b/SodiumTests/SodiumTests.swift
index bb52a0f..b10c899 100644
--- a/SodiumTests/SodiumTests.swift
+++ b/SodiumTests/SodiumTests.swift
@@ -355,4 +355,12 @@ class SodiumTests: XCTestCase {
         XCTAssertEqual(b64_nopad, "dGVzdA")
         XCTAssertEqual(bin2_nopad, bin)
     }
+    
+    func testPad() {
+        var data = "test".toData()!
+        sodium.utils.pad(data: &data, blockSize: 16)!
+        XCTAssertTrue(data.count % 16 == 0)
+        sodium.utils.unpad(data: &data, blockSize: 16)!
+        XCTAssertTrue(data.count == 4)
+    }
 }