spb/os-vault Public

1import Foundation2import TweetNacl34public struct Transaction: Encodable, Equatable {5    public var signatures = [Signature]()6    public var feePayer: PublicKey?7    public var instructions = [TransactionInstruction]()8    public var recentBlockhash: String?9//        TODO: nonceInfo1011    public init() {}12    public init(13        instructions: [TransactionInstruction],14        recentBlockhash: String? = nil,15        feePayer: PublicKey? = nil16    ) {17        self.init()18        self.instructions = instructions19        self.recentBlockhash = recentBlockhash20        self.feePayer = feePayer21    }2223    // MARK: - Methods2425    public mutating func sign(signers: [KeyPair]) throws {26        guard !signers.isEmpty else {27            throw VersionedTransactionError.noSigner28        }2930        // unique signers31        let signers = signers.reduce([KeyPair]()) { signers, signer in32            var uniqueSigners = signers33            if !uniqueSigners.contains(where: { $0.publicKey == signer.publicKey }) {34                uniqueSigners.append(signer)35            }36            return uniqueSigners37        }3839        // map signatures40        signatures = signers.map { Signature(signature: nil, publicKey: $0.publicKey) }4142        // construct message43        let message = try compile()4445        try partialSign(message: message, signers: signers)46    }4748    public func calculateTransactionFee(lamportsPerSignatures: UInt64) throws -> UInt64 {49        var transaction = self50        transaction.recentBlockhash = "BdA9gRatFvvwszr9uU5fznkHoMVQE8tf6ZFi8Mp6xdKs" // fake51        let message = try transaction.compile()52        return UInt64(message.header.numRequiredSignatures) * lamportsPerSignatures53    }5455    public mutating func serialize(56        requiredAllSignatures: Bool = true,57        verifySignatures: Bool = false58    ) throws -> Data {59        // message60        let serializedMessage = try serializeMessage()6162        // verification63        if verifySignatures, !_verifySignatures(64            serializedMessage: serializedMessage,65            requiredAllSignatures: requiredAllSignatures66        ) {67            throw VersionedTransactionError.signatureVerificationError68        }6970        return _serialize(serializedMessage: serializedMessage)71    }7273    // MARK: - Helpers7475    public mutating func addSignature(_ signature: Signature) throws {76        let _ = try compile() // Ensure signatures array is populated7778        try _addSignature(signature)79    }8081    mutating func serializeMessage() throws -> Data {82        try compile().serialize()83    }8485    mutating func verifySignatures() throws -> Bool {86        try _verifySignatures(serializedMessage: serializeMessage(), requiredAllSignatures: true)87    }8889    public func findSignature(pubkey: PublicKey) -> Signature? {90        signatures.first(where: { $0.publicKey == pubkey })91    }9293    // MARK: - Signing9495    public mutating func partialSign(signers: [KeyPair]) throws {96        guard !signers.isEmpty else {97            throw VersionedTransactionError.noSigner98        }99100        // unique signers101        let signers = signers.reduce([KeyPair]()) { signers, signer in102            var uniqueSigners = signers103            if !uniqueSigners.contains(where: { $0.publicKey == signer.publicKey }) {104                uniqueSigners.append(signer)105            }106            return uniqueSigners107        }108109        // construct message110        let message = try compile()111112        try partialSign(message: message, signers: signers)113    }114115    private mutating func partialSign(message: Message, signers: [KeyPair]) throws {116        let signData = try message.serialize()117118        for signer in signers {119            let data = try NaclSign.signDetached(message: signData, secretKey: signer.secretKey)120            try _addSignature(Signature(signature: data, publicKey: signer.publicKey))121        }122    }123124    private mutating func _addSignature(_ signature: Signature) throws {125        guard let data = signature.signature,126              data.count == 64,127              let index = signatures.firstIndex(where: { $0.publicKey == signature.publicKey })128        else {129            throw VersionedTransactionError.invalidSigner(signature.publicKey.base58EncodedString)130        }131132        signatures[index] = signature133    }134135    // MARK: - Compiling136137    private mutating func compile() throws -> Message {138        let message = try compileMessage()139        let signedKeys = message.accountKeys[0 ..< Int(message.header.numRequiredSignatures)]140141        if signatures.count == signedKeys.count {142            var isValid = true143            for (index, signature) in signatures.enumerated() {144                if signedKeys[index] != signature.publicKey {145                    isValid = false146                    break147                }148            }149            if isValid {150                return message151            }152        }153154        signatures = signedKeys.map { Signature(signature: nil, publicKey: $0) }155        return message156    }157158    public func compileMessage() throws -> Message {159        // verify instructions160        guard !instructions.isEmpty else {161            throw VersionedTransactionError.noInstructionProvided162        }163        guard let feePayer = feePayer else {164            throw VersionedTransactionError.feePayerNotFound165        }166        guard let recentBlockhash = recentBlockhash else {167            throw VersionedTransactionError.recentBlockhashNotFound168        }169170        // programIds & accountMetas171        var programIds = [PublicKey]()172        var accountMetas = [AccountMeta]()173174        for instruction in instructions {175            accountMetas.append(contentsOf: instruction.keys)176            if !programIds.contains(instruction.programId) {177                programIds.append(instruction.programId)178            }179        }180181        // Append programID account metas182        for programId in programIds {183            accountMetas.append(184                .init(publicKey: programId, isSigner: false, isWritable: false)185            )186        }187188        // sort accountMetas, first by signer, then by writable189        accountMetas.sort { x, y -> Bool in190            if x.isSigner != y.isSigner { return x.isSigner }191            if x.isWritable != y.isWritable { return x.isWritable }192            return false193        }194195        // filterOut duplicate account metas, keeps writable one196        accountMetas = accountMetas.reduce([AccountMeta]()) { result, accountMeta in197            var uniqueMetas = result198            if let index = uniqueMetas.firstIndex(where: { $0.publicKey == accountMeta.publicKey }) {199                // if accountMeta exists200                uniqueMetas[index].isWritable = uniqueMetas[index].isWritable || accountMeta.isWritable201            } else {202                uniqueMetas.append(accountMeta)203            }204            return uniqueMetas205        }206207        // Cull duplicate account metas208        var uniqueMetas: [AccountMeta] = []209        accountMetas.forEach { accountMeta in210            let pubkey = accountMeta.publicKey.base58EncodedString211            let uniqueIndex = uniqueMetas.firstIndex { x in x.publicKey.base58EncodedString == pubkey }212            if let uniqueIndex = uniqueIndex {213                uniqueMetas[uniqueIndex].isWritable = uniqueMetas[uniqueIndex].isWritable || accountMeta.isWritable214            } else {215                uniqueMetas.append(accountMeta)216            }217        }218219        // move fee payer to front220        let feePayerIndex = uniqueMetas.firstIndex { x in x.publicKey == feePayer }221        if let feePayerIndex = feePayerIndex {222            var payerMeta = uniqueMetas.remove(at: feePayerIndex)223            payerMeta.isSigner = true224            payerMeta.isWritable = true225            uniqueMetas.insert(payerMeta, at: 0)226        } else {227            uniqueMetas.insert(228                AccountMeta(229                    publicKey: feePayer,230                    isSigner: true,231                    isWritable: true232                ),233                at: 0234            )235        }236237        // accountMetas.removeAll(where: { $0.publicKey == feePayer })238        // accountMetas.insert(239        //    AccountMeta(publicKey: feePayer, isSigner: true, isWritable: true),240        //    at: 0241        // )242243        // verify signers244        for signature in signatures {245            if let index = try? uniqueMetas.index(ofElementWithPublicKey: signature.publicKey) {246                if !uniqueMetas[index].isSigner {247                    // TODO: check248                    uniqueMetas[index].isSigner = true249//                        accountMetas[index].isSigner = true250//                        Logger.log(message: "Transaction references a signature that is unnecessary, only the fee payer and instruction signer accounts should sign a transaction. This behavior is deprecated and will throw an error in the next major version release.", event: .warning)251                    print("WARN: Transaction references a signature that is unnecessary")252//                        throw Error.invalidRequest(reason: "Transaction references a signature that is unnecessary")253                }254            } else {255                throw VersionedTransactionError.unknownSigner(signature.publicKey.base58EncodedString)256            }257        }258259        // header260        var header = MessageHeader()261262        var signedKeys = [AccountMeta]()263        var unsignedKeys = [AccountMeta]()264265        uniqueMetas.forEach { accountMeta in266            // signed keys267            if accountMeta.isSigner {268                signedKeys.append(accountMeta)269                header.numRequiredSignatures += 1270271                if !accountMeta.isWritable {272                    header.numReadonlySignedAccounts += 1273                }274            }275276            // unsigned keys277            else {278                unsignedKeys.append(accountMeta)279280                if !accountMeta.isWritable {281                    header.numReadonlyUnsignedAccounts += 1282                }283            }284        }285286        accountMetas = signedKeys + unsignedKeys287        let accountKeys = accountMetas.map(\.publicKey)288        let instructions = instructions.compile(accountKeys: accountKeys)289        try instructions.forEach { instruction in290            try instruction.accounts.forEach { keyIndex in291                if keyIndex < 0 { throw VersionedTransactionError.unknown }292            }293        }294295        return Message(296            header: header,297            accountKeys: accountKeys,298            recentBlockhash: recentBlockhash,299            instructions: instructions300        )301    }302303    // MARK: - Verifying304305    private mutating func _verifySignatures(306        serializedMessage: Data,307        requiredAllSignatures: Bool308    ) -> Bool {309        for signature in signatures {310            if signature.signature == nil {311                if requiredAllSignatures {312                    return false313                }314            } else {315                if (try? NaclSign.signDetachedVerify(316                    message: serializedMessage,317                    sig: signature.signature!,318                    publicKey: signature.publicKey.data319                )) != true {320                    return false321                }322            }323        }324        return true325    }326327    // MARK: - Serializing328329    private mutating func _serialize(serializedMessage: Data) -> Data {330        // signature length331        let signaturesLength = signatures.count332333        // signature data334        let signaturesData = signatures.reduce(Data()) { result, signature in335            var data = result336            if let signature = signature.signature {337                data.append(signature)338            } else {339                data.append(Constants.defaultSignature)340            }341            return data342        }343344        let encodedSignatureLength = Data.encodeLength(signaturesLength)345346        // transaction length347        var data = Data(capacity: encodedSignatureLength.count + signaturesData.count + serializedMessage.count)348        data.append(encodedSignatureLength)349        data.append(signaturesData)350        data.append(serializedMessage)351        return data352    }353354    public static func from(data: Data) throws -> Transaction {355        var data = data356        var signatures: [String] = []357        let signatureCount = data.decodeLength()358359        for _ in stride(from: 0, through: signatureCount - 1, by: 1) {360            let signatureData = data.prefix(Constants.signatureLength)361            data = data.dropFirst(Constants.signatureLength)362            signatures.append(Base58.encode(signatureData))363        }364365        print(data.base64EncodedString())366        return try populate(Message.from(data: data), signatures)367    }368369    static func populate(_ message: Message, _ signatures: [String]) -> Transaction {370        var transaction = Transaction()371372        transaction.recentBlockhash = message.recentBlockhash373        if message.header.numRequiredSignatures > 0 {374            transaction.feePayer = message.accountKeys[0]375        }376        signatures.enumerated().forEach { index, signature in377            let sigPubkeyPair = Signature(378                signature: signature == Base58379                    .encode(Constants.defaultSignature) ? nil : Data(Base58.decode(signature)),380                publicKey: message.accountKeys[index]381            )382            transaction.signatures.append(sigPubkeyPair)383        }384385        message.instructions.forEach { instruction in386            let keys: [AccountMeta] = instruction.accounts.map { account in387                let pubkey = message.accountKeys[account]388                return AccountMeta(389                    publicKey: pubkey,390                    isSigner: transaction.signatures.contains { keyObj in keyObj.publicKey == pubkey } || message391                        .isAccountSigner(index: account),392                    isWritable: message.isAccountWritable(index: account)393                )394            }395396            transaction.instructions.append(397                TransactionInstruction(398                    keys: keys,399                    programId: message.accountKeys[instruction.programIdIndexValue],400                    data: instruction.data401                )402            )403        }404405        return transaction406    }407}408409public struct Signature: Encodable, Equatable {410    public var signature: Data?411    public var publicKey: PublicKey412413    enum CodingKeys: String, CodingKey {414        case signature, publicKey415    }416417    public init(signature: Data?, publicKey: PublicKey) {418        self.signature = signature419        self.publicKey = publicKey420    }421422    public func encode(to encoder: Encoder) throws {423        var container = encoder.container(keyedBy: CodingKeys.self)424        try container.encode(Base58.encode(signature?.bytes ?? []), forKey: .signature)425        try container.encode(publicKey.base58EncodedString, forKey: .publicKey)426    }427}428429private extension Array where Element == AccountMeta {430    func index(ofElementWithPublicKey publicKey: PublicKey) throws -> Int {431        guard let index = firstIndex(where: { $0.publicKey == publicKey })432        else { throw VersionedTransactionError.unknown }433        return index434    }435}436