From 44e94a355dc6233f81c1f40d0146f2f27adfdb41 Mon Sep 17 00:00:00 2001 From: Filip Jankovic Date: Thu, 27 Aug 2026 06:50:59 -0400 Subject: [PATCH] Fix two OOM error-path defects in the CPU PNG encoder Both are only reachable when an allocation fails, and both make that failure worse than it needs to be. 1. `torch_png_write_data` assigned the result of `realloc` straight back to `p->buffer`. `realloc` does not free the original block when it returns NULL, so this discarded the only pointer to it. The `png_error` -> `longjmp` cleanup path then sees `buf_info.buffer == nullptr` and skips the `free`, leaking whatever had been encoded so far (up to roughly the size of the finished PNG). The leak does not grow without bound -- each leaked block reduces what the next encode can obtain, so failures happen progressively earlier -- but it permanently destroys the process's ability to encode. Under a fixed memory cap, encode capacity collapsed to zero within three failed calls and never recovered; with the temporary pointer it holds steady indefinitely instead. `realloc(NULL, n)` is equivalent to `malloc(n)`, so the first-call branch is no longer needed. 2. `png_create_write_struct` and `png_create_info_struct` were not checked for NULL. Every subsequent libpng call NULL-guards its `png_ptr` and returns early, so on allocation failure `encode_png` ran to completion and returned an empty tensor: a silent, zero-byte "PNG" with no error raised. `decode_png` already handles the mirror case; this brings the encoder in line with it. Co-Authored-By: Claude Opus 5 (1M context) --- torchvision/csrc/io/image/cpu/encode_png.cpp | 25 +++++++++++++------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/torchvision/csrc/io/image/cpu/encode_png.cpp b/torchvision/csrc/io/image/cpu/encode_png.cpp index a3d5a8e7eb4..d403218733f 100644 --- a/torchvision/csrc/io/image/cpu/encode_png.cpp +++ b/torchvision/csrc/io/image/cpu/encode_png.cpp @@ -55,16 +55,19 @@ void torch_png_write_data( (struct torch_mem_encode*)png_get_io_ptr(png_ptr); size_t nsize = p->size + length; - /* allocate or grow buffer */ - if (p->buffer) { - p->buffer = (char*)realloc(p->buffer, nsize); - } else { - p->buffer = (char*)malloc(nsize); - } - - if (!p->buffer) { + /* Allocate or grow the buffer. The result has to land in a temporary: on + * failure realloc() returns NULL without freeing the original block, so + * assigning it straight back to p->buffer would discard the only pointer to + * that block and leak it. The error path below frees p->buffer, but only sees + * it if we leave it intact here. realloc(NULL, n) is equivalent to malloc(n), + * so the first call needs no special case. */ + char* nbuf = (char*)realloc(p->buffer, nsize); + + if (!nbuf) { + /* p->buffer still owns the original block and is freed by the caller. */ png_error(png_ptr, "Write Error"); } + p->buffer = nbuf; /* copy new bytes to end of buffer */ memcpy(p->buffer + p->size, data, length); @@ -145,8 +148,14 @@ torch::stable::Tensor encode_png( // Initialize PNG structures png_write = png_create_write_struct( PNG_LIBPNG_VER_STRING, &err_ptr, torch_png_error, nullptr); + STD_TORCH_CHECK(png_write, "libpng write structure allocation failed!"); info_ptr = png_create_info_struct(png_write); + if (!info_ptr) { + png_destroy_write_struct(&png_write, nullptr); + // Seems redundant with the if statement. done here to avoid leaking memory. + STD_TORCH_CHECK(info_ptr, "libpng info structure allocation failed!"); + } // Define custom buffer output png_set_write_fn(png_write, &buf_info, torch_png_write_data, nullptr);