Fix two OOM error-path defects in the CPU PNG encoder

[fjankovi]

Two defects in encode_png, both reachable only when an allocation fails, and both of which make that failure worse than it needs to be. They share a precondition, so they're fixed together.

1. realloc result assigned back to the pointer it reallocates

if (p->buffer) {
  p->buffer = (char*)realloc(p->buffer, nsize);   // NULL on failure
} else {
  p->buffer = (char*)malloc(nsize);
}
if (!p->buffer) {
  png_error(png_ptr, "Write Error");
}

realloc does not free the original block when it returns NULL, but the only pointer to that block has just been overwritten. The png_error → longjmp cleanup path then reads buf_info.buffer == nullptr and skips its free, so everything encoded so far is leaked — up to roughly the size of the finished PNG.

The compiler makes this unconditional; there's no window in which the old pointer survives:

call   realloc@plt
mov    %rax,(%rbx)      # p->buffer = result, NULL included
test   %rax,%rax        # only now is it checked

The leak does not grow without bound — each leaked block reduces what the next encode can obtain, so subsequent failures occur progressively earlier and the series converges. What it does instead is permanently destroy the process's ability to encode. Measured under a fixed RLIMIT_AS cap, encoding 1024×1024 RGB noise in a loop:

iteration	before	after
0	2,772,993 B encoded before failure	2,772,993 B
1	41 B	2,092,061 B
2	41 B	2,092,061 B
3–40	0 B — never recovers	2,092,061 B — steady

So the practical effect is that one failed encode turns a recoverable, transient OOM into a permanent one. With the fix, capacity holds flat indefinitely and every failure is fully recoverable.

Since realloc(NULL, n) is equivalent to malloc(n) (C99 7.22.3.5), the first-call branch is no longer needed and the whole thing collapses to a temporary plus a commit-on-success.

2. png_create_write_struct / png_create_info_struct not checked for NULL

Every subsequent libpng call NULL-guards its png_ptr and returns early, so when the create call fails, encode_png runs all the way to completion and returns an empty tensor — a silent, zero-byte "PNG" with no error raised. Reproduced by interposing a NULL-returning png_create_write_struct.

decode_png already handles the mirror case (decode_png.cpp:40-48); this brings the encoder in line with it, reusing the same message wording.

Validation

There's no way to force an allocation failure from the Python test suite, so this isn't covered by a new unit test. It was validated against a faithful reduction of the function built on real libpng 1.6.43, with a deterministic allocation-failure shim:

• Before: LeakSanitizer reports Direct leak of 98489 byte(s) in 1 object(s) in torch_png_write_data. After: clean.
• Leak magnitude before the fix ranges from 8 B to 777,809 B on a 512×512 RGB image whose PNG is 777,813 B — i.e. bounded by the encoded size.
• Post-fix suite, all clean under ASan + UBSan + LSan and warning-free under -Wall -Wextra: happy path (output decodes back to the correct dimensions), realloc failure mid-encode, realloc failure on the first call, NULL create (now raises instead of returning 0 bytes), zero-sized dimensions (the reachable longjmp path, still Invalid IHDR data), and invalid arguments.
• clang-format 18.1.3 clean, matching the .pre-commit-config.yaml pin.

The happy path is unchanged — these are error-path-only edits.

🤖 Generated with Claude Code

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9635

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-cla]

Hi @fjankovi!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!