Add mechanism for 'safe' memory reads for complex types

[clintropolis]

Description

This PR fixes an accidental kill switch introduced by #13332, where by exposing the ability to construct sketches via expressions we also expose the ability to crash any jvm process running them by feeding them bad input.

For example, the query SELECT COMPLEX_DECODE_BASE64('HLLSketch', 'AgEH') which is a truncated base64 encoded sketch object, will immediately crash the broker which evaluates it when planning the SQL query with an error of the form:

#
# A fatal error has been detected by the Java Runtime Environment:
#
#  SIGSEGV (0xb) at pc=0x000000010424616c, pid=77140, tid=0x0000000000010303
#
# JRE version: OpenJDK Runtime Environment (8.0_332-b08) (build 1.8.0_332-b08)
# Java VM: OpenJDK 64-Bit Server VM (25.332-b08 mixed mode bsd-aarch64 compressed oops)
# Problematic frame:
# j  org.apache.datasketches.hll.PreambleUtil.insertInt(Lorg/apache/datasketches/memory/WritableMemory;JI)V+3
#

After the changes in this PR, we just get a regular out of bounds exception:

2022-11-12T12:48:10,545 WARN [sql[a54c6476-ed24-4e00-a87b-f01c90cdae2c]] org.apache.druid.sql.http.SqlResource - Failed to handle query: SqlQuery{query='SELECT COMPLEX_DECODE_BASE64('HLLSketch', 'AgEH')', resultFormat=array, header=true, typesHeader=true, sqlTypesHeader=true, context={sqlOuterLimit=1001, sqlQueryId=a54c6476-ed24-4e00-a87b-f01c90cdae2c, queryId=a54c6476-ed24-4e00-a87b-f01c90cdae2c}, parameters=[]}
java.lang.IndexOutOfBoundsException: null
	at java.nio.Buffer.checkIndex(Buffer.java:545) ~[?:1.8.0_332]
	at java.nio.HeapByteBuffer.get(HeapByteBuffer.java:142) ~[?:1.8.0_332]
	at org.apache.druid.segment.data.SafeMemoryBase.getByte(SafeMemoryBase.java:65) ~[classes/:?]
	at org.apache.datasketches.hll.PreambleUtil.extractLgK(PreambleUtil.java:280) ~[datasketches-java-3.2.0.jar:?]
	at org.apache.datasketches.hll.HllSketch.wrap(HllSketch.java:247) ~[datasketches-java-3.2.0.jar:?]
	at org.apache.druid.query.aggregation.datasketches.hll.HllSketchObjectStrategy.fromByteBufferSafe(HllSketchObjectStrategy.java:64) ~[?:?]
	at org.apache.druid.query.aggregation.datasketches.hll.HllSketchObjectStrategy.fromByteBufferSafe(HllSketchObjectStrategy.java:31) ~[?:?]
	at org.apache.druid.segment.column.ObjectStrategyComplexTypeStrategy.fromBytes(ObjectStrategyComplexTypeStrategy.java:93) ~[classes/:?]
	at org.apache.druid.math.expr.BuiltInExprMacros$ComplexDecodeBase64ExprMacro$ComplexDecodeBase64Expression.eval(BuiltInExprMacros.java:110) ~[classes/:?]
	at org.apache.druid.math.expr.Parser.lambda$flatten$1(Parser.java:181) ~[classes/:?]
	at org.apache.druid.math.expr.BuiltInExprMacros$ComplexDecodeBase64ExprMacro$ComplexDecodeBase64Expression.visit(BuiltInExprMacros.java:117) ~[classes/:?]
	at org.apache.druid.math.expr.Parser.flatten(Parser.java:152) ~[classes/:?]
	at org.apache.druid.math.expr.Parser.parse(Parser.java:144) ~[classes/:?]
	at org.apache.druid.math.expr.Parser.parse(Parser.java:125) ~[classes/:?]
	at org.apache.druid.sql.calcite.planner.DruidRexExecutor.reduce(DruidRexExecutor.java:74) ~[classes/:?]
	at org.apache.calcite.rel.rules.ReduceExpressionsRule.reduceExpressionsInternal(ReduceExpressionsRule.java:695) ~[calcite-core-1.21.0.jar:1.21.0]

In fairness, this problem existed prior to #13332, since the underlying problem has to do with anything using Memory on untrusted input bytes for anything that gets instructions of where to read or write from those bytes. The reason is that Memory for performance reasons uses Unsafe for all read and write operations, even when wrapping ByteBuffer, so if the code tries to read out of bounds locations it can crash the jvm process and do who knows what else.

I think this behavior is absolutely correct for Memory for most uses, but to solve our problem in a bit of an unconventional way, this PR introduces SafeWritableMemory and SafeWritableBuffer that delegate all operations to an underlying plain java ByteBuffer so that it gets the bounds checking that goes with that, which we can use for anything that uses Memory when loading from an untrusted source. When reading from segments we still use standard Memory calls.

Fortunately due to a bug that wasn't fixed until #13332, i wasn't able to trigger this exact failure in older versions of Druid with the native complex_decode_base64 expression because the complex type mapping wasn't correctly wired up to the function, but i believe it is still possible when ingesting pre-aggregated sketches that use Memory.

To wire up to the complex_decode_base64 function, I've added a new method to ObjectStrategy,

  @Nullable
  default T fromByteBufferSafe(ByteBuffer buffer, int numBytes)
  {
    return fromByteBuffer(buffer, numBytes);
  }

which should be implemented by anything which uses Unsafe or other potentially dangerous operations. I debated having a default implementation but ultimately decided to add one since this interface is marked as an @ExtensionPoint. This is erring on the side of danger in favor of compatibility, so I welcome discussion here.

Besides the added tests, I also did an experiment to replace all calls to Memory.wrap with the "safe" version and all of the tests passed, so I think the "safe" implementation of Memory should be correct.

Release note

A new method has been added to the ObjectStrategy extension point

  @Nullable
  default T fromByteBufferSafe(ByteBuffer buffer, int numBytes)
  {
    return fromByteBuffer(buffer, numBytes);
  }

to allow extension writers that implement complex types which typically use "unsafe" memory operations (such as using Java Unsafe or Datasketches Memory) to optionally provide a "safe" read method to use when processing untrusted inputs such as those decoded from base64 string inputs at ingest time and avoid bad inputs crashing the process. A default implementation is provided to be compatibility focused, but if your extension fits this description and uses unsafe memory operations, you should implement this method. For users of Datasketches Memory, Druid provides a SafeWritableMemory.wrap method which can create a 'safe' Memory from wrapping a ByteBuffer.

---

This PR has:

• been self-reviewed.
• added Javadocs for most classes and all non-trivial methods. Linked related entities via Javadoc links.
• added comments explaining the "why" and the intent of the code wherever would not be obvious for an unfamiliar reader.
• added unit tests or modified existing tests to cover new code paths, ensuring the threshold for code coverage is met.
• been tested in a test Druid cluster.

[rohangarg]

Thanks for the changes @clintropolis! 👍

I had a doubt regarding the usage of SafeWritableMemory in future. Are there any plans to move it to datasketches-memory library as well so that it can be maintained in the long run?
Also, there are assertions in the Memory methods as well but since they aren't enabled in deployments, hence we don't get to use them. Maybe, the Memory object in future could itself expose safe wrapping methods which internally runs those assertions as code validations. WDYT?

Regarding the compatibility, I think since the method has not been released yet (is it released in 24.0.1?) the default implementation could also throw UOE, but it would mean that all complex types would have to implement it. Further, the function could be put behind a cluster level feature flag (not a good thing as well) so that the users/admins are aware of the scope. I was trying to see if this problem could get solved by newer java internals strong encapsulation feature in future but don't have any concrete update/thoughts on it yet.

[clintropolis]

Also, there are assertions in the Memory methods as well but since they aren't enabled in deployments, hence we don't get to use them. Maybe, the Memory object in future could itself expose safe wrapping methods which internally runs those assertions as code validations. WDYT?

This isn't really up to me, but we have offered to contribute if there is wider use beyond Druids needs, relevant slack thread https://the-asf.slack.com/archives/CP0930GKG/p1666290716186309.

Maybe, the Memory object in future could itself expose safe wrapping methods which internally runs those assertions as code validations. WDYT?

I guess I just used a ByteBuffer for this because I think such checks largely would negate the reasons of using unsafe, though I guess it would still allow larger memory objects than a java ByteBuffer allows to be done in a 'safe' manner if instead I was wrapping another Memory object, so maybe that is a better way if we were to merge this functionality in the upstream library since it would be a bit more flexible. I was mostly just being lazy as possible so I just used ByteBuffer since the checking was "free" to me. Is that what you meant?

Regarding the compatibility, I think since the method has not been released yet (is it released in 24.0.1?) the default implementation could also throw UOE, but it would mean that all complex types would have to implement it.

My personal feelings are that this is at least safe as of this PR to have a default implementation with all core and contrib extensions. Extension implementors have to go out of their way to use Unsafe or something like it since they are provided with a ByteBuffer, so I think as long as we advise in the release notes and javadocs that any extensions that do this should implement the new fromByteBufferSafe method then it is ok. But otoh, it wouldn't be very hard to fill in all the implementations with what I made the default if people feel strongly the other way to force implementors to adapt.

Further, the function could be put behind a cluster level feature flag (not a good thing as well) so that the users/admins are aware of the scope.

I'd rather avoid such a flag if possible, the PR is intended to prevent having to have something like that. And as mentioned, it isn't strictly a problem with COMPLEX_DECODE_BASE64, such a flag would also need to apply to all ingest time aggregators that read from byte[] and map it into Memory, which seems complicated. It would definitely be not needed if we did not provide a default implementation of fromByteBufferSafe. If it comes down to it, I would much rather drop the default implementation than add the flag.

[rohangarg]

This isn't really up to me, but we have offered to contribute if there is wider use beyond Druids needs, relevant slack thread https://the-asf.slack.com/archives/CP0930GKG/p1666290716186309.
I guess I just used a ByteBuffer for this because I think such checks largely would negate the reasons of using unsafe, though I guess it would still allow larger memory objects than a java ByteBuffer allows to be done in a 'safe' manner if instead I was wrapping another Memory object, so maybe that is a better way if we were to merge this functionality in the upstream library since it would be a bit more flexible. I was mostly just being lazy as possible so I just used ByteBuffer since the checking was "free" to me. Is that what you meant?

Thanks a lot for the link to the slack thread - it was relevant to the assertion/exception doubt that I had. Regarding the safe implementation, what I meant was to have a function like Memory.safeWrap which only our safe implementations use - the internal query system can keep on using the unsafe objects for better performance.

Regarding the compatibility and safety, I agree with your points that the current modules are safe and we could put advisory in the docs for extension implementors to consider implementing this method and also ensure safety while ingesting intermediate states objects. My reasoning for suggesting flag was to not affect users who aren't using these functionalities but do have extensions - but agree that a default throwing implementation would be a better option than flag.

[rohangarg]

LGTM since all the modules in apache distribution are safe now while deserializing complex objects.
For extensions, it is decided to put a release note with this change which requests the implementors to verify the safety behavior in the extension.