trezor-suite/packages/protobuf/scripts/protobuf-definitions.ts

/* eslint-disable no-console */
import fs from 'fs';
import path from 'path';
import * as protobuf from 'protobufjs';

// protobuf.ReflectionObject to JSON
type Definition = {
    reserved?: unknown[];
    options?: Record<string, unknown>;
    valuesOptions?: Record<string, unknown>;
    values?: Record<string, unknown>;
    rule?: string;
    type?: string;
    extend?: string;
    nested?: Record<string, Definition>;
    fields?: Record<string, Definition>;
};

const modifyDefinitionsJSON = (root: protobuf.Root, def: Definition) => {
    // remove "reserved" fields
    delete def.reserved;
    // remove "valuesOptions" fields
    delete def.valuesOptions;
    // remove unused "options"
    if (def.options) {
        const ignoreOptions = [
            '(experimental_field)',
            '(experimental_message)',
            '(has_bitcoin_only_values)',
            'deprecated',
        ];
        const options = Object.keys(def.options);
        // compatibility with json generated by `node_modules/.bin/pbjs`
        // `packed` option is not set for custom types, it is set only for primitives like uint32
        if (def.type && options.includes('packed')) {
            const obj = root.lookup(def.type);
            if (obj) {
                try {
                    root.lookupType(def.type);
                    ignoreOptions.push('packed');
                } catch {
                    /* empty */
                }
            }
        }
        const opts = options
            .filter(opt => !ignoreOptions.includes(opt))
            .reduce((prev, curr) => {
                prev[curr] = def.options?.[curr];

                return prev;
            }, {});

        if (Object.keys(opts).length < 1) {
            delete def.options;
        }
    }

    // replace types pointing to different packages like "hw.trezor.messages.common"
    if (def.type && def.type.includes('.')) {
        def.type = def.type.split('.').pop();
    }

    // modify recursively for nested types and fields
    const { nested } = def;
    if (nested) {
        Object.keys(nested).forEach(key => {
            const item = nested[key];
            const isInternal =
                item.options && Object.keys(item.options).includes('(internal_only)');
            const isExtendingGoogle = item.extend && item.extend.startsWith('google');
            if (isInternal || isExtendingGoogle) {
                delete nested[key];
            } else {
                modifyDefinitionsJSON(root, item);
            }
        });
    }

    if (def.fields) {
        Object.values(def.fields).forEach(item => modifyDefinitionsJSON(root, item));
    }

    return def;
};

type BuildOptions = {
    skipPackages?: string[];
    onlyPackages?: string[];
    includeImports?: boolean;
    messageType?: string; // enum name, default MessageType
};

const modifyMessageType = (proto: protobuf.Root, name: string) => {
    const messageTypeEnum = proto.lookupEnum(name);
    if (messageTypeEnum) {
        Object.keys(messageTypeEnum.values).forEach(key => {
            // replace key `MessageType_Initialize` > `Initialize`
            const newKey = key.replace(name + '_', '');
            const value = messageTypeEnum.values[key];
            // remove old key
            messageTypeEnum.remove(key);
            // check if MessageType is needed, package could be excluded
            if (proto.lookup(newKey)) {
                // add new key
                messageTypeEnum.add(newKey, value);
            }
        });
    }

    if (name !== 'MessageType') {
        // rename custom enum to be MessageType
        const { parent } = messageTypeEnum;
        parent?.remove(messageTypeEnum);

        messageTypeEnum.name = 'MessageType';
        parent?.add(messageTypeEnum);
    }
};

export const buildDefinitions = (protoDir: string, args: BuildOptions) => {
    // https://github.com/protobufjs/protobuf.js/blob/master/README.md#compatibility
    // Because the internals of this package do not rely on google/protobuf/descriptor.proto, options are parsed and presented literally.
    const root = new protobuf.Root({
        common: protobuf.common('descriptor', {}),
    });

    const files: string[] = [];
    const packages: string[] = [];
    const { skipPackages, onlyPackages, includeImports } = args;

    // read all messages*.proto files from directory
    fs.readdirSync(protoDir).forEach(fileName => {
        if (!/^messages.*.proto$/.test(fileName)) {
            return;
        }
        // messages.proto file => empty pkg
        const pkg = fileName.replace(/messages-?(.+)?.proto$/, '$1').replace('-', '_');
        if (skipPackages?.includes(pkg)) {
            return console.log('Skipping', pkg);
        }
        if (onlyPackages && !onlyPackages.includes(pkg)) {
            return console.log('Skipping', pkg);
        }

        if (pkg) {
            packages.push(pkg);
        }
        files.push(path.join(protoDir, fileName));
    });

    console.log('Loading files:', files);

    const proto = root.loadSync(files, { keepCase: true });
    const messages = proto.lookup('hw.trezor.messages');
    if (!messages) {
        throw new Error('hw.trezor.messages not found');
    }

    modifyMessageType(proto, args.messageType || 'MessageType');

    const result = {};
    // remove deep nesting (hw.trezor.messages.*)
    packages.forEach(p => {
        const pkg = proto.lookup(`hw.trezor.messages.${p}`);
        if (!pkg) {
            throw new Error(`hw.trezor.messages.${p} not found`);
        }
        const json = pkg.toJSON();
        Object.assign(result, json.nested);
    });

    // @ts-expect-error typed as protobuf.Reflection but in fact it is a protobuf.Namespace
    const topLevelMessages = includeImports ? messages.nested : {};
    // hw.trezor.messages Namespace contains all the packages, ignore already processed
    Object.keys(topLevelMessages).forEach(name => {
        if (!packages.includes(name)) {
            Object.assign(result, { [name]: topLevelMessages[name].toJSON() });
        }
    });

    return modifyDefinitionsJSON(proto, { nested: result });
};

if (require.main === module) {
    // called directly, otherwise required as a module
    const [protoDir, ...args] = process.argv.slice(2);

    // get --arg=X
    const getArgValue = (args2: string[], arg: string) =>
        args2.find(a => a.startsWith(arg))?.substring(arg.length + 1);

    const json = buildDefinitions(protoDir, {
        includeImports: true,
        skipPackages: getArgValue(args, '--skip')?.split(','),
        onlyPackages: getArgValue(args, '--only')?.split(','),
    });

    const distDir = path.join(__dirname, '../');
    fs.writeFile(`${distDir}/messages.json`, JSON.stringify(json, null, 2), err => {
        if (err) throw err;
    });
}