revanced-bots/packages/shared/src/schemas/Packet.ts

import {
    array,
    boolean,
    custom,
    enum_,
    type InferOutput,
    null_,
    object,
    optional,
    parse,
    pipe,
    string,
    url,
    // merge
} from 'valibot'
import DisconnectReason from '../constants/DisconnectReason'
import { ClientOperation, Operation, ServerOperation } from '../constants/Operation'

/**
 * Schema to validate packets
 */
export const PacketSchema = custom<Packet>(input => {
    if (
        typeof input === 'object' &&
        input &&
        'op' in input &&
        typeof input.op === 'number' &&
        input.op in Operation &&
        'd' in input
    ) {
        if (input.op in ServerOperation && !('s' in input && typeof input.s === 'number')) return false

        try {
            parse(PacketDataSchemas[input.op as Operation], input.d)
            return true
        } catch {
            return false
        }
    }
    return false
}, 'Invalid packet data')

/**
 * Schema to validate packet data for each possible operations
 */
export const PacketDataSchemas = {
    [ServerOperation.Hello]: null_(),
    [ServerOperation.ParsedText]: object({
        labels: array(
            object({
                name: string(),
                confidence: custom<number>(input => typeof input === 'number' && input >= 0 && input <= 1),
            }),
        ),
    }),
    [ServerOperation.ParsedImage]: object({
        text: string(),
    }),
    [ServerOperation.ParseTextFailed]: null_(),
    [ServerOperation.ParseImageFailed]: null_(),
    [ServerOperation.Disconnect]: object({
        reason: enum_(DisconnectReason),
    }),
    [ServerOperation.TrainedMessage]: boolean(),
    [ServerOperation.TrainMessageFailed]: null_(),

    [ClientOperation.ParseText]: object({
        text: string(),
    }),
    [ClientOperation.ParseImage]: object({
        image_url: pipe(string(), url()),
    }),
    [ClientOperation.TrainMessage]: object({
        text: string(),
        label: optional(string()),
    }),
} as const

export type Packet<TOp extends Operation = Operation> = TOp extends ServerOperation
    ? PacketWithSequenceNumber<TOp>
    : Omit<PacketWithSequenceNumber<TOp>, 's'>

type PacketWithSequenceNumber<TOp extends Operation> = {
    op: TOp
    d: InferOutput<(typeof PacketDataSchemas)[TOp]>
    s: number
}