Skip to content

Commit 0cd927d

Browse files
PierreMarieCurieroot
authored andcommitted
#4448 - Fail early if engine file path is unwritable before ONNX parsing
Signed-off-by: Paul Pacheco <154653205+PierreMarieCurie@users.noreply.github.com>
1 parent 40efe7e commit 0cd927d

File tree

3 files changed

+37
-0
lines changed

3 files changed

+37
-0
lines changed

samples/common/sampleOptions.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1321,6 +1321,10 @@ void BuildOptions::parse(Arguments& arguments)
13211321

13221322
if (getAndDelOption(arguments, "--saveEngine", engine))
13231323
{
1324+
if (!canWriteFile(engine))
1325+
{
1326+
throw std::invalid_argument(std::string("Cannot write engine file to path: ") + engine);
1327+
}
13241328
save = true;
13251329
}
13261330
if (load && save)

samples/common/sampleUtils.cpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,36 @@ void loadFromFile(std::string const& fileName, char* dst, size_t size)
116116
}
117117
}
118118

119+
// Check if the file at the given path can be written to.
120+
bool canWriteFile(const std::string& path)
121+
{
122+
// Verify that the target directory exists
123+
namespace fs = std::filesystem;
124+
fs::path p(path);
125+
fs::path dir = p.has_parent_path() ? p.parent_path() : fs::current_path();
126+
if (!fs::exists(dir) || !fs::is_directory(dir))
127+
{
128+
return false;
129+
}
130+
131+
// Try creating and writing to a temporary file in the directory
132+
const fs::path tempFilePath = dir / ".writetest.tmp";
133+
std::ofstream test(tempFilePath.string(), std::ios::out | std::ios::trunc);
134+
if (!test.is_open())
135+
{
136+
return false;
137+
}
138+
test << "test";
139+
const bool ok = test.good();
140+
test.close();
141+
142+
// Clean up the temporary file without throwing on failure
143+
std::error_code ec;
144+
fs::remove(tempFilePath, ec);
145+
146+
return ok;
147+
}
148+
119149
std::vector<std::string> splitToStringVec(std::string const& s, char separator, int64_t maxSplit)
120150
{
121151
std::vector<std::string> splitted;

samples/common/sampleUtils.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#define TRT_SAMPLE_UTILS_H
2020

2121
#include <fstream>
22+
#include <filesystem>
2223
#include <iostream>
2324
#include <memory>
2425
#include <numeric>
@@ -74,6 +75,8 @@ void dumpBuffer(void const* buffer, std::string const& separator, std::ostream&
7475

7576
void loadFromFile(std::string const& fileName, char* dst, size_t size);
7677

78+
bool canWriteFile(const std::string& path);
79+
7780
std::vector<std::string> splitToStringVec(std::string const& option, char separator, int64_t maxSplit = -1);
7881

7982
bool broadcastIOFormats(std::vector<IOFormat> const& formats, size_t nbBindings, bool isInput = true);

0 commit comments

Comments
 (0)