feat(cli):support more format type to publish protobuf

emqx · Jun 27, 2023 · a127fd9 · a127fd9
1 parent b3e7a99
commit a127fd9
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 8 deletions.
diff --git a/cli/src/index.ts b/cli/src/index.ts
@@ -194,6 +194,11 @@ export class Commander {
       )
       .option('-Pp, --protobuf-path <PATH>', 'the .proto file that defines the message format of protobuf')
       .option('-Pmn, --protobuf-message-name <NAME>', 'the name of the protobuf message type')
+      .option(
+        '-Pf, --protobuf-format <TYPE>',
+        'the format type of message body, support base64, json, hex',
+        parseFormat,
+      )
       .action(pub)
 
     this.program

diff --git a/cli/src/lib/pub.ts b/cli/src/lib/pub.ts
@@ -18,17 +18,18 @@ const send = (
     message: string | Buffer
     protobufPath: string | undefined
     protobufMessageName: string | undefined
+    protobufFormat: FormatType | undefined
     opts: IClientPublishOptions
   },
 ) => {
   const client = mqtt.connect(connOpts)
   basicLog.connecting(config, connOpts.hostname!, connOpts.port, pubOpts.topic, pubOpts.message.toString())
   client.on('connect', () => {
     basicLog.connected()
-    const { topic, message, protobufPath, protobufMessageName } = pubOpts
+    const { topic, message, protobufPath, protobufMessageName, protobufFormat } = pubOpts
     basicLog.publishing()
 
-    let bufferMessage = serializeProtobufToBuffer(message, protobufPath, protobufMessageName)
+    let bufferMessage = serializeProtobufToBuffer(message, protobufPath, protobufMessageName, protobufFormat)
     client.publish(topic, bufferMessage, pubOpts.opts, (err) => {
       if (err) {
         signale.warn(err)
@@ -59,6 +60,7 @@ const multisend = (
     message: string | Buffer
     protobufPath: string | undefined
     protobufMessageName: string | undefined
+    protobufFormat: FormatType | undefined
     opts: IClientPublishOptions
   },
   maximumReconnectTimes: number,
@@ -71,9 +73,9 @@ const multisend = (
     objectMode: true,
   })
   sender._write = (line, _enc, cb) => {
-    const { topic, opts, protobufPath, protobufMessageName } = pubOpts
+    const { topic, opts, protobufPath, protobufMessageName, protobufFormat } = pubOpts
 
-    let bufferMessage = serializeProtobufToBuffer(line.trim(), protobufPath, protobufMessageName)
+    let bufferMessage = serializeProtobufToBuffer(line.trim(), protobufPath, protobufMessageName, protobufFormat)
     client.publish(topic, bufferMessage, opts, cb)
   }
 

diff --git a/cli/src/types/global.d.ts b/cli/src/types/global.d.ts
@@ -74,6 +74,7 @@ declare global {
     connUserProperties?: Record<string, string | string[]>
     protobufPath?: string
     protobufMessageName?: string
+    protobufFormat?: FormatType
   }
 
   interface SubscribeOptions extends ConnectOptions {
@@ -97,7 +98,10 @@ declare global {
     interval: number
   }
 
-  type OmitPublishOptions = Omit<PublishOptions, 'stdin' | 'multiline' | 'protobufPath' | 'protobufMessageName'>
+  type OmitPublishOptions = Omit<
+    PublishOptions,
+    'stdin' | 'multiline' | 'protobufPath' | 'protobufMessageName' | 'protobufFormat'
+  >
 
   interface BenchPublishOptions extends OmitPublishOptions {
     count: number

diff --git a/cli/src/utils/convertPayload.ts b/cli/src/utils/convertPayload.ts
@@ -2,7 +2,7 @@ import chalk from 'chalk'
 
 const convertJSON = (value: Buffer) => {
   try {
-    return JSON.parse(value.toString())
+    return JSON.stringify(JSON.parse(value.toString()), null, 2)
   } catch (err) {
     return chalk.red(err)
   }

diff --git a/cli/src/utils/parse.ts b/cli/src/utils/parse.ts
@@ -291,6 +291,7 @@ const parsePublishOptions = (options: PublishOptions) => {
     contentType,
     protobufPath,
     protobufMessageName,
+    protobufFormat,
   } = options
 
   const publishOptions: IClientPublishOptions = {
@@ -316,7 +317,7 @@ const parsePublishOptions = (options: PublishOptions) => {
     )
   }
 
-  return { topic, message, protobufPath, protobufMessageName, opts: publishOptions }
+  return { topic, message, protobufPath, protobufMessageName, protobufFormat, opts: publishOptions }
 }
 
 const parseSubscribeOptions = (options: SubscribeOptions) => {

diff --git a/cli/src/utils/protobuf.ts b/cli/src/utils/protobuf.ts
@@ -1,16 +1,33 @@
 import protobuf from 'protobufjs'
 
+const convertObject = (raw: string | Buffer, format?: FormatType | undefined) => {
+  switch (format) {
+    case 'base64':
+      return JSON.parse(Buffer.from(raw.toString('utf-8'), 'base64').toString('utf-8'))
+    case 'hex':
+      return JSON.parse(Buffer.from(raw.toString('utf-8').replaceAll(' ', ''), 'hex').toString('utf-8'))
+    default:
+      return JSON.parse(raw.toString('utf-8'))
+  }
+}
+
 export const serializeProtobufToBuffer = (
   raw: string | Buffer,
   protobufPath: string | undefined,
   protobufMessageName: string | undefined,
+  format?: FormatType | undefined,
 ): Buffer => {
   let bufferMessage = Buffer.from(raw)
   if (protobufPath && protobufMessageName) {
     try {
       const root = protobuf.loadSync(protobufPath)
       const Message = root.lookupType(protobufMessageName)
-      const data = Message.create(JSON.parse(raw.toString()))
+      const rawData = convertObject(raw, format)
+      const err = Message.verify(rawData)
+      if (err) {
+        throw Error(err)
+      }
+      const data = Message.create(rawData)
       const serializedMessage = Message.encode(data).finish()
       bufferMessage = Buffer.from(serializedMessage)
     } catch (err) {