#a. tác dụng

rpc_tensor:

1
struct rpc_tensor {
2
  // u: unsigned (0-> +) khi dùng -1 (tức 0xffffff...) sẽ rất lớn
3
    uint64_t id;
4
    uint32_t type;
5
    uint64_t buffer;
6
    uint32_t ne[GGML_MAX_DIMS];
7
    uint32_t nb[GGML_MAX_DIMS];
8
    uint32_t op;
9
    int32_t  op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
10
    int32_t  flags;
11
    uint64_t src[GGML_MAX_SRC];
12
    uint64_t view_src;
13
    uint64_t view_offs;
14
    uint64_t data;
15
    char name[GGML_MAX_NAME];
16

17
    char padding[4];
18
};

Struct này đóng vai trò làm metadata cho tensor như các kiểu dữ liệu uint64_t cho các thông số khai báo buffer, ne, nb,… Đây chính là thứ client gửi đi phía server theo port

rpc_cmd:

1
enum rpc_cmd {
2
    RPC_CMD_ALLOC_BUFFER = 0,
3
    RPC_CMD_GET_ALIGNMENT,
4
    RPC_CMD_GET_MAX_SIZE,
5
    RPC_CMD_BUFFER_GET_BASE,
6
    RPC_CMD_FREE_BUFFER,
7
    RPC_CMD_BUFFER_CLEAR,
8
    RPC_CMD_SET_TENSOR,
9
    RPC_CMD_SET_TENSOR_HASH,
10
    RPC_CMD_GET_TENSOR,
11
    RPC_CMD_COPY_TENSOR,
12
    RPC_CMD_GRAPH_COMPUTE,
13
    RPC_CMD_GET_DEVICE_MEMORY,
14
    RPC_CMD_INIT_TENSOR,
15
    RPC_CMD_GET_ALLOC_SIZE,
16
    RPC_CMD_HELLO,
17
    RPC_CMD_DEVICE_COUNT,
18
    RPC_CMD_GRAPH_RECOMPUTE,
19
    RPC_CMD_COUNT,
20
};

Khai báo sẵn các command để server và client hiểu nhau, các lệnh sẽ được config ở sau. Có một số lệnh khá quan trọng như RPC_CMD_ALLOC_BUFFER yêu cầu server cấp phát RAM/VRAM,…

deserialize_tensor:

1
ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor) {
2
// Kiểu dữ liệu trả về của hàm là ggml_tensor (hay tensor) | thuộc layer rpc_server, tên hàm deserialize_tensor
3
// ctx (context): con trỏ dùng để cấp phát RAM cho tensor được tạo ra
4
// tensor: con trỏ gói dữ liệu thô nhận được từ client qua mạng, const (chỉ đọc không sửa - đảm bảo tính toàn vẹn)
5

6
    // Kiểm tra xem kiểu dữ liệu từ client gửi lên có nằm trong danh sách được hỗ trợ (GGML_TYPE_COUNT) không, không thì trả về nulll
7
    if (tensor->type >= GGML_TYPE_COUNT) {
8
        GGML_LOG_ERROR("[%s] invalid tensor type received: %u\n", __func__, tensor->type);
9
        return nullptr;
10
    }
11
    // dấu -> có tác dụng truy cập vào một thành phần cụ thể bên trong một struct. Ở đây truy cập vào type (kiểu dữ liệu) trong
12
    // tensor và kiểm tra nếu nó khác GGML_TYPE_COUNT
13

14

15
    // mỗi kiểu dữ liệu AI có một blocksize, nếu nó = 0 thì máy tính sẽ bị lỗi chia hết cho 0 nên dòng này ngăn lỗi này xảy ra
16
    if (ggml_blck_size((enum ggml_type)tensor->type) == 0) {
17
        GGML_LOG_ERROR("[%s] invalid tensor type received (blck_size is 0): %u\n", __func__, tensor->type);
18
        return nullptr;
19
    // __func__: lấy tên hàm hiện tại (deserialize_tensor)
20
    // blocksize có thể bằng 0 khi kiểu dữ liệu sai lệch (do cơ chế tính toán số block trong tensor)
21
    // nên khi chia cho 0 sẽ lỗi nên đoạn này dùng điều kiện nếu cái ggml_type của cái type trong tensor = 0 thì in lỗi rồi
22
    // trả về null
23
    }

Hàm này có tác dụng khởi tạo một tensor cấu trúc ggml_tensor (dạng ma trận) dựa trên thông tin nháp được gửi từ mạng cấu trúc rpc_tensor (dạng byte)

rpc_server: là một class(thường có tác dụng gom các thứ liên quan lại một chỗ cho dễ nhận diện) các hàm xử lý từng loại lệnh từ client như alloc_buffer, set_tensor, graph_compute,…

Các hàm được tạo trong class rpc_server sử dụng deserialize_tensor set_tensor: Đổ dữ liệu từ client vào RAM server set_tensor_hash: Check lỗi init_tensor: Khởi tạo thông số get_tensor: Cho phép server gửi dữ liệu từ RAM -> client copy_tensor: Bắt server copy dữ liệu từ vùng này sang vùng khác create_node: Tạo liên kết các tensor để tìm tensor cha mẹ của nó rồi từ đó xây nên graph tính toán (logic)

Các hàm này đều sử dụng hàm deserialize trong cấu trúc của nó, hầu hết đều có những chức năng khác nhau nhằm tạo ra cấu trúc ggml_tensor toàn vẹn

Nhưng tại sao init_tensor lại không check dữ liệu ngặt như những hàm khác? Vì init_tensor có tác dụng có sự ảnh hưởng yếu tới cấu trúc thật sự, chủ yếu chức năng là thiết lập các thông số cơ bản để khởi tạo một tensor mới nên nó không ảnh hưởng gì đến buffer. Nhưng khi không check result kẻ xấu có thể lợi dụng kẽ hở này để tạo các thông số yêu cầu tạo tensor lớn khiến crash RAM server

Luồng hoạt động: Giai đoạn: Thiết lập kết nối (Handshake)

1. Client: Gọi hàm get_socket() để tạo kết nối TCP đến IP
   của Server.
2. Client: Gửi lệnh RPC_CMD_HELLO (ID: 14).
3. Server: Nhận lệnh trong hàm rpc_serve_client, kiểm tra nếu đúng là HELLO thì gọi server.hello().
4. Server: Phản hồi lại cấu trúc rpc_msg_hello_rsp chứa thông tin phiên bản (Major, Minor, Patch).
5. Client: Kiểm tra phiên bản. Nếu khớp, kết nối chính thức được xác lập và lưu vào buft_map (danh sách các bộ đệm).

Giai đoạn: Workflow

Đây là luồng dữ liệu khi Client thực hiện một phép tính AI thực tế trên Server

1: Cấp phát Buffer

• Client: Gọi ggml_backend_buft_alloc_buffer
• Giao thức: Gửi RPC_CMD_ALLOC_BUFFER kèm kích thước yêu cầu
• Server: Hàm alloc_buffer gọi backend nội bộ (như CUDA hoặc CPU) để lấy RAM, sau đó trả về một remote_ptr (địa chỉ RAM trên Server)

2: Lập danh sách Tensor

• Client: Trước khi tính toán, nó biến các ggml_tensor (C++) thành rpc_tensor (Byte thô) thông qua hàm serialize_tensor.
• Lưu ý: Lúc này rpc_tensor.buffer sẽ chứa cái remote_ptr đã nhận ở Bước 1

3: Đổ dữ liệu vào RAM Server (Set Tensor)

• Client: Gọi ggml_backend_rpc_buffer_set_tensor
• Giao thức: Gửi gói tin gồm: [rpc_tensor metadata] + [offset] + [Dữ liệu byte thô].
• Server:
  1. Nhận gói tin trong set_tensor
  2. Gọi deserialize_tensor để dựng lại khung Tensor
  3. Check bảo mật: Kiểm tra xem data có nằm trong buffer không
  4. Thực thi: Gọi ggml_backend_tensor_set để chép dữ liệu vào RAM

4: Graph Building

• Client: Gọi ggml_backend_rpc_graph_compute
• Server:
  1. Nhận danh sách các Tensor ID.
  2. Gọi create_node liên tục. Hàm này sẽ dùng deserialize_tensor để biến các ID thành các vật thể Tensor thực trong RAM Server và kết nối chúng lại (Node A nối với Node B)

5: Kích hoạt tính toán (Compute)

• Server: Sau khi Graph đã dựng xong, nó gọi ggml_backend_graph_compute (của backend thật như CUDA). Server sẽ “suy nghĩ” dựa trên dữ liệu đã nhận

6: Lấy kết quả về (Get Tensor)

• Client: Gửi RPC_CMD_GET_TENSOR
• Server: Tìm Tensor tương ứng, đọc dữ liệu từ RAM Server và gửi ngược lại cho Client qua socket.

Chú ý:

• #pragma pack(push, 1): Toàn bộ struct RPC được nén chặt không có khoảng trắng thừa. Điều này giúp tính toán Offset cực kỳ chính xác khi viết script Python
• tensor_map & tensor_ptrs: Server dùng std::unordered_map để quản lý các Tensor. Nếu bạn gửi sai ID, Server sẽ không tìm thấy và trả về nullptr (Dẫn đến logic bug).
• deserialize_tensor là trái tim: Mọi hàm (set, get, copy, graph) đều phải đi qua nơi này để biến Byte từ mạng thành vật thể trong RAM.

Tóm lại: Luồng hoạt động là một chuỗi: Kết nối -> Thuê RAM -> Gửi Metadata -> Đổ dữ liệu -> Ra lệnh tính -> Lấy kết quả

#b. Hiểu đoạn code gây bug

1
struct ggml_tensor * result = deserialize_tensor(ctx, tensor);
2
if (result == nullptr) {
3
  return nullptr;
4
}

1
struct ggml_tensor * result = deserialize_tensor(ctx, tensor);
2
// tạo tensor rỗng
3
if (result == nullptr || result->buffer == nullptr) {
4
// trước đó chỉ có dòng nếu result = 0 thì về null
5
// nhưng giờ thêm điều kiện là nếu cái buffer của result = 0 thì có lỗi và về null
6
  GGML_LOG_ERROR("[%s] invalid tensor: null %s (id=%" PRIu64 ")\n",
7
           __func__, result == nullptr ? "tensor" : "buffer", id);
8
  return nullptr;
9
}

#Arbitrary Write

Gửi một gói tin vào server sao cho khi server xử lý đến hàm set_tensor , có thể thấy adr mà mình muốn Dưới đây là chi tiết: Nhìn vào đoạn code này trong file bug

1
if (result->buffer) {
2

3
        // require that the tensor data does not go beyond the buffer end
4
        uint64_t tensor_size = (uint64_t) ggml_nbytes(result);
5
        uint64_t buffer_start = (uint64_t) ggml_backend_buffer_get_base(result->buffer);
6
        uint64_t buffer_size = (uint64_t) ggml_backend_buffer_get_size(result->buffer);
7
        GGML_ASSERT(tensor->data + tensor_size >= tensor->data); // check for overflow
8
        GGML_ASSERT(tensor->data >= buffer_start && tensor->data + tensor_size <= buffer_start + buffer_size);
9
}
10

11
    result->op = (ggml_op) tensor->op;
12
    for (uint32_t i = 0; i < GGML_MAX_OP_PARAMS / sizeof(int32_t); i++) {
13
        result->op_params[i] = tensor->op_params[i];
14
    }
15
    result->flags = tensor->flags;
16
    result->data = reinterpret_cast<void *>(tensor->data);
17
    ggml_set_name(result, tensor->name);
18
    return result;
19

20
}

Bài này có logic là nếu result->buffer: con trỏ result trỏ vào một cái buffer tồn tại (buffer > 0) thì nó sẽ thực hiện các lệnh trong đoạn điều kiện

Lấy size, address bắt đầu cái buffer, kích thước buffer rồi dùng GGML_ASSERT để kiểm tra:

• Kiểm tra xem tensor->data + tensor_size >= tensor->data thì cho qua: tức là bạn có thể hiểu trong bộ nhớ máy tính địa chỉ là một vòng tròn (từ 0x00 -> 0xff… giống số nhà bạn khi đi một vòng trái đất lại về nhà bạn), địa chỉ cũng vậy nên dùng cái này để tránh nó sử dụng một địa chỉ tensor->data pass điều kiện một cái bên dưới nhưng dùng một kích thước rất lớn để trỏ về các vùng địa chỉ thấp hơn mà không vướng phải điều kiện check

• Kiểm tra nếu cái tensor->data (address trỏ đến data trong metadata) >= buffer_start (address start của buffer) và tensor->data + tensor_size <= buffer_start + buffer_size thì cho qua: Tức là cái đoạn đầu để tránh cái địa chỉ khai báo nằm thấp hơn cái địa chỉ buffer, còn cái sau là địa chỉ trỏ đến data trong metadata + kích thước tensor phải nhỏ hơn buffer size

Theo đó, khi mình set được cái buffer = 0 thì cái điều kiện kia nó không xảy ra, code sẽ chạy tiếp xuống các dòng khác Đặc biệt có một lệnh:

1
result->data = reinterpret_cast<void *>(tensor->data);

nó sẽ set data trong ggml_tensor = data trong rpc_tensor. nếu ta khai báo cái tensor->data là một vùng địa chỉ ngoài buffer (bypass condition checking) nó sẽ đặt địa chỉ ta trỏ vào làm vùng bắt đầu dữ liệu của tensor

#Cần làm tiếp theo

1. Tìm cách dùng leak information trong lab này
2. Cần hiểu code rõ hơn, các đoạn code thực thi các lệnh để có thể áp dụng leak information
3. Học cách viết script exploiting, biết lý do tại sao lại dùng các lệnh để xử lí logic