use crate::binary::ArgType;
use crate::dispatch::SharedArgs;
use crate::DCPError;
#[derive(Debug, Clone)]
pub struct ToolSchema {
pub name: &'static str,
pub id: u16,
pub description: &'static str,
pub input: InputSchema,
}
#[derive(Debug, Clone)]
pub struct InputSchema {
pub required: u64,
pub fields: Vec<FieldDef>,
}
#[derive(Debug, Clone)]
pub struct FieldDef {
pub name: &'static str,
pub field_type: ArgType,
pub offset: u16,
pub size: u16,
pub enum_range: Option<(u8, u8)>,
}
pub struct SchemaValidator;
impl SchemaValidator {
pub fn validate_required(schema: &InputSchema, present_mask: u64) -> Result<(), DCPError> {
let missing = schema.required & !present_mask;
if missing != 0 {
return Err(DCPError::ValidationFailed);
}
Ok(())
}
pub fn validate_enum(field: &FieldDef, value: u8) -> Result<(), DCPError> {
if let Some((min, max)) = field.enum_range {
if value < min || value > max {
return Err(DCPError::ValidationFailed);
}
}
Ok(())
}
pub fn validate_input(
schema: &InputSchema,
present_mask: u64,
field_values: &[(usize, u8)], ) -> Result<(), DCPError> {
Self::validate_required(schema, present_mask)?;
for &(field_idx, value) in field_values {
let field = schema
.fields
.get(field_idx)
.ok_or(DCPError::ValidationFailed)?;
Self::validate_enum(field, value)?;
}
Ok(())
}
pub fn validate_shared_args(schema: &InputSchema, args: &SharedArgs) -> Result<(), DCPError> {
let mut present_mask = 0u64;
let mut declared_data_len = 0usize;
for (idx, field) in schema.fields.iter().enumerate() {
if idx >= 16 {
return Err(DCPError::ValidationFailed);
}
let actual_type = arg_type_from_layout(args.layout(), idx)?;
if actual_type == ArgType::Null {
continue;
}
if actual_type != field.field_type {
return Err(DCPError::ValidationFailed);
}
let offset = field.offset as usize;
let size = field.size as usize;
let end = offset.checked_add(size).ok_or(DCPError::OutOfBounds)?;
if end > args.data().len() {
return Err(DCPError::OutOfBounds);
}
declared_data_len = declared_data_len.max(end);
if field.enum_range.is_some() {
let value = *args.data().get(offset).ok_or(DCPError::OutOfBounds)?;
Self::validate_enum(field, value)?;
}
present_mask |= 1 << idx;
}
for idx in schema.fields.len()..16 {
if arg_type_from_layout(args.layout(), idx)? != ArgType::Null {
return Err(DCPError::ValidationFailed);
}
}
if args.data().len() > declared_data_len {
return Err(DCPError::ValidationFailed);
}
Self::validate_required(schema, present_mask)
}
}
fn arg_type_from_layout(layout: u64, index: usize) -> Result<ArgType, DCPError> {
let shift = index.checked_mul(4).ok_or(DCPError::ValidationFailed)?;
let type_bits = ((layout >> shift) & 0xF) as u8;
ArgType::from_u8(type_bits).ok_or(DCPError::ValidationFailed)
}
impl InputSchema {
pub fn new() -> Self {
Self {
required: 0,
fields: Vec::new(),
}
}
pub fn add_field(&mut self, field: FieldDef) -> &mut Self {
self.fields.push(field);
self
}
pub fn set_required(&mut self, field_index: usize) -> &mut Self {
if field_index < 64 {
self.required |= 1 << field_index;
}
self
}
pub fn is_required(&self, field_index: usize) -> bool {
if field_index >= 64 {
return false;
}
self.required & (1 << field_index) != 0
}
pub fn required_count(&self) -> u32 {
self.required.count_ones()
}
}
impl Default for InputSchema {
fn default() -> Self {
Self::new()
}
}
impl FieldDef {
pub fn new(name: &'static str, field_type: ArgType, offset: u16, size: u16) -> Self {
Self {
name,
field_type,
offset,
size,
enum_range: None,
}
}
pub fn new_enum(name: &'static str, offset: u16, size: u16, min: u8, max: u8) -> Self {
Self {
name,
field_type: ArgType::I32, offset,
size,
enum_range: Some((min, max)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_required_validation() {
let mut schema = InputSchema::new();
schema.set_required(0);
schema.set_required(2);
assert!(SchemaValidator::validate_required(&schema, 0b101).is_ok());
assert!(SchemaValidator::validate_required(&schema, 0b111).is_ok());
assert_eq!(
SchemaValidator::validate_required(&schema, 0b100),
Err(DCPError::ValidationFailed)
);
assert_eq!(
SchemaValidator::validate_required(&schema, 0b001),
Err(DCPError::ValidationFailed)
);
}
#[test]
fn test_enum_validation() {
let field = FieldDef::new_enum("status", 0, 1, 1, 5);
assert!(SchemaValidator::validate_enum(&field, 1).is_ok());
assert!(SchemaValidator::validate_enum(&field, 3).is_ok());
assert!(SchemaValidator::validate_enum(&field, 5).is_ok());
assert_eq!(
SchemaValidator::validate_enum(&field, 0),
Err(DCPError::ValidationFailed)
);
assert_eq!(
SchemaValidator::validate_enum(&field, 6),
Err(DCPError::ValidationFailed)
);
}
#[test]
fn test_complete_validation() {
let mut schema = InputSchema::new();
schema.add_field(FieldDef::new("name", ArgType::String, 0, 32));
schema.add_field(FieldDef::new_enum("type", 32, 1, 1, 3));
schema.set_required(0);
schema.set_required(1);
assert!(SchemaValidator::validate_input(&schema, 0b11, &[(1, 2)]).is_ok());
assert_eq!(
SchemaValidator::validate_input(&schema, 0b01, &[(1, 2)]),
Err(DCPError::ValidationFailed)
);
assert_eq!(
SchemaValidator::validate_input(&schema, 0b11, &[(1, 5)]),
Err(DCPError::ValidationFailed)
);
}
#[test]
fn test_schema_helpers() {
let mut schema = InputSchema::new();
schema.set_required(0);
schema.set_required(3);
assert!(schema.is_required(0));
assert!(!schema.is_required(1));
assert!(!schema.is_required(2));
assert!(schema.is_required(3));
assert_eq!(schema.required_count(), 2);
}
}