search

openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators
Apache License 2.0
2.7k stars 432 forks source link

Integer overflow in ZlibOutputBuffer causes crash #11668

Open jtotzid opened 6 months ago

jtotzid commented 6 months ago

In https://github.com/openxla/xla/blob/main/third_party/tsl/tsl/lib/io/zlib_outputbuffer.cc#L158 we are truncating a 64-bit size to 32-bit and comparing it to a threshold.

When the size of input data is too big it will wrap around and the check will pass: the truncated size is smaller than the available space (when it actually is not). This will then cause a crash at https://github.com/openxla/xla/blob/main/third_party/tsl/tsl/lib/io/zlib_outputbuffer.cc#L77.

May I suggest instead the integer cast to be:

if (bytes_to_write <= static_cast<size_t>(AvailableInputSpace())) {

i.e. you are casting to a larger, not a smaller, type.

cheshire commented 6 months ago

Thanks ! Could you provide a test case?

akuegel commented 6 months ago

I have fixed it in the way you suggested (https://github.com/openxla/xla/commit/2abc3fd2aab39733f2777af5d3c9e1efdd4d7ed5). Please let me know whether it works now as expected.

jtotzid commented 6 months ago

thanks for the quick fix! trying to build from source but encountering some hiccups...