Segwit Support

Sources/BinaryCoding.swift

@@ -14,6 +14,18 @@ public protocol BinaryEncoding {
     func encode(into data: inout Data)
 }
 
+public extension BinaryEncoding {
+    var encoded: Data {
+        var data = Data()
+        encode(into: &data)
+        return data
+    }
+
+    var hexEncoded: String {
+        return self.encoded.hexString
+    }
+}
+
 public typealias BinaryCoding = BinaryDecoding & BinaryEncoding
 
 extension Data: BinaryEncoding {

Sources/Bitcoin/BitcoinPrivateKeyProvider.swift

@@ -0,0 +1,49 @@
+// Copyright © 2017-2018 Trust.
+//
+// This file is part of Trust. The full Trust copyright notice, including
+// terms governing use, modification, and redistribution, is contained in the
+// file LICENSE at the root of the source code distribution tree.
+
+import Foundation
+
+public protocol BitcoinPrivateKeyProvider {
+    /// Should return the private key for a public key hash
+    func key(forPublicKeyHash: Data) -> PrivateKey?
+
+    /// Should return the private key for a script hash
+    func key(forScriptHash: Data) -> PrivateKey?
+
+    /// Should return the full redeem script for a script hash
+    func script(forScriptHash: Data) -> BitcoinScript?
+}
+
+extension BitcoinPrivateKeyProvider {
+    func key(forPublicKey pubkey: PublicKey) -> PrivateKey? {
+        return key(forPublicKeyHash: pubkey.bitcoinKeyHash)
+    }
+}
+
+public final class BitcoinDefaultPrivateKeyProvider: BitcoinPrivateKeyProvider {
+    public var keys: [PrivateKey]
+    public var keysByScriptHash = [Data: PrivateKey]()
+    public var scriptsByScriptHash = [Data: BitcoinScript]()
+
+    public init(keys: [PrivateKey]) {
+        self.keys = keys
+    }
+
+    public func key(forPublicKeyHash hash: Data) -> PrivateKey? {
+        return keys.first { key in
+            let publicKey = key.publicKey(compressed: true)
+            return publicKey.bitcoinKeyHash == hash
+        }
+    }
+
+    public func key(forScriptHash hash: Data) -> PrivateKey? {
+        return keysByScriptHash[hash]
+    }
+
+    public func script(forScriptHash hash: Data) -> BitcoinScript? {
+        return scriptsByScriptHash[hash]
+    }
+}

Sources/Bitcoin/BitcoinScript+Parsing.swift

@@ -0,0 +1,76 @@
+// Copyright © 2017-2018 Trust.
+//
+// This file is part of Trust. The full Trust copyright notice, including
+// terms governing use, modification, and redistribution, is contained in the
+// file LICENSE at the root of the source code distribution tree.
+
+import Foundation
+
+public extension BitcoinScript {
+    /// Returns a single operation at the given index including its operand.
+    ///
+    /// - Parameters:
+    ///   - index: index where the operation starts, on return the index of the next operation.
+    /// - Returns: a tuple with the opcode and its operand or `nil` if's the end of the script or the script is invalid.
+    func getScriptOp(index: inout Int) -> (opcode: UInt8, operand: Data?)? {
+        // Read instruction
+        if index >= bytes.endIndex {
+            return nil
+        }
+
+        let opcode = bytes[index]
+        index += 1
+
+        if opcode > OpCode.OP_PUSHDATA4 {
+            return (opcode: opcode, operand: nil)
+        }
+
+        // Immediate operand
+        var size = 0
+        if opcode < OpCode.OP_PUSHDATA1 {
+            size = Int(opcode)
+        } else if opcode == OpCode.OP_PUSHDATA1 {
+            if bytes.endIndex - index < 1 {
+                return nil
+            }
+            size = index
+            index += 1
+        } else if opcode == OpCode.OP_PUSHDATA2 {
+            if bytes.endIndex - index < 2 {
+                return nil
+            }
+            size = Int(readLE16(at: index))
+            index += 2
+        } else if opcode == OpCode.OP_PUSHDATA4 {
+            if bytes.endIndex - index < 4 {
+                return nil
+            }
+            size = Int(readLE32(at: index))
+            index += 4
+        }
+        if bytes.endIndex - index < size {
+            return nil
+        }
+        index += size
+
+        return (opcode: opcode, operand: data[index ..< index + size])
+    }
+
+    /// Reads a little-endian `UInt16` from the script.
+    private func readLE16(at index: Int) -> UInt16 {
+        return bytes.withUnsafeBufferPointer { ptr in
+            (ptr.baseAddress! + index).withMemoryRebound(to: UInt16.self, capacity: 1) { intptr in
+                UInt16(littleEndian: intptr.pointee)
+            }
+        }
+    }
+
+    /// Reads a little-endian `UInt32` from the script.
+    private func readLE32(at index: Int) -> UInt32 {
+        return bytes.withUnsafeBufferPointer { ptr in
+            (ptr.baseAddress! + index).withMemoryRebound(to: UInt32.self, capacity: 1) { intptr in
+                UInt32(littleEndian: intptr.pointee)
+            }
+        }
+    }
+}

Sources/Bitcoin/BitcoinScript.swift

@@ -150,7 +150,7 @@ enum OpCode {
 }
 
 /// Serialized Bitcoin script.
-public final class BitcoinScript: BinaryEncoding {
+public final class BitcoinScript: BinaryEncoding, CustomDebugStringConvertible {
     public var bytes: [UInt8]
 
     public var data: Data {
@@ -170,6 +170,10 @@ public final class BitcoinScript: BinaryEncoding {
         self.bytes = Array(data)
     }
 
+    public var debugDescription: String {
+        return data.hexString
+    }
+
     var isPayToScriptHash: Bool {
         // Extra-fast test for pay-to-script-hash
         return bytes.count == 23 &&
@@ -180,33 +184,30 @@ public final class BitcoinScript: BinaryEncoding {
 
     var isPayToWitnessScriptHash: Bool {
         // Extra-fast test for pay-to-witness-script-hash
-        return bytes.count == 34 &&
+        return bytes.count == 22 &&
             bytes[0] == OpCode.OP_0 &&
-            bytes[1] == 0x20
+            bytes[1] == 0x14
     }
 
-    // A witness program is any valid CScript that consists of a 1-byte push opcode
-    // followed by a data push between 2 and 40 bytes.
-    func isWitnessProgram(version: inout Int, program: inout Data) -> Bool {
+    /// Returns the version and witness programm if this is a witness script.
+    func witnessProgram() -> (version: Int, program: Data)? {
         if bytes.count < 4 || bytes.count > 42 {
-            return false
+            return nil
         }
         if bytes[0] != OpCode.OP_0 && (bytes[0] < OpCode.OP_1 || bytes[0] > OpCode.OP_16) {
-            return false
+            return nil
         }
         if bytes[1] + 2 == bytes.count {
-            version = BitcoinScript.decodeNumber(opcode: bytes[0])
-            program = Data(bytes: bytes[2...])
-            return true
+            let version = BitcoinScript.decodeNumber(opcode: bytes[0])
+            let program = Data(bytes: bytes[2...])
+            return (version: version, program: program)
         }
-        return false
+        return nil
     }
 
     func isPushOnly(at index: inout Int) -> Bool {
         while index < bytes.endIndex {
-            var opcode = 0 as UInt8
-            var contents = Data()
-            if !getScriptOp(index: &index, opcode: &opcode, contents: &contents) {
+            guard let (opcode, _) = getScriptOp(index: &index) else {
                 return false
             }
             if opcode > OpCode.OP_16 {
@@ -217,15 +218,41 @@ public final class BitcoinScript: BinaryEncoding {
     }
 
     /// Builds a standard 'pay to public key hash' script.
-    public static func buildPayToPublicKeyHash(_ pubKeyHash: Data) -> BitcoinScript {
+    public static func buildPayToPublicKeyHash(address: BitcoinAddress) -> BitcoinScript {
+        let pubKeyHash = address.data.dropFirst()
         var data = Data(capacity: 5 + pubKeyHash.count)
         data.append(contentsOf: [OpCode.OP_DUP, OpCode.OP_HASH160])
         data.append(UInt8(pubKeyHash.count))
-        data.append(pubKeyHash)
+        data.append([UInt8](pubKeyHash), count: pubKeyHash.count)
         data.append(contentsOf: [OpCode.OP_EQUALVERIFY, OpCode.OP_CHECKSIG])
         return BitcoinScript(data: data)
     }
 
+    /// Builds a standard 'pay to script hash' script.
+    public static func buildPayToScriptHash(script: BitcoinScript) -> BitcoinScript {
+        return buildPayToScriptHash(scriptHash: Crypto.sha256ripemd160(script.data))
+    }
+
+    /// Builds a standard 'pay to script hash' script.
+    public static func buildPayToScriptHash(scriptHash: Data) -> BitcoinScript {
+        var data = Data()
+        data.append(OpCode.OP_HASH160)
+        data.append(UInt8(scriptHash.count))
+        data.append([UInt8](scriptHash), count: scriptHash.count)
+        data.append(OpCode.OP_EQUAL)
+        return BitcoinScript(data: data)
+    }
+
+    /// Builds a standard segwit script.
+    public static func buildSegWit(address: BitcoinAddress) -> BitcoinScript {
+        let pubKeyHash = address.data.dropFirst()
+        var data = Data()
+        data.append(OpCode.OP_0)
+        data.append(UInt8(pubKeyHash.count))
+        data.append([UInt8](pubKeyHash), count: pubKeyHash.count)
+        return BitcoinScript(data: data)
+    }
+
     /// Decodes a small integer
     static func decodeNumber(opcode: UInt8) -> Int {
         if opcode == OpCode.OP_0 {
@@ -244,71 +271,9 @@ public final class BitcoinScript: BinaryEncoding {
         return OpCode.OP_1 + UInt8(n - 1)
     }
 
-    func getScriptOp(index: inout Int, opcode opcodeRet: inout UInt8, contents: inout Data) -> Bool {
-        let end = bytes.endIndex
-        opcodeRet = OpCode.OP_INVALIDOPCODE
-        contents.removeAll()
-
-        // Read instruction
-        if index >= end {
-            return false
-        }
-        let opcode = bytes[index]
-        index += 1
-
-        // Immediate operand
-        if opcode <= OpCode.OP_PUSHDATA4 {
-            var size = 0
-            if opcode < OpCode.OP_PUSHDATA1 {
-                size = Int(opcode)
-            } else if opcode == OpCode.OP_PUSHDATA1 {
-                if end - index < 1 {
-                    return false
-                }
-                size = index
-                index += 1
-            } else if opcode == OpCode.OP_PUSHDATA2 {
-                if end - index < 2 {
-                    return false
-                }
-                size = Int(readLE16(at: index))
-                index += 2
-            } else if opcode == OpCode.OP_PUSHDATA4 {
-                if end - index < 4 {
-                    return false
-                }
-                size = Int(readLE32(at: index))
-                index += 4
-            }
-            if end - index < size {
-                return false
-            }
-            contents.append(data[index ..< index + size])
-            index += size
-        }
-
-        opcodeRet = opcode
-        return true
-    }
-
-    /// Reads a little-endian `UInt16` from the script.
-    private func readLE16(at index: Int) -> UInt16 {
-        return bytes.withUnsafeBufferPointer { ptr in
-            (ptr.baseAddress! + index).withMemoryRebound(to: UInt16.self, capacity: 1) { intptr in
-                UInt16(littleEndian: intptr.pointee)
-            }
-        }
-    }
-
-    /// Reads a little-endian `UInt32` from the script.
-    private func readLE32(at index: Int) -> UInt32 {
-        return bytes.withUnsafeBufferPointer { ptr in
-            (ptr.baseAddress! + index).withMemoryRebound(to: UInt32.self, capacity: 1) { intptr in
-                UInt32(littleEndian: intptr.pointee)
-            }
-        }
-    }
-
+    /// Matches the script to a pay-to-public-key (P2PK) script.
+    ///
+    /// - Returns: the public key.
     public func matchPayToPubkey() -> PublicKey? {
         if bytes.count == PublicKey.uncompressedSize + 2 && bytes[0] == PublicKey.uncompressedSize && bytes.last == OpCode.OP_CHECKSIG {
             let pubkeyData = Data(bytes: bytes[bytes.startIndex + 1 ..< bytes.startIndex + PublicKey.uncompressedSize + 1])
@@ -321,13 +286,46 @@ public final class BitcoinScript: BinaryEncoding {
         return nil
     }
 
+    /// Matches the script to a pay-to-public-key-hash (P2PKH).
+    ///
+    /// - Returns: the key hash.
     public func matchPayToPubkeyHash() -> Data? {
         if bytes.count == 25 && bytes[0] == OpCode.OP_DUP && bytes[1] == OpCode.OP_HASH160 && bytes[2] == 20 && bytes[23] == OpCode.OP_EQUALVERIFY && bytes[24] == OpCode.OP_CHECKSIG {
             return Data(bytes: bytes[bytes.startIndex + 3 ..< bytes.startIndex + 23])
         }
         return nil
     }
 
+    /// Matches the script to a pay-to-script-hash (P2SH).
+    ///
+    /// - Returns: the script hash.
+    public func matchPayToScriptHash() -> Data? {
+        guard isPayToScriptHash else {
+            return nil
+        }
+        return Data(bytes: bytes[2 ..< 22])
+    }
+
+    /// Matches the script to a pay-to-witness-public-key-hash (P2WPKH).
+    ///
+    /// - Returns: the key hash.
+    public func matchPayToWitnessPublicKeyHash() -> Data? {
+        if bytes.count == 22 && bytes[0] == OpCode.OP_0 && bytes[1] == 0x14 {
+            return Data(bytes: bytes[2...])
+        }
+        return nil
+    }
+
+    /// Matches the script to a pay-to-witness-script-hash (P2WSH).
+    ///
+    /// - Returns: the script hash, a SHA256 of the witness script.
+    public func matchPayToWitnessScriptHash() -> Data? {
+        if bytes.count == 34 && bytes[0] == OpCode.OP_0 && bytes[1] == 0x20 {
+            return Data(bytes: bytes[2...])
+        }
+        return nil
+    }
+
     public func matchMultisig(required: inout Int) -> [PublicKey]? {
         if bytes.count < 1 || bytes.last != OpCode.OP_CHECKMULTISIG {
             return []
@@ -336,13 +334,11 @@ public final class BitcoinScript: BinaryEncoding {
         var keys = [PublicKey]()
 
         var it = bytes.startIndex
-        var opcode = 0 as UInt8
-        var contents = Data()
-        if !getScriptOp(index: &it, opcode: &opcode, contents: &contents) || !OpCode.isSmallInteger(opcode) {
+        guard let (opcode, _) = getScriptOp(index: &it), OpCode.isSmallInteger(opcode) else {
             return nil
         }
         required = BitcoinScript.decodeNumber(opcode: opcode)
-        while getScriptOp(index: &it, opcode: &opcode, contents: &contents), let key = PublicKey(data: contents) {
+        while case .some(_, let data?) = getScriptOp(index: &it), let key = PublicKey(data: data) {
             keys.append(key)
         }
         if !OpCode.isSmallInteger(opcode) {