diff --git a/source/serializer/tmfile/tm2_serializer.c b/source/serializer/tmfile/tm2_serializer.c index 3fc87d660..3d295c994 100644 --- a/source/serializer/tmfile/tm2_serializer.c +++ b/source/serializer/tmfile/tm2_serializer.c @@ -154,13 +154,36 @@ static int unregister_tm2_op_loader(struct tm2_serializer* s, int op_type, int o return -1; } +/* return 1 if [off, off+sz) is fully inside the mapped model file, else 0. + All offsets/counts in a tmfile are attacker-controlled; validate before dereferencing. */ +static inline int tm2_off_ok(const struct tm2_priv* priv, size_t off, size_t sz) +{ + if (priv->mem_len < 0) + return 0; + size_t len = (size_t)priv->mem_len; + return off <= len && sz <= len - off; +} + +/* Largest offsets[] element count that still fits inside the mapping after the + TM2_Vector_offsets header at . Expressed as a division so the bound can + never overflow, mirroring tm2_off_ok above. */ +#define TM2_MAX_QUANTPARAMS(priv, off) \ + ((((size_t)(priv)->mem_len) - (size_t)(off) - sizeof(TM2_Vector_offsets)) \ + / sizeof(tm_uoffset_t)) + static int load_graph_tensors(struct tm2_serializer* tm2_s, struct graph* graph, struct tm2_priv* priv) { char* mem_base = (char*)priv->base; const TM2_Subgraph* tm_graph = priv->subgraph; + if (!tm2_off_ok(priv, tm_graph->offset_vo_tensors, sizeof(TM2_Vector_offsets)) || + !tm2_off_ok(priv, tm_graph->offset_vo_buffers, sizeof(TM2_Vector_offsets))) + return -1; const TM2_Vector_offsets* v_tensors = (TM2_Vector_offsets*)(mem_base + tm_graph->offset_vo_tensors); const TM2_Vector_offsets* v_buffers = (TM2_Vector_offsets*)(mem_base + tm_graph->offset_vo_buffers); + if (!tm2_off_ok(priv, tm_graph->offset_vo_tensors, sizeof(TM2_Vector_offsets) + (size_t)v_tensors->v_num * sizeof(v_tensors->offsets[0])) || + !tm2_off_ok(priv, tm_graph->offset_vo_buffers, sizeof(TM2_Vector_offsets) + (size_t)v_buffers->v_num * sizeof(v_buffers->offsets[0]))) + return -1; graph->graph_layout = tm_graph->graph_layout; graph->model_layout = tm_graph->model_layout; @@ -174,6 +197,8 @@ static int load_graph_tensors(struct tm2_serializer* tm2_s, struct graph* graph, for (int i = 0; i < v_tensors->v_num; i++) { + if (!tm2_off_ok(priv, v_tensors->offsets[i], sizeof(TM2_Tensor))) + return -1; const TM2_Tensor* tm_tensor = (TM2_Tensor*)(mem_base + v_tensors->offsets[i]); int flag_permute = 0; // flag the tensor has to be permute int dims_org[8] = {0}; @@ -192,14 +217,22 @@ static int load_graph_tensors(struct tm2_serializer* tm2_s, struct graph* graph, if (tm_tensor->offset_s_tname != TM2_NOT_SET) { // TODO: using update the TM2 model + if (!tm2_off_ok(priv, tm_tensor->offset_s_tname, sizeof(TM2_String))) + return -1; const TM2_String* tm_str = (TM2_String*)(mem_base + tm_tensor->offset_s_tname); + if (!tm2_off_ok(priv, tm_str->offset_data, tm_str->size)) + return -1; ir_tensor->name = strdup_name(mem_base + tm_str->offset_data, tm_str->size); } /* shape */ if (tm_tensor->offset_vd_dims != TM2_NOT_SET) { + if (!tm2_off_ok(priv, tm_tensor->offset_vd_dims, sizeof(TM2_Vector_dims))) + return -1; const TM2_Vector_dims* v_dims = (TM2_Vector_dims*)(mem_base + tm_tensor->offset_vd_dims); + if (!tm2_off_ok(priv, tm_tensor->offset_vd_dims, sizeof(TM2_Vector_dims) + (size_t)v_dims->v_num * sizeof(v_dims->dims[0]))) + return -1; if (tm_graph->model_layout == TENGINE_LAYOUT_NCHW) { @@ -235,6 +268,9 @@ static int load_graph_tensors(struct tm2_serializer* tm2_s, struct graph* graph, /* load const type of tensor, such as the weight or bias for convolution node */ if (ir_tensor->tensor_type == TENSOR_TYPE_CONST) { + if (tm_tensor->buffer_id >= v_buffers->v_num || + !tm2_off_ok(priv, v_buffers->offsets[tm_tensor->buffer_id], sizeof(TM2_Buffer))) + return -1; const TM2_Buffer* tm_buf = (TM2_Buffer*)(mem_base + v_buffers->offsets[tm_tensor->buffer_id]); /* fill temp data buffer to benchmark */ @@ -435,12 +471,23 @@ static int load_graph_tensors(struct tm2_serializer* tm2_s, struct graph* graph, /* load vector type of tensor */ if (tm_tensor->offect_vo_quantparams != TM2_NOT_SET) { + /* the vector header itself must lie inside the mapping before it is read */ + if (!tm2_off_ok(priv, tm_tensor->offect_vo_quantparams, sizeof(TM2_Vector_offsets))) + return -1; + const TM2_Vector_offsets* v_quantparams = (TM2_Vector_offsets*)(mem_base + tm_tensor->offect_vo_quantparams); + /* v_num is file controlled, so the offsets[] array it describes must fit too. + written as a division rather than a multiply so it cannot overflow. */ + if (v_quantparams->v_num > (TM2_MAX_QUANTPARAMS(priv, tm_tensor->offect_vo_quantparams))) + return -1; + /* currently only support one quant param */ ir_tensor->quant_param_num = v_quantparams->v_num; if (v_quantparams->v_num == 1) { + if (!tm2_off_ok(priv, v_quantparams->offsets[0], sizeof(TM2_QuantParam))) + return -1; const TM2_QuantParam* tm_qtparam = (TM2_QuantParam*)(mem_base + v_quantparams->offsets[0]); ir_tensor->scale = tm_qtparam->scale; ir_tensor->zero_point = tm_qtparam->zero_point; @@ -455,6 +502,8 @@ static int load_graph_tensors(struct tm2_serializer* tm2_s, struct graph* graph, for (int j = 0; j < v_quantparams->v_num; j++) { + if (!tm2_off_ok(priv, v_quantparams->offsets[j], sizeof(TM2_QuantParam))) + return -1; const TM2_QuantParam* tm_qtparam = (TM2_QuantParam*)(mem_base + v_quantparams->offsets[j]); ir_tensor->scale_list[j] = tm_qtparam->scale; ir_tensor->zp_list[j] = tm_qtparam->zero_point;