diff --git a/crates/ember-core/src/shard.rs b/crates/ember-core/src/shard.rs index 975f0455..5380745b 100644 --- a/crates/ember-core/src/shard.rs +++ b/crates/ember-core/src/shard.rs @@ -280,6 +280,14 @@ pub enum ShardRequest { ProtoType { key: String, }, + /// Writes a ProtoRegister AOF record (no keyspace mutation). + /// Broadcast to all shards after a schema registration so the + /// schema is recovered from any shard's AOF on restart. + #[cfg(feature = "protobuf")] + ProtoRegisterAof { + name: String, + descriptor: Bytes, + }, } /// The shard's response to a request. @@ -875,6 +883,10 @@ fn dispatch(ks: &mut Keyspace, req: &ShardRequest) -> ShardResponse { Ok(name) => ShardResponse::ProtoTypeName(name), Err(_) => ShardResponse::WrongType, }, + // ProtoRegisterAof is a no-op for the keyspace — the AOF record + // is written by the to_aof_record path after dispatch returns Ok. + #[cfg(feature = "protobuf")] + ShardRequest::ProtoRegisterAof { .. } => ShardResponse::Ok, // snapshot/rewrite/flush_async are handled in the main loop, not here ShardRequest::Snapshot | ShardRequest::RewriteAof | ShardRequest::FlushDbAsync => { ShardResponse::Ok @@ -1037,6 +1049,13 @@ fn to_aof_record(req: &ShardRequest, resp: &ShardResponse) -> Option expire_ms, }) } + #[cfg(feature = "protobuf")] + (ShardRequest::ProtoRegisterAof { name, descriptor }, ShardResponse::Ok) => { + Some(AofRecord::ProtoRegister { + name: name.clone(), + descriptor: descriptor.clone(), + }) + } _ => None, } } diff --git a/crates/ember-protocol/src/command.rs b/crates/ember-protocol/src/command.rs index 87f4100e..b6eaac8c 100644 --- a/crates/ember-protocol/src/command.rs +++ b/crates/ember-protocol/src/command.rs @@ -321,6 +321,36 @@ pub enum Command { /// PUBSUB NUMPAT. Returns the number of active pattern subscriptions. PubSubNumPat, + // --- protobuf commands --- + /// PROTO.REGISTER `name` `descriptor_bytes`. Registers a protobuf schema + /// (pre-compiled FileDescriptorSet) under the given name. + ProtoRegister { name: String, descriptor: Bytes }, + + /// PROTO.SET `key` `type_name` `data` \[EX s | PX ms\] \[NX | XX\]. + /// Stores a validated protobuf value. + ProtoSet { + key: String, + type_name: String, + data: Bytes, + expire: Option, + /// Only set the key if it does not already exist. + nx: bool, + /// Only set the key if it already exists. + xx: bool, + }, + + /// PROTO.GET `key`. Returns \[type_name, data\] or nil. + ProtoGet { key: String }, + + /// PROTO.TYPE `key`. Returns the message type name or nil. + ProtoType { key: String }, + + /// PROTO.SCHEMAS. Lists all registered schema names. + ProtoSchemas, + + /// PROTO.DESCRIBE `name`. Lists message types in a registered schema. + ProtoDescribe { name: String }, + /// AUTH \[username\] password. Authenticate the connection. Auth { /// Username for ACL-style auth. None for legacy AUTH. @@ -447,6 +477,12 @@ impl Command { Command::PubSubChannels { .. } => "pubsub", Command::PubSubNumSub { .. } => "pubsub", Command::PubSubNumPat => "pubsub", + Command::ProtoRegister { .. } => "proto.register", + Command::ProtoSet { .. } => "proto.set", + Command::ProtoGet { .. } => "proto.get", + Command::ProtoType { .. } => "proto.type", + Command::ProtoSchemas => "proto.schemas", + Command::ProtoDescribe { .. } => "proto.describe", Command::Auth { .. } => "auth", Command::Quit => "quit", Command::Unknown(_) => "unknown", @@ -544,6 +580,12 @@ impl Command { "PUNSUBSCRIBE" => parse_punsubscribe(&frames[1..]), "PUBLISH" => parse_publish(&frames[1..]), "PUBSUB" => parse_pubsub(&frames[1..]), + "PROTO.REGISTER" => parse_proto_register(&frames[1..]), + "PROTO.SET" => parse_proto_set(&frames[1..]), + "PROTO.GET" => parse_proto_get(&frames[1..]), + "PROTO.TYPE" => parse_proto_type(&frames[1..]), + "PROTO.SCHEMAS" => parse_proto_schemas(&frames[1..]), + "PROTO.DESCRIBE" => parse_proto_describe(&frames[1..]), "AUTH" => parse_auth(&frames[1..]), "QUIT" => parse_quit(&frames[1..]), _ => Ok(Command::Unknown(name)), @@ -1694,6 +1736,125 @@ fn parse_pubsub(args: &[Frame]) -> Result { } } +// --- proto command parsers --- + +fn parse_proto_register(args: &[Frame]) -> Result { + if args.len() != 2 { + return Err(ProtocolError::WrongArity("PROTO.REGISTER".into())); + } + let name = extract_string(&args[0])?; + let descriptor = extract_bytes(&args[1])?; + Ok(Command::ProtoRegister { name, descriptor }) +} + +fn parse_proto_set(args: &[Frame]) -> Result { + if args.len() < 3 { + return Err(ProtocolError::WrongArity("PROTO.SET".into())); + } + + let key = extract_string(&args[0])?; + let type_name = extract_string(&args[1])?; + let data = extract_bytes(&args[2])?; + + let mut expire = None; + let mut nx = false; + let mut xx = false; + let mut idx = 3; + + while idx < args.len() { + let flag = extract_string(&args[idx])?.to_ascii_uppercase(); + match flag.as_str() { + "NX" => { + nx = true; + idx += 1; + } + "XX" => { + xx = true; + idx += 1; + } + "EX" => { + idx += 1; + if idx >= args.len() { + return Err(ProtocolError::WrongArity("PROTO.SET".into())); + } + let amount = parse_u64(&args[idx], "PROTO.SET")?; + if amount == 0 { + return Err(ProtocolError::InvalidCommandFrame( + "invalid expire time in 'PROTO.SET' command".into(), + )); + } + expire = Some(SetExpire::Ex(amount)); + idx += 1; + } + "PX" => { + idx += 1; + if idx >= args.len() { + return Err(ProtocolError::WrongArity("PROTO.SET".into())); + } + let amount = parse_u64(&args[idx], "PROTO.SET")?; + if amount == 0 { + return Err(ProtocolError::InvalidCommandFrame( + "invalid expire time in 'PROTO.SET' command".into(), + )); + } + expire = Some(SetExpire::Px(amount)); + idx += 1; + } + _ => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "unsupported PROTO.SET option '{flag}'" + ))); + } + } + } + + if nx && xx { + return Err(ProtocolError::InvalidCommandFrame( + "XX and NX options at the same time are not compatible".into(), + )); + } + + Ok(Command::ProtoSet { + key, + type_name, + data, + expire, + nx, + xx, + }) +} + +fn parse_proto_get(args: &[Frame]) -> Result { + if args.len() != 1 { + return Err(ProtocolError::WrongArity("PROTO.GET".into())); + } + let key = extract_string(&args[0])?; + Ok(Command::ProtoGet { key }) +} + +fn parse_proto_type(args: &[Frame]) -> Result { + if args.len() != 1 { + return Err(ProtocolError::WrongArity("PROTO.TYPE".into())); + } + let key = extract_string(&args[0])?; + Ok(Command::ProtoType { key }) +} + +fn parse_proto_schemas(args: &[Frame]) -> Result { + if !args.is_empty() { + return Err(ProtocolError::WrongArity("PROTO.SCHEMAS".into())); + } + Ok(Command::ProtoSchemas) +} + +fn parse_proto_describe(args: &[Frame]) -> Result { + if args.len() != 1 { + return Err(ProtocolError::WrongArity("PROTO.DESCRIBE".into())); + } + let name = extract_string(&args[0])?; + Ok(Command::ProtoDescribe { name }) +} + fn parse_auth(args: &[Frame]) -> Result { match args.len() { 1 => { @@ -3922,4 +4083,159 @@ mod tests { let err = Command::from_frame(cmd(&["QUIT", "extra"])).unwrap_err(); assert!(matches!(err, ProtocolError::WrongArity(_))); } + + // --- PROTO.REGISTER --- + + #[test] + fn proto_register_basic() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.REGISTER", "myschema", "descriptor"])).unwrap(), + Command::ProtoRegister { + name: "myschema".into(), + descriptor: Bytes::from("descriptor"), + }, + ); + } + + #[test] + fn proto_register_case_insensitive() { + assert!(Command::from_frame(cmd(&["proto.register", "s", "d"])).is_ok()); + } + + #[test] + fn proto_register_wrong_arity() { + let err = Command::from_frame(cmd(&["PROTO.REGISTER", "only"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + // --- PROTO.SET --- + + #[test] + fn proto_set_basic() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.SET", "key1", "my.Type", "data"])).unwrap(), + Command::ProtoSet { + key: "key1".into(), + type_name: "my.Type".into(), + data: Bytes::from("data"), + expire: None, + nx: false, + xx: false, + }, + ); + } + + #[test] + fn proto_set_with_ex() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.SET", "k", "t", "d", "EX", "60"])).unwrap(), + Command::ProtoSet { + key: "k".into(), + type_name: "t".into(), + data: Bytes::from("d"), + expire: Some(SetExpire::Ex(60)), + nx: false, + xx: false, + }, + ); + } + + #[test] + fn proto_set_with_px_and_nx() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.SET", "k", "t", "d", "PX", "5000", "NX"])).unwrap(), + Command::ProtoSet { + key: "k".into(), + type_name: "t".into(), + data: Bytes::from("d"), + expire: Some(SetExpire::Px(5000)), + nx: true, + xx: false, + }, + ); + } + + #[test] + fn proto_set_nx_xx_conflict() { + let err = Command::from_frame(cmd(&["PROTO.SET", "k", "t", "d", "NX", "XX"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + #[test] + fn proto_set_wrong_arity() { + let err = Command::from_frame(cmd(&["PROTO.SET", "k", "t"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + #[test] + fn proto_set_zero_expiry() { + let err = Command::from_frame(cmd(&["PROTO.SET", "k", "t", "d", "EX", "0"])).unwrap_err(); + assert!(matches!(err, ProtocolError::InvalidCommandFrame(_))); + } + + // --- PROTO.GET --- + + #[test] + fn proto_get_basic() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.GET", "key1"])).unwrap(), + Command::ProtoGet { key: "key1".into() }, + ); + } + + #[test] + fn proto_get_wrong_arity() { + let err = Command::from_frame(cmd(&["PROTO.GET"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + // --- PROTO.TYPE --- + + #[test] + fn proto_type_basic() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.TYPE", "key1"])).unwrap(), + Command::ProtoType { key: "key1".into() }, + ); + } + + #[test] + fn proto_type_wrong_arity() { + let err = Command::from_frame(cmd(&["PROTO.TYPE"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + // --- PROTO.SCHEMAS --- + + #[test] + fn proto_schemas_basic() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.SCHEMAS"])).unwrap(), + Command::ProtoSchemas, + ); + } + + #[test] + fn proto_schemas_wrong_arity() { + let err = Command::from_frame(cmd(&["PROTO.SCHEMAS", "extra"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } + + // --- PROTO.DESCRIBE --- + + #[test] + fn proto_describe_basic() { + assert_eq!( + Command::from_frame(cmd(&["PROTO.DESCRIBE", "myschema"])).unwrap(), + Command::ProtoDescribe { + name: "myschema".into() + }, + ); + } + + #[test] + fn proto_describe_wrong_arity() { + let err = Command::from_frame(cmd(&["PROTO.DESCRIBE"])).unwrap_err(); + assert!(matches!(err, ProtocolError::WrongArity(_))); + } } diff --git a/crates/ember-server/src/connection.rs b/crates/ember-server/src/connection.rs index efe2edcc..c28664aa 100644 --- a/crates/ember-server/src/connection.rs +++ b/crates/ember-server/src/connection.rs @@ -554,7 +554,10 @@ async fn cluster_slot_check(ctx: &ServerContext, cmd: &Command) -> Option | Command::SRem { ref key, .. } | Command::SMembers { ref key } | Command::SIsMember { ref key, .. } - | Command::SCard { ref key } => cluster.check_slot(key.as_bytes()).await, + | Command::SCard { ref key } + | Command::ProtoSet { ref key, .. } + | Command::ProtoGet { ref key } + | Command::ProtoType { ref key } => cluster.check_slot(key.as_bytes()).await, // multi-key commands — crossslot validation + slot ownership Command::Del { ref keys } @@ -1621,6 +1624,167 @@ async fn execute( Frame::Error("ERR subscribe commands should not reach execute".into()) } + // -- protobuf commands -- + #[cfg(feature = "protobuf")] + Command::ProtoRegister { name, descriptor } => { + let registry = match engine.schema_registry() { + Some(r) => r, + None => return Frame::Error("ERR protobuf support is not enabled".into()), + }; + let result = { + let mut reg = match registry.write() { + Ok(r) => r, + Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), + }; + reg.register(name.clone(), descriptor.clone()) + }; + match result { + Ok(types) => { + // persist the registration to all shards' AOF + let _ = engine + .broadcast(|| ShardRequest::ProtoRegisterAof { + name: name.clone(), + descriptor: descriptor.clone(), + }) + .await; + Frame::Array( + types + .into_iter() + .map(|t| Frame::Bulk(Bytes::from(t))) + .collect(), + ) + } + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + #[cfg(feature = "protobuf")] + Command::ProtoSet { + key, + type_name, + data, + expire, + nx, + xx, + } => { + let registry = match engine.schema_registry() { + Some(r) => r, + None => return Frame::Error("ERR protobuf support is not enabled".into()), + }; + // validate the bytes against the schema before storing + { + let reg = match registry.read() { + Ok(r) => r, + Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), + }; + if let Err(e) = reg.validate(&type_name, &data) { + return Frame::Error(format!("ERR {e}")); + } + } + let duration = expire.map(|e| match e { + SetExpire::Ex(secs) => Duration::from_secs(secs), + SetExpire::Px(millis) => Duration::from_millis(millis), + }); + let req = ShardRequest::ProtoSet { + key: key.clone(), + type_name, + data, + expire: duration, + nx, + xx, + }; + match engine.route(&key, req).await { + Ok(ShardResponse::Ok) => Frame::Simple("OK".into()), + Ok(ShardResponse::Value(None)) => Frame::Null, + Ok(ShardResponse::OutOfMemory) => oom_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + #[cfg(feature = "protobuf")] + Command::ProtoGet { key } => { + if engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let req = ShardRequest::ProtoGet { key: key.clone() }; + match engine.route(&key, req).await { + Ok(ShardResponse::ProtoValue(Some((type_name, data)))) => { + Frame::Array(vec![Frame::Bulk(Bytes::from(type_name)), Frame::Bulk(data)]) + } + Ok(ShardResponse::ProtoValue(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + #[cfg(feature = "protobuf")] + Command::ProtoType { key } => { + if engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let req = ShardRequest::ProtoType { key: key.clone() }; + match engine.route(&key, req).await { + Ok(ShardResponse::ProtoTypeName(Some(name))) => Frame::Bulk(Bytes::from(name)), + Ok(ShardResponse::ProtoTypeName(None)) => Frame::Null, + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + #[cfg(feature = "protobuf")] + Command::ProtoSchemas => { + let registry = match engine.schema_registry() { + Some(r) => r, + None => return Frame::Error("ERR protobuf support is not enabled".into()), + }; + let reg = match registry.read() { + Ok(r) => r, + Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), + }; + let names = reg.schema_names(); + Frame::Array( + names + .into_iter() + .map(|n| Frame::Bulk(Bytes::from(n))) + .collect(), + ) + } + + #[cfg(feature = "protobuf")] + Command::ProtoDescribe { name } => { + let registry = match engine.schema_registry() { + Some(r) => r, + None => return Frame::Error("ERR protobuf support is not enabled".into()), + }; + let reg = match registry.read() { + Ok(r) => r, + Err(_) => return Frame::Error("ERR schema registry lock poisoned".into()), + }; + match reg.describe(&name) { + Some(types) => Frame::Array( + types + .into_iter() + .map(|t| Frame::Bulk(Bytes::from(t))) + .collect(), + ), + None => Frame::Error(format!("ERR unknown schema '{name}'")), + } + } + + // when protobuf feature is disabled, proto commands are unknown + #[cfg(not(feature = "protobuf"))] + Command::ProtoRegister { .. } + | Command::ProtoSet { .. } + | Command::ProtoGet { .. } + | Command::ProtoType { .. } + | Command::ProtoSchemas + | Command::ProtoDescribe { .. } => { + Frame::Error("ERR unknown command (protobuf support not compiled)".into()) + } + // AUTH on an already-authenticated connection (re-auth) Command::Auth { username, password } => match &ctx.requirepass { None => Frame::Error( diff --git a/crates/ember-server/src/main.rs b/crates/ember-server/src/main.rs index ace4dab6..2c6dba6c 100644 --- a/crates/ember-server/src/main.rs +++ b/crates/ember-server/src/main.rs @@ -124,6 +124,13 @@ struct Args { #[arg(long, env = "EMBER_ENCRYPTION_KEY_FILE")] encryption_key_file: Option, + // -- protobuf support -- + /// enable protobuf value storage. when set, PROTO.* commands become + /// available for schema-validated structured data. + #[cfg(feature = "protobuf")] + #[arg(long, env = "EMBER_PROTOBUF")] + protobuf: bool, + // -- cluster options -- /// enable cluster mode with gossip-based discovery and slot routing #[arg(long, env = "EMBER_CLUSTER_ENABLED")] @@ -251,7 +258,15 @@ async fn main() { None }; - let engine_config = build_engine_config(max_memory, eviction_policy, shard_count, persistence); + #[allow(unused_mut)] + let mut engine_config = + build_engine_config(max_memory, eviction_policy, shard_count, persistence); + + #[cfg(feature = "protobuf")] + if args.protobuf { + engine_config.schema_registry = Some(ember_core::schema::SchemaRegistry::shared()); + info!("protobuf value storage enabled"); + } if let Some(limit) = max_memory { info!(