From 26ef9a982315ec8d50eb18a71e009c9b50e579dd Mon Sep 17 00:00:00 2001
From: Alexander Kasimir <alexander@kasimir.at>
Date: Tue, 30 Jul 2019 22:14:44 +0200
Subject: [PATCH] added canonical cbor map key sorting

---
 Sources/SwiftCBOR/CBOREncoder.swift         | 38 ++++++++++++++++++++-
 Tests/SwiftCBORTests/CBOREncoderTests.swift |  6 ++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/Sources/SwiftCBOR/CBOREncoder.swift b/Sources/SwiftCBOR/CBOREncoder.swift
index 03c2c69..e498d37 100644
--- a/Sources/SwiftCBOR/CBOREncoder.swift
+++ b/Sources/SwiftCBOR/CBOREncoder.swift
@@ -123,12 +123,12 @@ extension CBOR {
     }
 
     // MARK: - major 5: a map of pairs of data items
-
     public static func encodeMap<A: CBOREncodable, B: CBOREncodable>(_ map: [A: B]) -> [UInt8] {
         var res: [UInt8] = []
         res.reserveCapacity(1 + map.count * (MemoryLayout<A>.size + MemoryLayout<B>.size + 2))
         res = map.count.encode()
         res[0] = res[0] | 0b101_00000
+
         for (k, v) in map {
             res.append(contentsOf: k.encode())
             res.append(contentsOf: v.encode())
@@ -136,6 +136,42 @@ extension CBOR {
         return res
     }
 
+    // MARK: - major 5: a map of pairs of data items sorted canonically see: https://tools.ietf.org/html/rfc7049#section-3.9
+    public static func encodeMapCanonical<A: CBOREncodable, B: CBOREncodable>(_ map: [A: B]) -> [UInt8] {
+        var res: [UInt8] = []
+        res.reserveCapacity(1 + map.count * (MemoryLayout<A>.size + MemoryLayout<B>.size + 2))
+        res = map.count.encode()
+        res[0] = res[0] | 0b101_00000
+
+        let canonicalSortedEntries = map
+            .map { k, v in (k.encode(), v.encode()) }
+            .sorted { a, b in canonicalKeyOrder(a: a.0, b: b.0) }
+
+        for (k, v) in canonicalSortedEntries {
+            res.append(contentsOf: k)
+            res.append(contentsOf: v)
+        }
+        return res
+    }
+
+
+    // MARK: sort keys canonically, see: https://tools.ietf.org/html/rfc7049#section-3.9
+    private static func canonicalKeyOrder(a: [UInt8], b: [UInt8]) -> Bool {
+        if a.count == b.count {
+            for (index, byteA) in a.enumerated() {
+                let byteB = b[index]
+                if byteA == byteB {
+                    continue
+                } else {
+                    return byteA < byteB
+                }
+            }
+            return false
+        } else {
+            return a.count < b.count
+        }
+    }
+
     public static func encodeMap<A: CBOREncodable>(_ map: [A: Any?]) throws -> [UInt8] {
         var res: [UInt8] = []
         res = map.count.encode()
diff --git a/Tests/SwiftCBORTests/CBOREncoderTests.swift b/Tests/SwiftCBORTests/CBOREncoderTests.swift
index 4bb9dab..543d4c2 100644
--- a/Tests/SwiftCBORTests/CBOREncoderTests.swift
+++ b/Tests/SwiftCBORTests/CBOREncoderTests.swift
@@ -123,6 +123,12 @@ class CBOREncoderTests: XCTestCase {
         XCTAssertEqual(encodedMapToAny, [0xa2, 0x61, 0x61, 0x01, 0x61, 0x62, 0x82, 0x02, 0x03])
     }
 
+    func testCanonicallyEncodeMaps() {
+        XCTAssertEqual(CBOR.encode(Dictionary<Int, Int>()), [0xa0])
+        let encoded = CBOR.encodeMapCanonical([1: 2, 3: 4])
+        XCTAssert(encoded == [0xa2, 0x01, 0x02, 0x03, 0x04])
+    }
+
     func testEncodeTagged() {
         let bignum: [UInt8] = [0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00] // 2**64
         let bignumCBOR = CBOR.byteString(bignum)