spb/os-vault Public

1import Foundation23public struct MessageV0: IMessage, Equatable {4    public var version: TransactionVersion { .v0 }56    public var header: MessageHeader7    public var staticAccountKeys: [PublicKey]8    public var recentBlockhash: BlockHash9    public var compiledInstructions: [MessageCompiledInstruction]10    public let addressTableLookups: [MessageAddressTableLookup]1112    public init(13        header: MessageHeader,14        staticAccountKeys: [PublicKey],15        recentBlockhash: BlockHash,16        compiledInstructions: [MessageCompiledInstruction],17        addressTableLookups: [MessageAddressTableLookup]18    ) {19        self.header = header20        self.staticAccountKeys = staticAccountKeys21        self.recentBlockhash = recentBlockhash22        self.compiledInstructions = compiledInstructions23        self.addressTableLookups = addressTableLookups24    }2526    public var numAccountKeysFromLookups: Int {27        var count = 028        for lookup in addressTableLookups {29            count += lookup.readonlyIndexes.count + lookup.writableIndexes.count30        }31        return count32    }3334    public func getAccountKeys(accountKeysFromLookups: AccountKeysFromLookups) -> MessageAccountKeys {35        .init(36            staticAccountKeys: staticAccountKeys,37            accountKeysFromLookups: accountKeysFromLookups38        )39    }4041    public func getAccountKeys(addressLookupTableAccounts: [AddressLookupTableAccount]) throws -> MessageAccountKeys {42        let accountKeysFromLookups = try resolveAddressTableLookups(43            addressLookupTableAccounts: addressLookupTableAccounts44        )4546        return .init(47            staticAccountKeys: staticAccountKeys,48            accountKeysFromLookups: accountKeysFromLookups49        )50    }5152    public func isAccountSigner(index: Int) -> Bool {53        index < header.numRequiredSignatures54    }5556    public func isAccountWritable(index: Int) -> Bool {57        let numSignedAccounts = header.numRequiredSignatures58        let numStaticAccountKeys = staticAccountKeys.count5960        if index >= numStaticAccountKeys {61            let lookupAccountKeysIndex = index - numStaticAccountKeys62            let numWritableLookupAccountKeys = addressTableLookups.reduce(0) { count, lookup in63                count + lookup.writableIndexes.count64            }65            return lookupAccountKeysIndex < numWritableLookupAccountKeys66        } else if index >= header.numRequiredSignatures {67            let unsignedAccountIndex = index - numSignedAccounts68            let numUnsignedAccounts = numStaticAccountKeys - numSignedAccounts69            let numWritableUnsignedAccounts = numUnsignedAccounts - header.numReadonlyUnsignedAccounts70            return unsignedAccountIndex < numWritableUnsignedAccounts71        } else {72            let numWritableSignedAccounts = numSignedAccounts - header.numReadonlySignedAccounts73            return index < numWritableSignedAccounts74        }75    }7677    public func resolveAddressTableLookups(78        addressLookupTableAccounts: [AddressLookupTableAccount]79    ) throws -> AccountKeysFromLookups {80        var accountKeysFromLookups = AccountKeysFromLookups(readonly: [], writable: [])8182        for tableLookup in addressTableLookups {83            let tableAccount = addressLookupTableAccounts.first { account in84                account.key == tableLookup.accountKey85            }8687            guard let tableAccount = tableAccount else {88                throw VersionedMessageError89                    .other(90                        "Failed to find address lookup table account for table key \(tableLookup.accountKey.base58EncodedString)"91                    )92            }9394            for index in tableLookup.writableIndexes {95                if index < tableAccount.state.addresses.count {96                    accountKeysFromLookups.writable.append(97                        tableAccount.state.addresses[Int(index)]98                    )99                } else {100                    throw VersionedMessageError101                        .other(102                            "Failed to find address for index \(index) in address lookup table \(tableLookup.accountKey.base58EncodedString)"103                        )104                }105            }106107            for index in tableLookup.readonlyIndexes {108                if index < tableAccount.state.addresses.count {109                    accountKeysFromLookups.readonly.append(110                        tableAccount.state.addresses[Int(index)]111                    )112                } else {113                    throw VersionedMessageError114                        .other(115                            "Failed to find address for index \(index) in address lookup table \(tableLookup.accountKey.base58EncodedString)"116                        )117                }118            }119        }120121        return accountKeysFromLookups122    }123124    public func serialize() throws -> Data {125        let messageVersion0Prefix: UInt8 = 1 << 7126127        var message = Data()128129        // Header130        message.append(Data(messageVersion0Prefix.bytes))131        message.append(serializeHeader())132133        // Accounts134        message.append(Data.encodeLength(staticAccountKeys.count))135        message.append(staticAccountKeys.map(\.data).reduce(Data(), +))136137        // Blockhash138        message.append(serializeRecentBlockhash())139140        // Instructions141        message.append(Data.encodeLength(compiledInstructions.count))142        message.append(serializeInstructions())143//        message.append(Data.encodeLength(0))144//        message.append(serializeInstructions())145146        // Table look uolana147        message.append(Data.encodeLength(addressTableLookups.count))148        try message.append(serializeAddressTableLookups())149150        return message151    }152153    internal func serializeHeader() -> Data {154        Data(header.bytes)155    }156157    internal func serializeInstructions() -> Data {158//        var serializedInstructions: Data = .init()159//160//        for instruction in compiledInstructions {161//            let encodedAccountKeyIndexesLength = Data.encodeLength(instruction.accountKeyIndexes.count)162//            let encodedDataLength = Data.encodeLength(instruction.data.count)163//164//            serializedInstructions.append(Data(instruction.programIdIndex.bytes))165//            serializedInstructions.append(encodedAccountKeyIndexesLength)166//            serializedInstructions.append(Data(Array(instruction.accountKeyIndexes.reduce([], +))))167//            serializedInstructions.append(encodedDataLength)168//            serializedInstructions.append(Data(instruction.data))169//        }170171//        let instructionsLength = Data.encodeLength(compiledInstructions.count)172        Data(compiledInstructions.map(\.serializedData).reduce([], +))173    }174175    internal func serializeAddressTableLookups() throws -> Data {176        var serializedAddressTableLookups = Data()177178        for lookup in addressTableLookups {179            serializedAddressTableLookups.append(lookup.accountKey.data)180181            serializedAddressTableLookups.append(Data.encodeLength(lookup.writableIndexes.count))182            try serializedAddressTableLookups.append(lookup.writableIndexes.rawSerialized())183184            serializedAddressTableLookups.append(Data.encodeLength(lookup.readonlyIndexes.count))185            try serializedAddressTableLookups.append(lookup.readonlyIndexes.rawSerialized())186        }187188        return serializedAddressTableLookups189    }190191    internal func serializeRecentBlockhash() -> Data {192        Data(Base58.decode(recentBlockhash))193    }194195    public static func compile(196        payerKey: PublicKey,197        instructions: [TransactionInstruction],198        recentBlockHash: BlockHash,199        addressLookupTableAccounts: [AddressLookupTableAccount]?200    ) throws -> Self {201        var compiledKeys = CompiledKeys.compile(instructions: instructions, payer: payerKey)202203        var addressTableLookups: [MessageAddressTableLookup] = []204        var accountKeysFromLookups: AccountKeysFromLookups = .init(readonly: [], writable: [])205206        let lookupTableAccounts = addressLookupTableAccounts ?? []207208        for lookupTable in lookupTableAccounts {209            if let extractResult = try compiledKeys.extractTableLookup(lookupTable: lookupTable) {210                addressTableLookups.append(extractResult.0)211                accountKeysFromLookups.writable.append(contentsOf: extractResult.1.writable)212                accountKeysFromLookups.readonly.append(contentsOf: extractResult.1.readonly)213            }214        }215216        let (header, staticAccountKeys) = compiledKeys.getMessageComponents()217        let accountKeys = MessageAccountKeys(218            staticAccountKeys: staticAccountKeys,219            accountKeysFromLookups: accountKeysFromLookups220        )221        let compiledInstructions = try accountKeys.compileInstructions(instructions: instructions)222223        return MessageV0(224            header: header,225            staticAccountKeys: staticAccountKeys,226            recentBlockhash: recentBlockHash,227            compiledInstructions: compiledInstructions,228            addressTableLookups: addressTableLookups229        )230    }231232    public static func deserialize(serializedMessage: Data) throws -> Self {233        var byteArray = BinaryReader(bytes: serializedMessage.bytes)234235        let prefix: UInt8 = try byteArray.read()236        let maskedPrefix: UInt8 = prefix & Constants.versionPrefixMask237        guard prefix != maskedPrefix else {238            throw VersionedMessageError.expectedVersionedMessageButReceivedLegacyMessage239        }240241        let version = maskedPrefix242        guard version == 0 else {243            throw VersionedMessageError.invalidMessageVersion(expectedVersion: 0, receivedVersion: version)244        }245246        let header = try MessageHeader(247            numRequiredSignatures: Int(byteArray.read()),248            numReadonlySignedAccounts: Int(byteArray.read()),249            numReadonlyUnsignedAccounts: Int(byteArray.read())250        )251252        var staticAccountKeys: [PublicKey] = []253        let staticAccountKeysLength = try byteArray.decodeLength()254        for _ in 0 ..< staticAccountKeysLength {255            let account: [UInt8] = try byteArray.read(count: PublicKey.numberOfBytes)256            try staticAccountKeys.append(PublicKey(string: Base58.encode(account)))257        }258259        let recentBlockhash = try Base58.encode(byteArray.read(count: PublicKey.numberOfBytes))260261        let instructionCount = try byteArray.decodeLength()262        var compiledInstructions: [MessageCompiledInstruction] = []263        for _ in 0 ..< instructionCount {264            let programIdIndex = try byteArray.read()265            let accountKeyIndexesLength = try byteArray.decodeLength()266            let accountKeyIndexes = try byteArray.read(count: accountKeyIndexesLength)267            let dataLength = try byteArray.decodeLength()268            let data = try byteArray.read(count: dataLength)269            compiledInstructions.append(270                .init(271                    programIdIndex: programIdIndex,272                    accountKeyIndexes: accountKeyIndexes,273                    data: data274                )275            )276        }277278        let addressTableLookupsCount = try byteArray.decodeLength()279        var addressTableLookups: [MessageAddressTableLookup] = []280        for _ in 0 ..< addressTableLookupsCount {281            let accountKey = try PublicKey(string: Base58.encode(byteArray.read(count: PublicKey.numberOfBytes)))282            let writableIndexesLength = try byteArray.decodeLength()283            let writableIndexes = try byteArray.read(count: writableIndexesLength)284            let readonlyIndexesLength = try byteArray.decodeLength()285            let readonlyIndexes = try byteArray.read(count: readonlyIndexesLength)286            addressTableLookups.append(287                .init(288                    accountKey: accountKey,289                    writableIndexes: writableIndexes,290                    readonlyIndexes: readonlyIndexes291                )292            )293        }294295        return .init(296            header: header,297            staticAccountKeys: staticAccountKeys,298            recentBlockhash: recentBlockhash,299            compiledInstructions: compiledInstructions,300            addressTableLookups: addressTableLookups301        )302    }303}304