#Rust Add buffer verification

google · Jan 6, 2025 · e7f5d13 · e7f5d13
1 parent 977da58
commit e7f5d13
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 110 deletions.
diff --git a/rust/reflection/src/lib.rs b/rust/reflection/src/lib.rs
@@ -32,6 +32,8 @@ use num::traits::float::Float;
 use num::traits::int::PrimInt;
 use num::traits::FromPrimitive;
 use stdint::uintmax_t;
+use std::error;
+use stdint::uintmax_t;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
@@ -338,9 +340,14 @@ pub unsafe fn get_any_field_string_in_struct(
 /// # Safety
 ///
 /// [buf] must contain a valid root table and valid offset to it.
-pub unsafe fn set_any_field_integer(buf: &mut [u8], field: &Field, v: i64) -> FlatbufferResult<()> {
+pub unsafe fn set_any_field_integer(
+    buf: &mut [u8],
+    table_loc: usize,
+    field: &Field,
+    v: i64,
+) -> FlatbufferResult<()> {
     let field_type = field.type_().base_type();
-    let table = get_any_root(buf);
+    let table = Table::follow(buf, table_loc);
 
     let Some(field_loc) = get_field_loc(&table, field) else {
         return Err(FlatbufferError::SetValueNotSupported());
@@ -358,9 +365,14 @@ pub unsafe fn set_any_field_integer(buf: &mut [u8], field: &Field, v: i64) -> Fl
 /// # Safety
 ///
 /// [buf] must contain a valid root table and valid offset to it.
-pub unsafe fn set_any_field_float(buf: &mut [u8], field: &Field, v: f64) -> FlatbufferResult<()> {
+pub unsafe fn set_any_field_float(
+    buf: &mut [u8],
+    table_loc: usize,
+    field: &Field,
+    v: f64,
+) -> FlatbufferResult<()> {
     let field_type = field.type_().base_type();
-    let table = get_any_root(buf);
+    let table = Table::follow(buf, table_loc);
 
     let Some(field_loc) = get_field_loc(&table, field) else {
         return Err(FlatbufferError::SetValueNotSupported());
@@ -378,9 +390,14 @@ pub unsafe fn set_any_field_float(buf: &mut [u8], field: &Field, v: f64) -> Flat
 /// # Safety
 ///
 /// [buf] must contain a valid root table and valid offset to it.
-pub unsafe fn set_any_field_string(buf: &mut [u8], field: &Field, v: &str) -> FlatbufferResult<()> {
+pub unsafe fn set_any_field_string(
+    buf: &mut [u8],
+    table_loc: usize,
+    field: &Field,
+    v: &str,
+) -> FlatbufferResult<()> {
     let field_type = field.type_().base_type();
-    let table = get_any_root(buf);
+    let table = Table::follow(buf, table_loc);
 
     let Some(field_loc) = get_field_loc(&table, field) else {
         return Err(FlatbufferError::SetValueNotSupported());
@@ -400,11 +417,12 @@ pub unsafe fn set_any_field_string(buf: &mut [u8], field: &Field, v: &str) -> Fl
 /// [buf] must contain a valid root table and valid offset to it.
 pub unsafe fn set_field<T: EndianScalar>(
     buf: &mut [u8],
+    table_loc: usize,
     field: &Field,
     v: T,
 ) -> FlatbufferResult<()> {
     let field_type = field.type_().base_type();
-    let table = get_any_root(buf);
+    let table = Table::follow(buf, table_loc);
 
     if !is_scalar(field_type) {
         return Err(FlatbufferError::SetValueNotSupported());
@@ -444,6 +462,7 @@ pub unsafe fn set_field<T: EndianScalar>(
 /// [buf] must contain a valid root table and valid offset to it and conform to the [schema].
 pub unsafe fn set_string(
     buf: &mut Vec<u8>,
+    table_loc: usize,
     field: &Field,
     v: &str,
     schema: &Schema,
@@ -460,8 +479,7 @@ pub unsafe fn set_string(
         ));
     }
 
-    let table = get_any_root(buf);
-    let table_loc = table.loc();
+    let table = Table::follow(buf, table_loc);
 
     let Some(field_loc) = get_field_loc(&table, field) else {
         return Err(FlatbufferError::SetValueNotSupported());