NV12toBGRandResize: Fix potential buffer overflow in file output functions

- Increased filename buffer sizes from 120 to 256 characters
- Replaced sprintf() with snprintf() to prevent potential buffer overflows
This commit is contained in:
Rob Armstrong 2025-02-12 11:41:33 -08:00
parent dcce6e1f14
commit 8b2b51e2a5

View File

@ -38,14 +38,17 @@
#include "utils.h" #include "utils.h"
__global__ void floatToChar(float *src, unsigned char *dst, int height, __global__ void floatToChar(float *src, unsigned char *dst, int height,
int width, int batchSize) { int width, int batchSize)
{
int x = threadIdx.x + blockIdx.x * blockDim.x; int x = threadIdx.x + blockIdx.x * blockDim.x;
if (x >= height * width) return; if (x >= height * width)
return;
int offset = height * width * 3; int offset = height * width * 3;
for (int j = 0; j < batchSize; j++) { for (int j = 0; j < batchSize; j++)
{
// b // b
*(dst + j * offset + x * 3 + 0) = *(dst + j * offset + x * 3 + 0) =
(unsigned char)*(src + j * offset + height * width * 0 + x); (unsigned char)*(src + j * offset + height * width * 0 + x);
@ -59,13 +62,15 @@ __global__ void floatToChar(float *src, unsigned char *dst, int height,
} }
void floatPlanarToChar(float *src, unsigned char *dst, int height, int width, void floatPlanarToChar(float *src, unsigned char *dst, int height, int width,
int batchSize) { int batchSize)
{
floatToChar<<<(height * width - 1) / 1024 + 1, 1024, 0, NULL>>>( floatToChar<<<(height * width - 1) / 1024 + 1, 1024, 0, NULL>>>(
src, dst, height, width, batchSize); src, dst, height, width, batchSize);
} }
void dumpRawBGR(float *d_srcBGR, int pitch, int width, int height, void dumpRawBGR(float *d_srcBGR, int pitch, int width, int height,
int batchSize, char *folder, char *tag) { int batchSize, char *folder, char *tag)
{
float *bgr, *d_bgr; float *bgr, *d_bgr;
int frameSize; int frameSize;
char directory[120]; char directory[120];
@ -82,22 +87,25 @@ void dumpRawBGR(float *d_srcBGR, int pitch, int width, int height,
frameSize = width * height * 3 * sizeof(float); frameSize = width * height * 3 * sizeof(float);
bgr = (float *)malloc(frameSize); bgr = (float *)malloc(frameSize);
if (bgr == NULL) { if (bgr == NULL)
{
std::cerr << "Failed malloc for bgr\n"; std::cerr << "Failed malloc for bgr\n";
return; return;
} }
d_bgr = d_srcBGR; d_bgr = d_srcBGR;
for (int i = 0; i < batchSize; i++) { for (int i = 0; i < batchSize; i++)
char filename[120]; {
char filename[256];
std::ofstream *outputFile; std::ofstream *outputFile;
checkCudaErrors(cudaMemcpy((void *)bgr, (void *)d_bgr, frameSize, checkCudaErrors(cudaMemcpy((void *)bgr, (void *)d_bgr, frameSize,
cudaMemcpyDeviceToHost)); cudaMemcpyDeviceToHost));
sprintf(filename, "%s/%s_%d.raw", directory, tag, (i + 1)); snprintf(filename, sizeof(filename), "%s/%s_%d.raw", directory, tag, (i + 1));
outputFile = new std::ofstream(filename); outputFile = new std::ofstream(filename);
if (outputFile) { if (outputFile)
{
outputFile->write((char *)bgr, frameSize); outputFile->write((char *)bgr, frameSize);
delete outputFile; delete outputFile;
} }
@ -109,14 +117,16 @@ void dumpRawBGR(float *d_srcBGR, int pitch, int width, int height,
} }
void dumpBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize, void dumpBGR(float *d_srcBGR, int pitch, int width, int height, int batchSize,
char *folder, char *tag) { char *folder, char *tag)
{
dumpRawBGR(d_srcBGR, pitch, width, height, batchSize, folder, tag); dumpRawBGR(d_srcBGR, pitch, width, height, batchSize, folder, tag);
} }
void dumpYUV(unsigned char *d_nv12, int size, char *folder, char *tag) { void dumpYUV(unsigned char *d_nv12, int size, char *folder, char *tag)
{
unsigned char *nv12Data; unsigned char *nv12Data;
std::ofstream *nv12File; std::ofstream *nv12File;
char filename[120]; char filename[256];
char directory[60]; char directory[60];
char mkdir_cmd[256]; char mkdir_cmd[256];
#if !defined(_WIN32) #if !defined(_WIN32)
@ -129,16 +139,18 @@ void dumpYUV(unsigned char *d_nv12, int size, char *folder, char *tag) {
int ret = system(mkdir_cmd); int ret = system(mkdir_cmd);
sprintf(filename, "%s/%s.nv12", directory, tag); snprintf(filename, sizeof(filename), "%s/%s.nv12", directory, tag);
nv12File = new std::ofstream(filename); nv12File = new std::ofstream(filename);
if (nv12File == NULL) { if (nv12File == NULL)
{
std::cerr << "Failed to new " << filename; std::cerr << "Failed to new " << filename;
return; return;
} }
nv12Data = (unsigned char *)malloc(size * (sizeof(char))); nv12Data = (unsigned char *)malloc(size * (sizeof(char)));
if (nv12Data == NULL) { if (nv12Data == NULL)
{
std::cerr << "Failed to allcoate memory\n"; std::cerr << "Failed to allcoate memory\n";
return; return;
} }