diff --git a/crates/inputs/inputs/src/lib.rs b/crates/inputs/inputs/src/lib.rs index d8c55695e..3e9858927 100644 --- a/crates/inputs/inputs/src/lib.rs +++ b/crates/inputs/inputs/src/lib.rs @@ -42,7 +42,8 @@ pub mod prelude { #[cfg(feature = "server")] pub mod server { pub use crate::server::{ - InputRebroadcaster, InputSystems, ServerInputConfig, ServerInputPlugin, + InputRebroadcaster, InputSystems, InputValidationAppExt, ServerInputConfig, + ServerInputPlugin, }; } } diff --git a/crates/inputs/inputs/src/server.rs b/crates/inputs/inputs/src/server.rs index af2023164..a449bb76c 100644 --- a/crates/inputs/inputs/src/server.rs +++ b/crates/inputs/inputs/src/server.rs @@ -116,12 +116,48 @@ pub type InputSet = InputSystems; #[derive(SystemSet, Debug, Hash, PartialEq, Eq, Clone, Copy)] pub enum InputSystems { + /// Validate / sanitize received [`InputMessage`]s before they are applied to + /// the [`InputBuffer`]. Runs after `MessageSystems::Receive` and before + /// [`Self::ReceiveInputs`]. Empty by default — add systems here (see + /// [`InputValidationAppExt::add_input_validator`]) that mutate or drop + /// messages via [`MessageReceiver::retain_messages`]. A game that wants to + /// authorize input targets against `ControlledBy` can do so here. + ValidateInputs, /// Receive the latest ActionDiffs from the client ReceiveInputs, /// Use the ActionDiff received from the client to update the `ActionState` UpdateActionState, } +/// App-builder helper to register a server-side input-validation system. +/// +/// The system runs in [`InputSystems::ValidateInputs`] — after messages are +/// received, before they are buffered — so it can mutate or drop them with full +/// ECS access (any `SystemParam`). It typically queries +/// `Query<&mut MessageReceiver>>` and calls +/// [`MessageReceiver::retain_messages`]. This is sugar for +/// `app.add_systems(PreUpdate, system.in_set(InputSystems::ValidateInputs))`. +/// +/// Validators in the set are unordered relative to each other. To make one run +/// before another, pass an ordered config — e.g. +/// `app.add_input_validator(my_validator.after(other_validator))`. +pub trait InputValidationAppExt { + fn add_input_validator( + &mut self, + systems: impl IntoScheduleConfigs, + ) -> &mut Self; +} + +impl InputValidationAppExt for App { + fn add_input_validator( + &mut self, + systems: impl IntoScheduleConfigs, + ) -> &mut Self { + self.add_systems(PreUpdate, systems.in_set(InputSystems::ValidateInputs)); + self + } +} + /// Component that is used to customize how inputs will be rebroadcasted /// /// If absent, the inputs received on a given `ClientOf` entity will be rebroadcasted to all other clients @@ -169,7 +205,12 @@ impl Plugin for ServerInputPlugin { // - but host-server broadcasting their inputs only updates `state` app.configure_sets( PreUpdate, - (MessageSystems::Receive, InputSystems::ReceiveInputs).chain(), + ( + MessageSystems::Receive, + InputSystems::ValidateInputs, + InputSystems::ReceiveInputs, + ) + .chain(), ); app.configure_sets(FixedPreUpdate, InputSystems::UpdateActionState); diff --git a/crates/tests/src/client_server/input/leafwing.rs b/crates/tests/src/client_server/input/leafwing.rs index 9e1d488b7..b0e5c4a8c 100644 --- a/crates/tests/src/client_server/input/leafwing.rs +++ b/crates/tests/src/client_server/input/leafwing.rs @@ -562,3 +562,298 @@ fn test_input_message_with_huge_end_tick_does_not_allocate_unbounded_buffer() { `is_input_within_lookahead` in lightyear_inputs::server.", ); } + +/// Example + test for the game-side input-validation seam: a normal Bevy system +/// registered with `add_input_validator` runs in `InputSystems::ValidateInputs` +/// (after receive, before buffering) with **full ECS access**, and mutates/drops +/// received `InputMessage`s in place via `MessageReceiver::retain_messages`. +/// +/// Here the validator reads a `Res` (proving arbitrary `SystemParam` +/// access) and drops every input message while the flag is set, so a legitimate, +/// authorized key press never reaches the server's `ActionState`. +#[test] +fn test_input_validator_system_can_drop_messages() { + use bevy::ecs::resource::Resource; + use bevy::ecs::system::{Query, Res}; + use lightyear::input::leafwing::input_message::LeafwingSequence; + use lightyear_inputs::input_message::InputMessage; + use lightyear_inputs::prelude::server::InputValidationAppExt; + use lightyear_messages::prelude::MessageReceiver; + + #[derive(Resource)] + struct RejectInputs(bool); + + // A game-side validation system: full ECS access (reads a resource), drops + // the input messages in place. A real validator would clamp/inspect against + // game state rather than reject wholesale. + fn reject_inputs( + reject: Res, + mut receivers: Query<&mut MessageReceiver>>>, + ) { + if !reject.0 { + return; + } + for mut receiver in &mut receivers { + receiver.retain_messages(|_msg| false); + } + } + + let mut stepper = ClientServerStepper::from_config(StepperConfig::with_netcode_clients(1)); + stepper.server_app.insert_resource(RejectInputs(true)); + stepper.server_app.add_input_validator(reject_inputs); + + let server_entity = stepper + .server_app + .world_mut() + .spawn(( + ActionState::::default(), + Replicate::to_clients(NetworkTarget::All), + )) + .id(); + stepper.frame_step(2); + + let local = stepper + .client(0) + .get::() + .unwrap() + .entity_mapper + .get_local(server_entity) + .expect("entity replicated to client 0"); + stepper.client_apps[0] + .world_mut() + .entity_mut(local) + .insert(InputMap::::new([( + LeafwingInput1::Jump, + KeyCode::KeyA, + )])); + stepper.frame_step(1); + stepper.client_apps[0] + .world_mut() + .resource_mut::>() + .press(KeyCode::KeyA); + stepper.frame_step(10); + + let server_state = stepper + .server_app + .world() + .entity(server_entity) + .get::>() + .expect("entity has ActionState"); + assert!( + !server_state.pressed(&LeafwingInput1::Jump), + "input reached the server even though the validation system dropped \ + every message in ValidateInputs — the seam isn't running before \ + ReceiveInputs, or retain_messages didn't take effect.", + ); +} + +/// Example: a *game-supplied* `ValidateInputs` system implements input-target +/// authorization itself — lightyear does not enforce `ControlledBy` (it's an +/// optional helper). The validator drops any `InputTarget::Entity` the sender +/// doesn't control (via `ControlledByRemote` + `retain_messages`). +/// +/// Client 0 controls entity A and forges an input also targeting entity B +/// (uncontrolled). The validator must let A's input through (non-overblocking) +/// and drop B's — so A's server `ActionState` is pressed and B's is not. +#[test] +fn test_user_validator_can_authorize_targets() { + use bevy::ecs::relationship::RelationshipTarget; + use bevy::ecs::system::Query; + use lightyear::input::leafwing::input_message::LeafwingSequence; + use lightyear_core::id::RemoteId; + use lightyear_inputs::input_message::{InputMessage, InputTarget}; + use lightyear_inputs::prelude::server::InputValidationAppExt; + use lightyear_messages::prelude::MessageReceiver; + use lightyear_replication::control::ControlledByRemote; + use lightyear_replication::prelude::ControlledBy; + + // Game-side authorization, expressed as an ordinary ValidateInputs system. + fn authorize_targets( + mut receivers: Query<( + &RemoteId, + Option<&ControlledByRemote>, + &mut MessageReceiver>>, + )>, + ) { + for (client_id, controlled, mut receiver) in &mut receivers { + if client_id.is_local() { + continue; + } + receiver.retain_messages(|msg| { + msg.inputs.retain(|data| match data.target { + InputTarget::Entity(e) => { + controlled.is_some_and(|c| c.collection().contains(&e)) + } + InputTarget::PreSpawned(_) => true, + }); + !msg.inputs.is_empty() + }); + } + } + + let mut stepper = ClientServerStepper::from_config(StepperConfig::with_netcode_clients(1)); + stepper.server_app.add_input_validator(authorize_targets); + + let client_of_0 = stepper.client_of(0).id(); + // Entity A: controlled by client 0. + let entity_a = stepper + .server_app + .world_mut() + .spawn(( + ActionState::::default(), + Replicate::to_clients(NetworkTarget::All), + ControlledBy { + owner: client_of_0, + lifetime: Default::default(), + }, + )) + .id(); + // Entity B: replicated to client 0 but NOT controlled by it (the spoof victim). + let entity_b = stepper + .server_app + .world_mut() + .spawn(( + ActionState::::default(), + Replicate::to_clients(NetworkTarget::All), + )) + .id(); + stepper.frame_step(10); + + let local_a = stepper + .client(0) + .get::() + .unwrap() + .entity_mapper + .get_local(entity_a) + .expect("A replicated to client 0"); + let local_b = stepper + .client(0) + .get::() + .unwrap() + .entity_mapper + .get_local(entity_b) + .expect("B replicated to client 0"); + + // Client 0 puts an InputMap on BOTH its own entity and the victim's, so its + // outgoing message targets A (legit) and B (spoofed). + for local in [local_a, local_b] { + stepper.client_apps[0].world_mut().entity_mut(local).insert( + InputMap::::new([(LeafwingInput1::Jump, KeyCode::KeyA)]), + ); + } + stepper.frame_step(1); + stepper.client_apps[0] + .world_mut() + .resource_mut::>() + .press(KeyCode::KeyA); + stepper.frame_step(10); + + // Non-overblocking: A's authorized input reached the server. + assert!( + stepper + .server_app + .world() + .entity(entity_a) + .get::>() + .unwrap() + .pressed(&LeafwingInput1::Jump), + "the authorized input for A did not land — the validator over-stripped", + ); + // The spoofed input for B was dropped by the validator. + assert!( + !stepper + .server_app + .world() + .entity(entity_b) + .get::>() + .unwrap() + .pressed(&LeafwingInput1::Jump), + "spoofed input landed on victim B's ActionState", + ); +} + +/// `retain_received_messages` exposes per-message metadata (`remote_tick`, +/// `channel_kind`, `message_id`) that `retain_messages` hides — needed for +/// rate-limit / tick-window / replay validators. Here a validator reads +/// `remote_tick` and drops the message; we assert both that the metadata was +/// readable and that the drop took effect. +#[test] +fn test_validator_can_read_message_metadata() { + use bevy::ecs::resource::Resource; + use bevy::ecs::system::{Query, ResMut}; + use lightyear::input::leafwing::input_message::LeafwingSequence; + use lightyear_inputs::input_message::InputMessage; + use lightyear_inputs::prelude::server::InputValidationAppExt; + use lightyear_messages::prelude::MessageReceiver; + + #[derive(Resource, Default)] + struct SeenRemoteTick(Option); + + fn inspect_metadata( + mut seen: ResMut, + mut receivers: Query<&mut MessageReceiver>>>, + ) { + for mut receiver in &mut receivers { + receiver.retain_received_messages(|metadata, _data| { + // Metadata is reachable here (read-only), unlike `retain_messages`. + seen.0 = Some(metadata.remote_tick.0); + false // drop the message + }); + } + } + + let mut stepper = ClientServerStepper::from_config(StepperConfig::with_netcode_clients(1)); + stepper.server_app.init_resource::(); + stepper.server_app.add_input_validator(inspect_metadata); + + let server_entity = stepper + .server_app + .world_mut() + .spawn(( + ActionState::::default(), + Replicate::to_clients(NetworkTarget::All), + )) + .id(); + stepper.frame_step(2); + + let local = stepper + .client(0) + .get::() + .unwrap() + .entity_mapper + .get_local(server_entity) + .expect("entity replicated to client 0"); + stepper.client_apps[0] + .world_mut() + .entity_mut(local) + .insert(InputMap::::new([( + LeafwingInput1::Jump, + KeyCode::KeyA, + )])); + stepper.frame_step(1); + stepper.client_apps[0] + .world_mut() + .resource_mut::>() + .press(KeyCode::KeyA); + stepper.frame_step(10); + + assert!( + stepper + .server_app + .world() + .resource::() + .0 + .is_some(), + "validator never observed a message's remote_tick metadata", + ); + assert!( + !stepper + .server_app + .world() + .entity(server_entity) + .get::>() + .unwrap() + .pressed(&LeafwingInput1::Jump), + "input landed even though retain_received_messages dropped the message", + ); +} diff --git a/crates/transport/messages/src/receive.rs b/crates/transport/messages/src/receive.rs index 57c6e3047..56ba655d3 100644 --- a/crates/transport/messages/src/receive.rs +++ b/crates/transport/messages/src/receive.rs @@ -67,6 +67,19 @@ pub struct ReceivedMessage { pub message_id: Option, } +/// Read-only per-message metadata handed to +/// [`MessageReceiver::retain_received_messages`] validators alongside `&mut` +/// access to the message data. +#[derive(Debug, Clone, Copy)] +pub struct MessageMetadata { + /// Tick on the remote peer when the message was sent. + pub remote_tick: Tick, + /// Channel the message was sent on. + pub channel_kind: ChannelKind, + /// MessageId of the message, if the channel assigns one. + pub message_id: Option, +} + impl Default for MessageReceiver { fn default() -> Self { Self { recv: Vec::new() } @@ -90,6 +103,42 @@ impl MessageReceiver { self.recv.drain(..) } + /// Mutate and/or drop the buffered messages in place, *before* they are + /// consumed by the receiving system. + /// + /// This is the hook for validation/sanitization systems that run between + /// message receipt and whatever consumes the messages (e.g. server-side + /// input validation between `MessageSystems::Receive` and the input-buffer + /// apply). Returning `false` from `keep` drops that message; mutating the + /// `&mut M` rewrites it. Per-message metadata (remote tick, channel, + /// message id) is preserved automatically — unlike drain-then-re-push. + pub fn retain_messages(&mut self, mut keep: impl FnMut(&mut M) -> bool) { + self.recv.retain_mut(|received| keep(&mut received.data)); + } + + /// Like [`retain_messages`](Self::retain_messages), but the predicate also + /// gets the per-message [`MessageMetadata`] (`remote_tick`, `channel_kind`, + /// `message_id`) that `retain_messages` hides. + /// + /// Use this when validation needs the metadata, e.g. rate limiting, + /// tick-window / staleness checks, replay diagnostics, or per-channel + /// policy. The metadata is passed **by value (read-only)** — only the + /// message data is `&mut` (mutate to rewrite, return `false` to drop) — so a + /// validator can't accidentally rewrite the wire metadata. + pub fn retain_received_messages( + &mut self, + mut keep: impl FnMut(MessageMetadata, &mut M) -> bool, + ) { + self.recv.retain_mut(|received| { + let metadata = MessageMetadata { + remote_tick: received.remote_tick, + channel_kind: received.channel_kind, + message_id: received.message_id, + }; + keep(metadata, &mut received.data) + }); + } + pub fn num_messages(&self) -> usize { self.recv.len() }