// Isolated reproduction of stable-diffusion.cpp's conditioner load_embedding()
// heap OOB write (src/conditioning/conditioner.hpp:250-253).
//
// The write OFFSET is computed with the CURRENT embedding's dtype size
// (ggml_type_size(embd->type)), but `num_custom_embeddings` counts rows that
// were previously appended using their OWN (possibly smaller) dtype size. Load a
// small-dtype embedding (F16) first, then a large-dtype one (F32), and the offset
// over-counts the earlier rows -> memcpy writes past the just-resized vector.
//
// This file replicates ONLY the vulnerable arithmetic + the two memcpy()s, verbatim,
// to show the out-of-bounds write is real and attacker-controlled. Build at -O0 so the
// stores are not elided (in sd.cpp the buffer is read later, so the writes are live).
#include <vector>
#include <cstring>
#include <cstdint>
#include <cstdio>

// ggml semantics: type_size(F16)=2, type_size(F32)=4; nbytes = ne0*ne1*type_size (contiguous)
static inline size_t type_size(int is_f32) { return is_f32 ? 4 : 2; }

int main() {
    const int64_t hidden_size = 768;          // text_model->model.hidden_size (SD1.x CLIP-L)
    std::vector<uint8_t> token_embed_custom;   // member buffer, persists across load_embedding()
    int32_t num_custom_embeddings = 0;         // member counter, persists

    // ---- load_embedding() #1 : malicious embedding, dtype F16, R0 rows (attacker-controlled ne[1]) ----
    {
        const int is_f32 = 0;                  // embd->type == F16
        const int64_t R0 = 60000;              // rows = tensor_storage.ne[1] straight from the file
        const size_t nbytes = (size_t)hidden_size * R0 * type_size(is_f32); // ggml_nbytes(embd)
        std::vector<uint8_t> embd_data(nbytes, 0xAA);                       // embd->data
        token_embed_custom.resize(token_embed_custom.size() + nbytes);     // conditioner.hpp:250
        std::memcpy((void*)(token_embed_custom.data()
                        + num_custom_embeddings * hidden_size * type_size(is_f32)), // :251 offset (=0 here)
                    embd_data.data(), nbytes);                                       // :252-253
        num_custom_embeddings += (int32_t)R0;  // :254-257 (one bump per row)
        printf("[embd0 F16] buf=%zu bytes, num_custom_embeddings=%d\n",
               token_embed_custom.size(), num_custom_embeddings);
    }

    // ---- load_embedding() #2 : dtype F32, just 1 row ----
    {
        const int is_f32 = 1;                  // embd->type == F32  (larger dtype!)
        const int64_t R1 = 1;
        const size_t nbytes = (size_t)hidden_size * R1 * type_size(is_f32); // ggml_nbytes(embd)
        std::vector<uint8_t> embd_data(nbytes, 0x41);                       // 0x41414141 sentinel
        token_embed_custom.resize(token_embed_custom.size() + nbytes);     // :263 grow by F32 nbytes only

        const size_t off = (size_t)num_custom_embeddings * hidden_size * type_size(is_f32); // :264 BUG
        printf("[embd1 F32] write offset=%zu, len=%zu, buffer_size=%zu  => writes %lld bytes PAST end\n",
               off, nbytes, token_embed_custom.size(),
               (long long)((off + nbytes) - token_embed_custom.size()));
        std::memcpy((void*)(token_embed_custom.data() + off), embd_data.data(), nbytes); // :264-266 OOB WRITE
        num_custom_embeddings += (int32_t)R1;
    }

    printf("SURVIVED (unexpected — no fault)\n");
    return 0;
}