rmontanana/Pyclassifiers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix Security vulnerabilities and Thread safety

Browse Source
This commit is contained in:
Ricardo Montañana Gómez
2025-07-03 23:23:52 +02:00
parent 91225207f2
commit 707be33097
3 changed files with 460 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

471
pyclfs/PyWrap.cc
View File

@@ -12,7 +12,7 @@ namespace pywrap {
PyWrap* PyWrap::wrapper = nullptr;
std::mutex PyWrap::mutex;
CPyInstance* PyWrap::pyInstance = nullptr;
auto moduleClassMap = std::map<std::pair<std::string, std::string>, std::tuple<PyObject*, PyObject*, PyObject*>>();
// moduleClassMap is now an instance member - removed global declaration
PyWrap* PyWrap::GetInstance()
{
@@ -39,24 +39,48 @@ namespace pywrap {
}
void PyWrap::importClass(const clfId_t id, const std::string& moduleName, const std::string& className)
{
// Validate input parameters for security
validateModuleName(moduleName);
validateClassName(className);
std::lock_guard<std::mutex> lock(mutex);
auto result = moduleClassMap.find(id);
if (result != moduleClassMap.end()) {
return;
}
PyObject* module = PyImport_ImportModule(moduleName.c_str());
if (PyErr_Occurred()) {
errorAbort("Couldn't import module " + moduleName);
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
PyObject* module = PyImport_ImportModule(moduleName.c_str());
if (PyErr_Occurred()) {
PyGILState_Release(gstate);
errorAbort("Couldn't import module " + moduleName);
}
PyObject* classObject = PyObject_GetAttrString(module, className.c_str());
if (PyErr_Occurred()) {
Py_DECREF(module);
PyGILState_Release(gstate);
errorAbort("Couldn't find class " + className);
}
PyObject* instance = PyObject_CallObject(classObject, NULL);
if (PyErr_Occurred()) {
Py_DECREF(module);
Py_DECREF(classObject);
PyGILState_Release(gstate);
errorAbort("Couldn't create instance of class " + className);
}
moduleClassMap.insert({ id, { module, classObject, instance } });
PyGILState_Release(gstate);
}
PyObject* classObject = PyObject_GetAttrString(module, className.c_str());
if (PyErr_Occurred()) {
errorAbort("Couldn't find class " + className);
catch (...) {
PyGILState_Release(gstate);
throw;
}
PyObject* instance = PyObject_CallObject(classObject, NULL);
if (PyErr_Occurred()) {
errorAbort("Couldn't create instance of class " + className);
}
moduleClassMap.insert({ id, { module, classObject, instance } });
}
void PyWrap::clean(const clfId_t id)
{
@@ -82,64 +106,221 @@ namespace pywrap {
}
void PyWrap::errorAbort(const std::string& message)
{
std::cerr << message << std::endl;
PyErr_Print();
RemoveInstance();
exit(1);
// Clear Python error state
if (PyErr_Occurred()) {
PyErr_Clear();
}
// Sanitize error message to prevent information disclosure
std::string sanitizedMessage = sanitizeErrorMessage(message);
// Throw exception instead of terminating process
throw PyWrapException(sanitizedMessage);
}
void PyWrap::validateModuleName(const std::string& moduleName) {
// Whitelist of allowed module names for security
static const std::set<std::string> allowedModules = {
"sklearn.svm", "sklearn.ensemble", "sklearn.tree",
"xgboost", "numpy", "sklearn",
"stree", "odte", "adaboost"
};
if (moduleName.empty()) {
throw PyImportException("Module name cannot be empty");
}
// Check for path traversal attempts
if (moduleName.find("..") != std::string::npos ||
moduleName.find("/") != std::string::npos ||
moduleName.find("\\") != std::string::npos) {
throw PyImportException("Invalid characters in module name: " + moduleName);
}
// Check if module is in whitelist
if (allowedModules.find(moduleName) == allowedModules.end()) {
throw PyImportException("Module not in whitelist: " + moduleName);
}
}
void PyWrap::validateClassName(const std::string& className) {
if (className.empty()) {
throw PyClassException("Class name cannot be empty");
}
// Check for dangerous characters
if (className.find("__") != std::string::npos) {
throw PyClassException("Invalid characters in class name: " + className);
}
// Must be valid Python identifier
if (!std::isalpha(className[0]) && className[0] != '_') {
throw PyClassException("Invalid class name format: " + className);
}
for (char c : className) {
if (!std::isalnum(c) && c != '_') {
throw PyClassException("Invalid character in class name: " + className);
}
}
}
void PyWrap::validateHyperparameters(const json& hyperparameters) {
// Whitelist of allowed hyperparameter keys
static const std::set<std::string> allowedKeys = {
"random_state", "n_estimators", "max_depth", "learning_rate",
"C", "gamma", "kernel", "degree", "coef0", "probability",
"criterion", "splitter", "min_samples_split", "min_samples_leaf",
"min_weight_fraction_leaf", "max_features", "max_leaf_nodes",
"min_impurity_decrease", "bootstrap", "oob_score", "n_jobs",
"verbose", "warm_start", "class_weight"
};
for (const auto& [key, value] : hyperparameters.items()) {
if (allowedKeys.find(key) == allowedKeys.end()) {
throw PyWrapException("Hyperparameter not in whitelist: " + key);
}
// Validate value types and ranges
if (key == "random_state" && value.is_number_integer()) {
int val = value.get<int>();
if (val < 0 || val > 2147483647) {
throw PyWrapException("Invalid random_state value: " + std::to_string(val));
}
}
else if (key == "n_estimators" && value.is_number_integer()) {
int val = value.get<int>();
if (val < 1 || val > 10000) {
throw PyWrapException("Invalid n_estimators value: " + std::to_string(val));
}
}
else if (key == "max_depth" && value.is_number_integer()) {
int val = value.get<int>();
if (val < 1 || val > 1000) {
throw PyWrapException("Invalid max_depth value: " + std::to_string(val));
}
}
}
}
std::string PyWrap::sanitizeErrorMessage(const std::string& message) {
// Remove sensitive information from error messages
std::string sanitized = message;
// Remove file paths
std::regex pathRegex(R"([A-Za-z]:[\\/.][^\s]+|/[^\s]+)");
sanitized = std::regex_replace(sanitized, pathRegex, "[PATH_REMOVED]");
// Remove memory addresses
std::regex addrRegex(R"(0x[0-9a-fA-F]+)");
sanitized = std::regex_replace(sanitized, addrRegex, "[ADDR_REMOVED]");
// Limit message length
if (sanitized.length() > 200) {
sanitized = sanitized.substr(0, 200) + "...";
}
return sanitized;
}
PyObject* PyWrap::getClass(const clfId_t id)
{
std::lock_guard<std::mutex> lock(mutex); // Add thread safety
auto item = moduleClassMap.find(id);
if (item == moduleClassMap.end()) {
errorAbort("Module not found");
throw std::runtime_error("Module not found for id: " + std::to_string(id));
}
return std::get<2>(item->second);
}
std::string PyWrap::callMethodString(const clfId_t id, const std::string& method)
{
PyObject* instance = getClass(id);
PyObject* result;
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
if (!(result = PyObject_CallMethod(instance, method.c_str(), NULL)))
PyObject* instance = getClass(id);
PyObject* result;
if (!(result = PyObject_CallMethod(instance, method.c_str(), NULL))) {
PyGILState_Release(gstate);
errorAbort("Couldn't call method " + method);
}
std::string value = PyUnicode_AsUTF8(result);
Py_XDECREF(result);
PyGILState_Release(gstate);
return value;
}
catch (const std::exception& e) {
PyGILState_Release(gstate);
errorAbort(e.what());
return ""; // This line should never be reached due to errorAbort throwing
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
std::string value = PyUnicode_AsUTF8(result);
Py_XDECREF(result);
return value;
}
int PyWrap::callMethodInt(const clfId_t id, const std::string& method)
{
PyObject* instance = getClass(id);
PyObject* result;
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
if (!(result = PyObject_CallMethod(instance, method.c_str(), NULL)))
PyObject* instance = getClass(id);
PyObject* result;
if (!(result = PyObject_CallMethod(instance, method.c_str(), NULL))) {
PyGILState_Release(gstate);
errorAbort("Couldn't call method " + method);
}
int value = PyLong_AsLong(result);
Py_XDECREF(result);
PyGILState_Release(gstate);
return value;
}
catch (const std::exception& e) {
PyGILState_Release(gstate);
errorAbort(e.what());
return 0; // This line should never be reached due to errorAbort throwing
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
int value = PyLong_AsLong(result);
Py_XDECREF(result);
return value;
}
std::string PyWrap::sklearnVersion()
{
PyObject* sklearnModule = PyImport_ImportModule("sklearn");
if (sklearnModule == nullptr) {
errorAbort("Couldn't import sklearn");
}
PyObject* versionAttr = PyObject_GetAttrString(sklearnModule, "__version__");
if (versionAttr == nullptr || !PyUnicode_Check(versionAttr)) {
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
// Validate module name for security
validateModuleName("sklearn");
PyObject* sklearnModule = PyImport_ImportModule("sklearn");
if (sklearnModule == nullptr) {
PyGILState_Release(gstate);
errorAbort("Couldn't import sklearn");
}
PyObject* versionAttr = PyObject_GetAttrString(sklearnModule, "__version__");
if (versionAttr == nullptr || !PyUnicode_Check(versionAttr)) {
Py_XDECREF(sklearnModule);
PyGILState_Release(gstate);
errorAbort("Couldn't get sklearn version");
}
std::string result = PyUnicode_AsUTF8(versionAttr);
Py_XDECREF(versionAttr);
Py_XDECREF(sklearnModule);
errorAbort("Couldn't get sklearn version");
PyGILState_Release(gstate);
return result;
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
std::string result = PyUnicode_AsUTF8(versionAttr);
Py_XDECREF(versionAttr);
Py_XDECREF(sklearnModule);
return result;
}
std::string PyWrap::version(const clfId_t id)
{
@@ -147,80 +328,128 @@ namespace pywrap {
}
int PyWrap::callMethodSumOfItems(const clfId_t id, const std::string& method)
{
// Call method on each estimator and sum the results (made for RandomForest)
PyObject* instance = getClass(id);
PyObject* estimators = PyObject_GetAttrString(instance, "estimators_");
if (estimators == nullptr) {
errorAbort("Failed to get attribute: " + method);
}
int sumOfItems = 0;
Py_ssize_t len = PyList_Size(estimators);
for (Py_ssize_t i = 0; i < len; i++) {
PyObject* estimator = PyList_GetItem(estimators, i);
PyObject* result;
if (method == "node_count") {
PyObject* owner = PyObject_GetAttrString(estimator, "tree_");
if (owner == nullptr) {
Py_XDECREF(estimators);
errorAbort("Failed to get attribute tree_ for: " + method);
}
result = PyObject_GetAttrString(owner, method.c_str());
if (result == nullptr) {
Py_XDECREF(estimators);
Py_XDECREF(owner);
errorAbort("Failed to get attribute node_count: " + method);
}
Py_DECREF(owner);
} else {
result = PyObject_CallMethod(estimator, method.c_str(), nullptr);
if (result == nullptr) {
Py_XDECREF(estimators);
errorAbort("Failed to call method: " + method);
}
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
// Call method on each estimator and sum the results (made for RandomForest)
PyObject* instance = getClass(id);
PyObject* estimators = PyObject_GetAttrString(instance, "estimators_");
if (estimators == nullptr) {
PyGILState_Release(gstate);
errorAbort("Failed to get attribute: " + method);
}
sumOfItems += PyLong_AsLong(result);
Py_DECREF(result);
int sumOfItems = 0;
Py_ssize_t len = PyList_Size(estimators);
for (Py_ssize_t i = 0; i < len; i++) {
PyObject* estimator = PyList_GetItem(estimators, i);
PyObject* result;
if (method == "node_count") {
PyObject* owner = PyObject_GetAttrString(estimator, "tree_");
if (owner == nullptr) {
Py_XDECREF(estimators);
PyGILState_Release(gstate);
errorAbort("Failed to get attribute tree_ for: " + method);
}
result = PyObject_GetAttrString(owner, method.c_str());
if (result == nullptr) {
Py_XDECREF(estimators);
Py_XDECREF(owner);
PyGILState_Release(gstate);
errorAbort("Failed to get attribute node_count: " + method);
}
Py_DECREF(owner);
} else {
result = PyObject_CallMethod(estimator, method.c_str(), nullptr);
if (result == nullptr) {
Py_XDECREF(estimators);
PyGILState_Release(gstate);
errorAbort("Failed to call method: " + method);
}
}
sumOfItems += PyLong_AsLong(result);
Py_DECREF(result);
}
Py_DECREF(estimators);
PyGILState_Release(gstate);
return sumOfItems;
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
Py_DECREF(estimators);
return sumOfItems;
}
void PyWrap::setHyperparameters(const clfId_t id, const json& hyperparameters)
{
// Set hyperparameters as attributes of the class
PyObject* pValue;
PyObject* instance = getClass(id);
for (const auto& [key, value] : hyperparameters.items()) {
std::stringstream oss;
oss << value.type_name();
if (oss.str() == "string") {
pValue = Py_BuildValue("s", value.get<std::string>().c_str());
} else {
if (value.is_number_integer()) {
pValue = Py_BuildValue("i", value.get<int>());
// Validate hyperparameters for security
validateHyperparameters(hyperparameters);
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
// Set hyperparameters as attributes of the class
PyObject* pValue;
PyObject* instance = getClass(id);
for (const auto& [key, value] : hyperparameters.items()) {
std::stringstream oss;
oss << value.type_name();
if (oss.str() == "string") {
pValue = Py_BuildValue("s", value.get<std::string>().c_str());
} else {
pValue = Py_BuildValue("f", value.get<double>());
if (value.is_number_integer()) {
pValue = Py_BuildValue("i", value.get<int>());
} else {
pValue = Py_BuildValue("f", value.get<double>());
}
}
if (!pValue) {
PyGILState_Release(gstate);
throw PyWrapException("Failed to create Python value for hyperparameter: " + key);
}
int res = PyObject_SetAttrString(instance, key.c_str(), pValue);
if (res == -1 && PyErr_Occurred()) {
Py_XDECREF(pValue);
PyGILState_Release(gstate);
errorAbort("Couldn't set attribute " + key + "=" + value.dump());
}
}
int res = PyObject_SetAttrString(instance, key.c_str(), pValue);
if (res == -1 && PyErr_Occurred()) {
Py_XDECREF(pValue);
errorAbort("Couldn't set attribute " + key + "=" + value.dump());
}
Py_XDECREF(pValue);
PyGILState_Release(gstate);
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
}
void PyWrap::fit(const clfId_t id, CPyObject& X, CPyObject& y)
{
PyObject* instance = getClass(id);
CPyObject result;
CPyObject method = PyUnicode_FromString("fit");
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), y.getObject(), NULL)))
PyObject* instance = getClass(id);
CPyObject result;
CPyObject method = PyUnicode_FromString("fit");
if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), y.getObject(), NULL))) {
PyGILState_Release(gstate);
errorAbort("Couldn't call method fit");
}
PyGILState_Release(gstate);
}
catch (const std::exception& e) {
PyGILState_Release(gstate);
errorAbort(e.what());
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
}
PyObject* PyWrap::predict_proba(const clfId_t id, CPyObject& X)
{
@@ -232,32 +461,60 @@ namespace pywrap {
}
PyObject* PyWrap::predict_method(const std::string name, const clfId_t id, CPyObject& X)
{
PyObject* instance = getClass(id);
PyObject* result;
CPyObject method = PyUnicode_FromString(name.c_str());
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), NULL)))
PyObject* instance = getClass(id);
PyObject* result;
CPyObject method = PyUnicode_FromString(name.c_str());
if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), NULL))) {
PyGILState_Release(gstate);
errorAbort("Couldn't call method " + name);
}
PyGILState_Release(gstate);
// PyObject_CallMethodObjArgs already returns a new reference, no need for Py_INCREF
return result; // Caller must free this object
}
catch (const std::exception& e) {
PyGILState_Release(gstate);
errorAbort(e.what());
return nullptr; // This line should never be reached due to errorAbort throwing
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
// PyObject_CallMethodObjArgs already returns a new reference, no need for Py_INCREF
return result; // Caller must free this object
}
double PyWrap::score(const clfId_t id, CPyObject& X, CPyObject& y)
{
PyObject* instance = getClass(id);
CPyObject result;
CPyObject method = PyUnicode_FromString("score");
// Acquire GIL for Python operations
PyGILState_STATE gstate = PyGILState_Ensure();
try {
if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), y.getObject(), NULL)))
PyObject* instance = getClass(id);
CPyObject result;
CPyObject method = PyUnicode_FromString("score");
if (!(result = PyObject_CallMethodObjArgs(instance, method.getObject(), X.getObject(), y.getObject(), NULL))) {
PyGILState_Release(gstate);
errorAbort("Couldn't call method score");
}
double resultValue = PyFloat_AsDouble(result);
PyGILState_Release(gstate);
return resultValue;
}
catch (const std::exception& e) {
PyGILState_Release(gstate);
errorAbort(e.what());
return 0.0; // This line should never be reached due to errorAbort throwing
}
catch (...) {
PyGILState_Release(gstate);
throw;
}
double resultValue = PyFloat_AsDouble(result);
return resultValue;
}
}

38
pyclfs/PyWrap.h
View File

@@ -4,7 +4,10 @@
#include <map>
#include <tuple>
#include <mutex>
#include <regex>
#include <set>
#include <nlohmann/json.hpp>
#include <stdexcept>
#include "boost/python/detail/wrap_python.hpp"
#include "PyHelper.hpp"
#include "TypeId.h"
@@ -16,6 +19,36 @@ namespace pywrap {
Singleton class to handle Python/numpy interpreter.
*/
using json = nlohmann::json;
// Custom exception classes for PyWrap errors
class PyWrapException : public std::runtime_error {
public:
explicit PyWrapException(const std::string& message) : std::runtime_error(message) {}
};
class PyImportException : public PyWrapException {
public:
explicit PyImportException(const std::string& module)
: PyWrapException("Failed to import Python module: " + module) {}
};
class PyClassException : public PyWrapException {
public:
explicit PyClassException(const std::string& className)
: PyWrapException("Failed to find Python class: " + className) {}
};
class PyInstanceException : public PyWrapException {
public:
explicit PyInstanceException(const std::string& className)
: PyWrapException("Failed to create instance of Python class: " + className) {}
};
class PyMethodException : public PyWrapException {
public:
explicit PyMethodException(const std::string& method)
: PyWrapException("Failed to call Python method: " + method) {}
};
class PyWrap {
public:
PyWrap() = default;
@@ -37,6 +70,11 @@ namespace pywrap {
void importClass(const clfId_t id, const std::string& moduleName, const std::string& className);
PyObject* getClass(const clfId_t id);
private:
// Input validation and security
void validateModuleName(const std::string& moduleName);
void validateClassName(const std::string& className);
void validateHyperparameters(const json& hyperparameters);
std::string sanitizeErrorMessage(const std::string& message);
// Only call RemoveInstance from clean method
static void RemoveInstance();
PyObject* predict_method(const std::string name, const clfId_t id, CPyObject& X);