@layerzerolabs/lz-sui-sdk-v2/src/utils/package-allowlist-validator.ts

import { MoveCallWithObjectDetails } from './move-call-object-fetcher'
import { IPTBValidator } from './ptb-validator'
import { PackageWhitelistValidator as PackageWhitelistSDK } from '../modules/ptb-builders/package-whitelist-validator'

/**
 * PackageAllowlistValidator ensures all function calls and objects belong to allowlisted packages
 * This validator checks that:
 * 1. All function calls target approved packages
 * 2. All object arguments come from approved packages
 *
 * If local packageWhitelist is empty, it will use the PTB whitelist contract for validation
 * If local packageWhitelist is provided, it will use only the local whitelist (contract is ignored)
 */
export class PackageAllowlistValidator implements IPTBValidator {
    private packageWhitelist: Set<string>
    private validator?: PackageWhitelistSDK

    constructor(packageWhitelist: string[] = [], validator?: PackageWhitelistSDK) {
        if (packageWhitelist.length == 0 && validator == null) {
            throw new Error('PackageAllowlistValidator requires either a packageWhitelist or a validator')
        }
        this.packageWhitelist = new Set(packageWhitelist)
        this.validator = validator
    }

    /**
     * Validates that all function calls and object arguments belong to whitelisted packages
     * @param moveCallsWithDetails - Array of move calls with object details
     * @throws Error if any function call or object belongs to a non-whitelisted package
     */
    async validate(moveCallsWithDetails: MoveCallWithObjectDetails[]): Promise<void> {
        const allPackages = new Set<string>()

        // Collect packages from function calls
        for (const { moveCall } of moveCallsWithDetails) {
            const packageId = moveCall.function.package
            if (packageId === '') {
                throw new Error('Move call package is missing')
            }
            allPackages.add(packageId)
        }

        // Collect packages from objects
        for (const { objectDetails } of moveCallsWithDetails) {
            for (const [objectId, objectResponse] of objectDetails) {
                if (!objectResponse.data) {
                    throw new Error(`Object ${objectId} not found`)
                }

                const packageId = this.extractPackageIdFromType(objectResponse.data.type)
                if (packageId == null) {
                    throw new Error(`Could not extract package ID from object ${objectId}`)
                }
                allPackages.add(packageId)
            }
        }

        // Validate all packages at once
        const isAllValid = await this.validatePackages(Array.from(allPackages))
        // Check result - if any package failed, throw error
        if (!isAllValid) {
            throw new Error(
                `One or more packages are not allowlisted. Packages: [${Array.from(allPackages).join(', ')}]`
            )
        }
    }

    /**
     * Validate multiple packages at once
     * @param packages - Array of package IDs to validate
     * @returns Object mapping package ID to validation result
     */
    private async validatePackages(packages: string[]): Promise<boolean> {
        if (packages.length === 0) {
            return true
        }
        // If custom whitelist is provided, use it first
        if (this.packageWhitelist.size > 0) {
            // Check if ALL packages are in whitelist - if any fails, return false
            for (const packageId of packages) {
                if (!this.packageWhitelist.has(packageId)) {
                    return false
                }
            }
            return true
        } else {
            // Use the contract's validate function to validate all packages at once
            try {
                return await this.validator!.validate(packages)
            } catch (error) {
                console.warn(`Failed to check whitelist:`, error)
                return false
            }
        }
    }

    /**
     * Extracts package ID from object type string
     * Object types are in format: "packageId::module::struct<generics>"
     * @param objectType - The object type string
     * @returns The package ID or null if extraction fails
     */
    private extractPackageIdFromType(objectType: string | null | undefined): string | null {
        if (objectType == null) {
            return null
        }

        // Object type format: "0x123abc::module_name::struct_name<T1, T2>"
        // We want to extract the package ID (0x123abc)
        const parts = objectType.split('::')
        if (parts.length < 2) {
            return null
        }

        const packageId = parts[0]

        // Validate that it looks like a package ID (starts with 0x and has hex characters)
        if (!packageId.startsWith('0x') || packageId.length < 3) {
            return null
        }

        return packageId
    }

    /**
     * Get available packages for error messages
     */
    private getAvailablePackagesList(): string {
        if (this.packageWhitelist.size > 0) {
            return Array.from(this.packageWhitelist).join(', ')
        }
        return 'contract whitelist'
    }
}