1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.features.kryo.json
10 
11 import com.esotericsoftware.kryo.io.{Input, Output}
12 import com.fasterxml.jackson.databind.JsonNode
13 import com.fasterxml.jackson.databind.node.{ArrayNode, ObjectNode}
14 import com.typesafe.scalalogging.LazyLogging
15 import org.locationtech.geomesa.features.kryo.json.JsonPathParser._
16 
17 import java.nio.charset.StandardCharsets
18 import scala.util.control.NonFatal
19 
20 /**
21  * Serializes into bson (https://bsonspec.org/). Note this is a limited form of bson that only matches
22  * the existing json types - does not cover the bson extensions like binary data, etc. Also note endianness,
23  * etc might not match the spec 100%.
24  *
25  * The bson serialized value is preceded by a single byte, `\x00` to indicate null, or `\x01` to indicate non-null.
26  * Additionally, non-document top-level values are supported, indicated with the prefix byte `\x02`. Top-level,
27  * non-document values are encoded as a jsonb `element` with an empty name.
28  *
29  * Reduced BSON spec - only native JSON elements supported:
30  *
31  * byte    1 byte (8-bits)
32  * int32   4 bytes (32-bit signed integer, two's complement)
33  * int64   8 bytes (64-bit signed integer, two's complement)
34  * double  8 bytes (64-bit IEEE 754-2008 binary floating point)
35  *
36  * document  ::= int32 e_list "\x00" BSON Document. int32 is the total number of bytes comprising the document.
37  * e_list    ::= element e_list
38  *           |	""
39  * element   ::= "\x01" e_name double	64-bit binary floating point
40  *           | "\x02" e_name string	UTF-8 string
41  *           | "\x03" e_name document	Embedded document
42  *           |	"\x04" e_name document	Array
43  *           |	"\x08" e_name "\x00"	Boolean "false"
44  *           |	"\x08" e_name "\x01"	Boolean "true"
45  *           |	"\x09" e_name int64	UTC datetime
46  *           |	"\x0A" e_name	Null value
47  *           |	"\x10" e_name int32	32-bit integer
48  *           |	"\x11" e_name int64	Timestamp
49  *           |	"\x12" e_name int64	64-bit integer
50  * e_name    ::= cstring	Key name
51  * string    ::= int32 (byte*) "\x00"	String - The int32 is the number bytes in the (byte*) + 1 (for the trailing '\x00').
52  *                                              The (byte*) is zero or more UTF-8 encoded characters.
53  * cstring   ::= (byte*) "\x00"	Zero or more modified UTF-8 encoded characters followed by '\x00'. The (byte*)
54  *                               MUST NOT contain '\x00', hence it is not full UTF-8.
55  *
56  * Note:
57  *   Array - The document for an array is a normal BSON document with integer values for the keys,
58  *   starting with 0 and continuing sequentially. For example, the array ['red', 'blue'] would be
59  *   encoded as the document {'0': 'red', '1': 'blue'}. The keys must be in ascending numerical order.
60  */
61 object KryoJsonSerialization extends LazyLogging {
62 
63   private[json] val TerminalByte :Byte = 0x00
64   private[json] val DoubleByte   :Byte = 0x01
65   private[json] val StringByte   :Byte = 0x02
66   private[json] val DocByte      :Byte = 0x03
67   private[json] val ArrayByte    :Byte = 0x04
68   private[json] val BooleanByte  :Byte = 0x08
69   private[json] val NullByte     :Byte = 0x0A
70   private[json] val IntByte      :Byte = 0x10
71   private[json] val LongByte     :Byte = 0x12
72 
73   private[json] val BooleanFalse :Byte = 0x00
74   private[json] val BooleanTrue  :Byte = 0x01
75   private[json] val NonDoc       :Byte = 0x02
76 
77   private val nameBuffers = new ThreadLocal[Array[Byte]] {
78     override def initialValue(): Array[Byte] = Array.ofDim[Byte](32)
79   }
80 
81   /**
82     * Serialize a json object
83     *
84     * @param out output to write to
85     * @param json json string to serialize - must be a json object
86     */
87   def serialize(out: Output, json: String): Unit = {
88     val obj = if (json == null) { null } else {
89       try {
90         mapper.readTree(json)
91       } catch {
92         case NonFatal(e) =>
93           logger.warn(s"Error parsing json:\n$json", e)
94           null
95       }
96     }
97     serialize(out, obj)
98   }
99 
100   /**
101     * Serialize a json object
102     *
103     * @param out output to write to
104     * @param json object to serialize
105     */
106   def serialize(out: Output, json: JsonNode): Unit = {
107     if (json == null || json.isNull) {
108       out.write(BooleanFalse)
109     } else if (json.isObject) {
110       out.write(BooleanTrue); writeDocument(out, json)
111     } else {
112       out.write(NonDoc); writeValue(out, "", json)
113     }
114   }
115 
116   /**
117     * Deserialize the given input. The input should be pointing to the start of
118     * the bytes written by `serialize`. Upon completion, the input will be pointing
119     * to the first byte after the bytes written by `serialize`.
120     *
121     * @param in input, pointing to the start of the json object
122     * @return json as a string
123     */
124   def deserializeAndRender(in: Input): String = {
125     val json = deserialize(in)
126     if (json == null) {
127       null
128     } else {
129       mapper.writeValueAsString(json)
130     }
131   }
132 
133   /**
134     * Deserialize the given input. The input should be pointing to the start of
135     * the bytes written by `serialize`. Upon completion, the input will be pointing
136     * to the first byte after the bytes written by `serialize`.
137     *
138     * @param in input, pointing to the start of the json object
139     * @return parsed json object
140     */
141   def deserialize(in: Input): JsonNode = {
142     try {
143       in.readByte match {
144         case BooleanFalse => null
145         case BooleanTrue  => readDocument(in: Input)
146         case NonDoc       => readValue(in)._2
147       }
148     } catch {
149       case NonFatal(e) => logger.error("Error reading serialized kryo json", e); null
150     }
151   }
152 
153   /**
154     * Deserialize the result of a json-path. The input should be pointing to the start of
155     * the bytes written by `serialize`. There is no guarantee where the input will be
156     * pointing after completion.
157     *
158     * If the path selects leaf elements, they will be returned as primitives. If the path
159     * selects objects, they will be returned as json strings. If more than one item is
160     * selected, they will be returned in a java.util.List. If nothing is selected, it will return null.
161     *
162     * @param in input, pointing to the start of the json object
163     * @param path pre-parsed json path to evaluate
164     * @return result of the path, if any
165     */
166   def deserialize(in: Input, path: JsonPath): Any = {
167     if (path.isEmpty) {
168       deserializeAndRender(in)
169     } else {
170       try { new KryoJsonPath(in).deserialize(path) } catch {
171         case NonFatal(e) => logger.error("Error reading serialized kryo json", e); null
172       }
173     }
174   }
175 
176   // primitive writing functions - in general will write a byte identifying the type, the key and then the value
177 
178   private def writeDocument(out: Output, name: String, value: JsonNode): Unit = {
179     out.writeByte(DocByte)
180     out.writeName(name)
181     writeDocument(out, value)
182   }
183 
184   // write a document without a name - used for the outermost object which doesn't have a key
185   private def writeDocument(out: Output, value: JsonNode): Unit = {
186     val start = out.position()
187     // write a placeholder that we will overwrite when we go back to write total length
188     // note: don't just modify position, as that doesn't expand the buffer correctly
189     out.writeInt(0)
190     value.forEachEntry { case (name, elem) => writeValue(out, name, elem) }
191     out.writeByte(TerminalByte) // marks the end of our object
192     // go back and write the total length
193     val end = out.position()
194     out.setPosition(start)
195     out.writeInt(end - start)
196     out.setPosition(end)
197   }
198 
199   private def writeValue(out: Output, name: String, value: JsonNode): Unit = {
200     if (value.isTextual || value.isBinary) {
201       writeString(out, name, value.asText())
202     } else if (value.isObject) {
203       writeDocument(out, name, value)
204     } else if (value.isArray) {
205       writeArray(out, name, value)
206     } else if (value.isDouble || value.isFloat) {
207       writeDouble(out, name, value.asDouble())
208     } else if (value.isInt || value.isShort) {
209       writeInt(out, name, value.intValue())
210     } else if (value.isLong) {
211       writeLong(out, name, value.longValue())
212     } else if (value.isNull) {
213       writeNull(out, name)
214     } else if (value.isBoolean) {
215       writeBoolean(out, name, value.booleanValue())
216     } else if (value.isBigDecimal) {
217       writeDouble(out, name, value.asDouble())
218     } else if (value.isBigInteger) {
219       if (value.canConvertToInt) {
220         writeInt(out, name, value.intValue())
221       } else if (value.canConvertToLong) {
222         writeLong(out, name, value.longValue())
223       } else {
224         logger.warn(s"Skipping int value that does not fit in a long: $value")
225       }
226     } else {
227       logger.warn(s"Unhandled JsonNode: $value")
228     }
229   }
230 
231   private def writeArray(out: Output, name: String, value: JsonNode): Unit = {
232     out.writeByte(ArrayByte)
233     out.writeName(name)
234     // we store as an object where array index is the key
235     var i = 0
236     val obj = mapper.createObjectNode()
237     while (i < value.size()) {
238       obj.set(i.toString, value.get(i))
239       i += 1
240     }
241     writeDocument(out, obj)
242   }
243 
244   private def writeString(out: Output, name: String, value: String): Unit = {
245     out.writeByte(StringByte)
246     out.writeName(name)
247     val bytes = value.getBytes(StandardCharsets.UTF_8)
248     out.writeInt(bytes.length)
249     out.write(bytes)
250     out.writeByte(TerminalByte)
251   }
252 
253   private def writeDouble(out: Output, name: String, value: Double): Unit = {
254     out.writeByte(DoubleByte)
255     out.writeName(name)
256     out.writeDouble(value)
257   }
258 
259   private def writeInt(out: Output, name: String, value: Int): Unit = {
260     out.writeByte(IntByte)
261     out.writeName(name)
262     out.writeInt(value)
263   }
264 
265   private def writeLong(out: Output, name: String, value: Long): Unit = {
266     out.writeByte(LongByte)
267     out.writeName(name)
268     out.writeLong(value)
269   }
270 
271   private def writeBoolean(out: Output, name: String, value: Boolean): Unit = {
272     out.writeByte(BooleanByte)
273     out.writeName(name)
274     out.writeByte(if (value) BooleanTrue else BooleanFalse)
275   }
276 
277   private def writeNull(out: Output, name: String): Unit = {
278     out.writeByte(NullByte)
279     out.writeName(name)
280   }
281 
282   // primitive reading/skipping methods corresponding to the write methods above
283   // assumes that the indicator byte and name have already been read
284 
285   private[json] def readDocument(in: Input): ObjectNode = {
286     val end = in.position() + in.readInt() - 1 // last byte is the terminal byte
287     val obj = mapper.createObjectNode()
288     while (in.position() < end) {
289       val (k, v) = readValue(in)
290       obj.set(k, v)
291     }
292     in.skip(1) // skip over terminal byte
293     obj
294   }
295 
296   private[json] def readValue(in: Input): (String, JsonNode) = {
297     val switch = in.readByte()
298     val name = in.readName()
299     val value = switch match {
300       case StringByte   => mapper.getNodeFactory.textNode(readString(in))
301       case DocByte      => readDocument(in)
302       case ArrayByte    => readArray(in)
303       case DoubleByte   => mapper.getNodeFactory.numberNode(in.readDouble())
304       case IntByte      => mapper.getNodeFactory.numberNode(in.readInt())
305       case LongByte     => mapper.getNodeFactory.numberNode(in.readLong())
306       case NullByte     => mapper.getNodeFactory.nullNode()
307       case BooleanByte  => mapper.getNodeFactory.booleanNode(readBoolean(in))
308     }
309     (name, value)
310   }
311 
312   private[json] def readArray(in: Input): ArrayNode = {
313     val obj = readDocument(in)
314     val array = mapper.getNodeFactory.arrayNode(obj.size())
315     obj.forEachEntry { case (_, v) => array.add(v) }
316     array
317   }
318 
319   private[json] def skipDocument(in: Input): Unit = in.skip(in.readInt - 4) // length includes bytes storing length
320 
321   private[json] def readString(in: Input): String = {
322     val bytes = Array.ofDim[Byte](in.readInt())
323     in.read(bytes)
324     in.skip(1) // skip TerminalByte
325     new String(bytes, StandardCharsets.UTF_8)
326   }
327 
328   private[json] def skipString(in: Input): Unit = in.skip(in.readInt() + 1) // skip TerminalByte
329 
330   private[json] def readBoolean(in: Input): Boolean = in.readByte == BooleanTrue
331 
332   private[json] def skipBoolean(in: Input): Unit = in.skip(1)
333 
334   private[json] implicit class RichOutput(val out: Output) extends AnyRef {
335     def writeName(name: String): Unit = {
336       // note: names are not allowed to contain the terminal byte (0x00) but we don't check for it
337       out.write(name.getBytes(StandardCharsets.UTF_8))
338       out.writeByte(TerminalByte)
339     }
340   }
341 
342   private[json] implicit class RichInput(val in: Input) extends AnyRef {
343     def readName(): String = {
344       var buffer = nameBuffers.get()
345       var i = 0
346       var byte: Byte = in.readByte()
347       while (byte != TerminalByte) {
348         if (i == buffer.length) {
349           // expand our cached buffer to accommodate the name
350           val copy = Array.ofDim[Byte](buffer.length * 2)
351           System.arraycopy(buffer, 0, copy, 0, i)
352           buffer = copy
353           nameBuffers.set(buffer)
354         }
355         buffer(i) = byte
356         i += 1
357         byte = in.readByte()
358       }
359       new String(buffer, 0, i, StandardCharsets.UTF_8)
360     }
361 
362     def skipName(): Unit = while (in.readByte() != TerminalByte) {}
363   }
364 }