spb/os-vault Public

1import Foundation23public struct PublicKey: Codable, Equatable, CustomStringConvertible, Hashable {4    public static let NULL_PUBLICKEY_BYTES: [UInt8] = Array(repeating: UInt8(0), count: numberOfBytes)5    public static let numberOfBytes = 326    public let bytes: [UInt8]78    public func encode(to encoder: Encoder) throws {9        var container = encoder.singleValueContainer()10        try container.encode(base58EncodedString)11    }1213    public init(from decoder: Decoder) throws {14        let container = try decoder.singleValueContainer()15        let string = try container.decode(String.self)16        try self.init(string: string)17    }1819    public init(string: String?) throws {20        guard let string = string, string.utf8.count >= PublicKey.numberOfBytes21        else {22            throw PublicKeyError.invalidAddress(string)23        }24        let bytes = Base58.decode(string)25        self.bytes = bytes26    }2728    public init(data: Data) throws {29        guard data.count <= PublicKey.numberOfBytes else {30            throw PublicKeyError.invalidAddress(.init(data: data, encoding: .utf8))31        }32        bytes = [UInt8](data)33    }3435    public init(bytes: [UInt8]?) throws {36        guard let bytes = bytes, bytes.count <= PublicKey.numberOfBytes else {37            throw PublicKeyError.invalidAddress(.init(data: Data(bytes ?? []), encoding: .utf8))38        }39        self.bytes = bytes40    }4142    public var base58EncodedString: String {43        Base58.encode(bytes)44    }4546    public var data: Data {47        Data(bytes)48    }4950    public var description: String {51        base58EncodedString52    }5354    public func short(numOfSymbolsRevealed: Int = 4) -> String {55        let pubkey = base58EncodedString56        return pubkey.prefix(numOfSymbolsRevealed) + "..." + pubkey.suffix(numOfSymbolsRevealed)57    }5859    public func hash(into hasher: inout Hasher) {60        hasher.combine(bytes)61    }6263    public static func == (lhs: PublicKey, rhs: PublicKey) -> Bool {64        lhs.bytes == rhs.bytes65    }66}6768extension PublicKey: BytesEncodable {}6970// https://github.com/solana-labs/solana-web3.js/blob/dfb4497745c9fbf01e9633037bf9898dfd5adf94/src/publickey.ts#L2247172// MARK: - Constants7374private var maxSeedLength = 3275private let gf1 = NaclLowLevel.gf([1])7677private extension Int {78    func toBool() -> Bool {79        self != 080    }81}8283public enum PublicKeyError: Error, Equatable {84    case notFound85    case invalidAddress(String?)86    case maxSeedLengthExceeded87    case invalidSeed(reason: String?)88}8990public extension PublicKey {91    static func associatedTokenAddress(92        walletAddress: PublicKey,93        tokenMintAddress: PublicKey,94        tokenProgramId: PublicKey95    ) throws -> PublicKey {96        try findProgramAddress(97            seeds: [98                walletAddress.data,99                tokenProgramId.data,100                tokenMintAddress.data,101            ],102            programId: AssociatedTokenProgram.id103        ).0104    }105106    // MARK: - Helpers107108    static func findProgramAddress(109        seeds: [Data],110        programId: Self111    ) throws -> (Self, UInt8) {112        for nonce in stride(from: UInt8(255), to: 0, by: -1) {113            let seedsWithNonce = seeds + [Data([nonce])]114            do {115                let address = try createProgramAddress(116                    seeds: seedsWithNonce,117                    programId: programId118                )119                return (address, nonce)120            } catch {121                continue122            }123        }124        throw PublicKeyError.notFound125    }126127    static func createProgramAddress(128        seeds: [Data],129        programId: PublicKey130    ) throws -> PublicKey {131        // construct data132        var data = Data()133        for seed in seeds {134            if seed.bytes.count > maxSeedLength {135                throw PublicKeyError.maxSeedLengthExceeded136            }137            data.append(seed)138        }139        data.append(programId.data)140        data.append("ProgramDerivedAddress".data(using: .utf8)!)141142        // hash it143        let hash = data.sha256()144        let publicKeyBytes = Bignum(number: hash.hexString, withBase: 16).data145146        // check it147        if isOnCurve(publicKeyBytes: publicKeyBytes).toBool() {148            throw PublicKeyError.invalidSeed(reason: "address must fall off the curve")149        }150        return try PublicKey(data: publicKeyBytes)151    }152153    static func createWithSeed(154        fromPublicKey: PublicKey,155        seed: String,156        programId: PublicKey157    ) throws -> PublicKey {158        var data = Data()159        data += fromPublicKey.data160        guard let seedData = seed.data(using: .utf8) else {161            throw PublicKeyError.invalidSeed(reason: nil)162        }163        data += seedData164        data += programId.data165        let hash = data.sha256()166        return try PublicKey(data: hash)167    }168169    static func isOnCurve(publicKey: String) -> Int {170        let data = Base58.decode(publicKey)171        return isOnCurve(publicKeyBytes: Data(data))172    }173174    static func isOnCurve(publicKeyBytes: Data) -> Int {175        var r = [[Int64]](repeating: NaclLowLevel.gf(), count: 4)176177        var t = NaclLowLevel.gf(),178            chk = NaclLowLevel.gf(),179            num = NaclLowLevel.gf(),180            den = NaclLowLevel.gf(),181            den2 = NaclLowLevel.gf(),182            den4 = NaclLowLevel.gf(),183            den6 = NaclLowLevel.gf()184185        NaclLowLevel.set25519(&r[2], gf1)186        NaclLowLevel.unpack25519(&r[1], publicKeyBytes.bytes)187        NaclLowLevel.S(&num, r[1])188        NaclLowLevel.M(&den, num, NaclLowLevel.D)189        NaclLowLevel.Z(&num, num, r[2])190        NaclLowLevel.A(&den, r[2], den)191192        NaclLowLevel.S(&den2, den)193        NaclLowLevel.S(&den4, den2)194        NaclLowLevel.M(&den6, den4, den2)195        NaclLowLevel.M(&t, den6, num)196        NaclLowLevel.M(&t, t, den)197198        NaclLowLevel.pow2523(&t, t)199        NaclLowLevel.M(&t, t, num)200        NaclLowLevel.M(&t, t, den)201        NaclLowLevel.M(&t, t, den)202        NaclLowLevel.M(&r[0], t, den)203204        NaclLowLevel.S(&chk, r[0])205        NaclLowLevel.M(&chk, chk, den)206        if NaclLowLevel.neq25519(chk, num).toBool() {207            NaclLowLevel.M(&r[0], r[0], NaclLowLevel.I)208        }209210        NaclLowLevel.S(&chk, r[0])211        NaclLowLevel.M(&chk, chk, den)212213        if NaclLowLevel.neq25519(chk, num).toBool() {214            return 0215        }216        return 1217    }218}219220public extension PublicKey {221    static var sysvarRent: PublicKey { "SysvarRent111111111111111111111111111111111" }222223    static var wrappedSOLMint: PublicKey { "So11111111111111111111111111111111111111112" }224    static var solMint: PublicKey { "Ejmc1UB4EsES5oAaRN63SpoxMJidt3ZGBrqrZk49vjTZ"225    } // Arbitrary mint to represent SOL (not wrapped SOL).226227    static var swapHostFeeAddress: PublicKey { "AHLwq66Cg3CuDJTFtwjPfwjJhifiv6rFwApQNKgX57Yg" }228229    static var renBTCMint: PublicKey { "CDJWUqTcYTVAKXAVXoQZFes5JUFc7owSeq7eMQcDSbo5" }230    static var renBTCMintDevnet: PublicKey { "FsaLodPu4VmSwXGr3gWfwANe4vKf8XSZcCh1CEeJ3jpD" }231    static var fake: PublicKey { "BGcmLttQoYIw4Yfzc7RkZJCKR53IlAybgq8HK0vmovP0\n" }232233    static func orcaSwapId(version: Int = 2) -> PublicKey {234        switch version {235        case 2:236            return "9W959DqEETiGZocYWCQPaJ6sBmUzgfxXfqGeTEdp3aQP"237        default:238            return "DjVE6JNiYqPL2QXyCUUh8rNjHrbz9hXHNYt99MQ59qw1"239        }240    }241242    static var usdcMint: PublicKey { "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v" }243    static var usdtMint: PublicKey { "Es9vMFrzaCERmJfrF4H2FYD4KCoNkY11McCe8BenwNYB" }244    static var dexPID: PublicKey { "9xQeWvG816bUx9EPjHmaT23yvVM2ZWbrrpZb9PusVFin" }245    static var serumSwapPID: PublicKey { "22Y43yTVxuUkoRKdm9thyRhQ3SdgQS7c7kB6UNCiaczD" }246    var isUsdx: Bool {247        self == .usdcMint || self == .usdtMint248    }249}250251extension PublicKey: BorshCodable {252    public func serialize(to writer: inout Data) throws {253        try bytes.forEach { try $0.serialize(to: &writer) }254    }255256    public init(from reader: inout BinaryReader) throws {257        let byteArray = try Array(0 ..< PublicKey.numberOfBytes).map { _ in try UInt8(from: &reader) }258        bytes = byteArray259    }260}261