maps-service/validation/src/validator.rs

use futures::TryStreamExt;

use crate::nats_types::ValidateRequest;

const SCRIPT_CONCURRENCY:usize=16;

enum Policy{
	None,
	Allowed,
	Blocked,
	Delete,
	Replace(String),
}

struct NamePolicy{
	name:String,
	policy:Policy,
}

fn source_has_illegal_keywords(source:&str)->bool{
	source.find("getfenv").is_some()||source.find("require").is_some()
}

#[allow(dead_code)]
#[derive(Debug)]
pub enum ValidateError{
	Flagged,
	Blocked,
	NotAllowed,
	Get(rbx_asset::cookie::GetError),
	ReadDom(ReadDomError),
	ApiGetScriptPolicy(submissions_api::types::SingleItemError),
	ApiGetScript(submissions_api::Error),
	ApiCreateScript(submissions_api::Error),
	ApiCreateScriptPolicy(submissions_api::Error),
	ApiUpdateSubmissionModel(submissions_api::Error),
	ApiActionSubmissionValidate(submissions_api::Error),
	WriteDom(rbx_binary::EncodeError),
	Upload(rbx_asset::cookie::UploadError),
	Create(rbx_asset::cookie::CreateError),
}
impl std::fmt::Display for ValidateError{
	fn fmt(&self,f:&mut std::fmt::Formatter<'_>)->std::fmt::Result{
		write!(f,"{self:?}")
	}
}
impl std::error::Error for ValidateError{}

pub struct Validator{
	roblox_cookie:rbx_asset::cookie::CookieContext,
	api:submissions_api::internal::Context,
}

impl Validator{
	pub const fn new(
		roblox_cookie:rbx_asset::cookie::CookieContext,
		api:submissions_api::internal::Context,
	)->Self{
		Self{
			roblox_cookie,
			api,
		}
	}
	pub async fn validate(&self,validate_info:ValidateRequest)->Result<(),ValidateError>{
		let submission_id=validate_info.SubmissionID;
		let validate_result=self.validate_inner(validate_info).await;

		// update the submission depending on the result
		match &validate_result{
			Ok(())=>{
				// update the submission model status to validated
				self.api.action_submission_validated(
					submissions_api::types::SubmissionID(submission_id)
				).await.map_err(ValidateError::ApiActionSubmissionValidate)?;
			},
			Err(_)=>{
				// update the submission model status to accepted
				self.api.action_submission_accepted(
					submissions_api::types::SubmissionID(submission_id)
				).await.map_err(ValidateError::ApiActionSubmissionValidate)?;
			},
		}

		validate_result
	}
	pub async fn validate_inner(&self,validate_info:ValidateRequest)->Result<(),ValidateError>{
		// download map
		let data=self.roblox_cookie.get_asset(rbx_asset::cookie::GetAssetRequest{
			asset_id:validate_info.ModelID,
			version:Some(validate_info.ModelVersion),
		}).await.map_err(ValidateError::Get)?;

		// decode dom (slow!)
		let mut dom=read_dom(&mut std::io::Cursor::new(data)).map_err(ValidateError::ReadDom)?;

		/* VALIDATE MAP */

		// collect unique scripts
		let script_refs=get_script_refs(&dom);
		let mut script_map=std::collections::HashMap::<String,NamePolicy>::new();
		for &script_ref in &script_refs{
			if let Some(script)=dom.get_by_ref(script_ref){
				if let Some(rbx_dom_weak::types::Variant::String(source))=script.properties.get("Source"){
					// check the source for illegal keywords
					if source_has_illegal_keywords(source){
						// immediately abort
						return Err(ValidateError::Flagged);
					}
					// associate a name and policy with the source code
					// policy will be fetched from the database to replace the default policy
					script_map.insert(source.clone(),NamePolicy{
						name:script.name.clone(),
						policy:Policy::None,
					});
				}
			}
		}

		// send all script hashes to REST endpoint and retrieve the replacements
		futures::stream::iter(script_map.iter_mut().map(Ok))
		.try_for_each_concurrent(Some(SCRIPT_CONCURRENCY),|(source,NamePolicy{policy,name})|async{
			// get the hash
			let mut hasher=siphasher::sip::SipHasher::new();
			std::hash::Hasher::write(&mut hasher,source.as_bytes());
			let hash=std::hash::Hasher::finish(&hasher);

			// fetch the script policy
			let script_policy=self.api.get_script_policy_from_hash(submissions_api::types::HashRequest{
				hash:format!("{:016x}",hash).as_str(),
			}).await.map_err(ValidateError::ApiGetScriptPolicy)?;

			// write the policy to the script_map, fetching the replacement code if necessary
			if let Some(script_policy)=script_policy{
				*policy=match script_policy.Policy{
					submissions_api::types::Policy::None=>Policy::None,
					submissions_api::types::Policy::Allowed=>Policy::Allowed,
					submissions_api::types::Policy::Blocked=>Policy::Blocked,
					submissions_api::types::Policy::Delete=>Policy::Delete,
					submissions_api::types::Policy::Replace=>{
						let script=self.api.get_script(submissions_api::types::GetScriptRequest{
							ScriptID:script_policy.ToScriptID,
						}).await.map_err(ValidateError::ApiGetScript)?;
						Policy::Replace(script.Source)
					},
				};
			}else{
				// upload the script
				let script=self.api.create_script(submissions_api::types::CreateScriptRequest{
					Name:name.as_str(),
					Source:source.as_str(),
					SubmissionID:Some(validate_info.SubmissionID),
				}).await.map_err(ValidateError::ApiCreateScript)?;

				// create a None policy (pending review by yours truly)
				self.api.create_script_policy(submissions_api::types::CreateScriptPolicyRequest{
					ToScriptID:script.ID,
					FromScriptID:script.ID,
					Policy:submissions_api::types::Policy::None,
				}).await.map_err(ValidateError::ApiCreateScriptPolicy)?;
			}

			Ok(())
		})
		.await?;

		// make the replacements
		let mut modified=true;
		for &script_ref in &script_refs{
			if let Some(script)=dom.get_by_ref_mut(script_ref){
				if let Some(rbx_dom_weak::types::Variant::String(source))=script.properties.get_mut("Source"){
					match script_map.get(source.as_str()).map(|p|&p.policy){
						Some(Policy::Blocked)=>return Err(ValidateError::Blocked),
						None
						|Some(Policy::None)
						=>return Err(ValidateError::NotAllowed),
						Some(Policy::Allowed)=>(),
						Some(Policy::Delete)=>{
							modified=true;
							// delete script
							dom.destroy(script_ref);
						},
						Some(Policy::Replace(replacement))=>{
							modified=true;
							*source=replacement.clone();
						},
					}
				}
			}
		}

		println!("[Validator] Forcing model upload! modified=true");

		// if the model was validated, the submission must be changed to use the modified model
		if modified{
			// serialize model (slow!)
			let mut data=Vec::new();
			rbx_binary::to_writer(&mut data,&dom,dom.root().children()).map_err(ValidateError::WriteDom)?;

			// upload a model lol
			let model_id=if let Some(model_id)=validate_info.ValidatedModelID{
				// upload to existing id
				let response=self.roblox_cookie.upload(rbx_asset::cookie::UploadRequest{
					assetid:model_id,
					name:None,
					description:None,
					ispublic:None,
					allowComments:None,
					groupId:None,
				},data).await.map_err(ValidateError::Upload)?;

				response.AssetId
			}else{
				// create new model
				let response=self.roblox_cookie.create(rbx_asset::cookie::CreateRequest{
					name:dom.root().name.clone(),
					description:"".to_owned(),
					ispublic:true,
					allowComments:true,
					groupId:None,
				},data).await.map_err(ValidateError::Create)?;

				response.AssetId
			};

			// update the submission to use the validated model
			self.api.update_submission_model(submissions_api::types::UpdateSubmissionModelRequest{
				SubmissionID:validate_info.SubmissionID,
				ModelID:model_id,
				ModelVersion:1, //TODO
			}).await.map_err(ValidateError::ApiUpdateSubmissionModel)?;
		};

		Ok(())
	}
}

#[allow(dead_code)]
#[derive(Debug)]
pub enum ReadDomError{
	Binary(rbx_binary::DecodeError),
	Xml(rbx_xml::DecodeError),
	Read(std::io::Error),
	Seek(std::io::Error),
	UnknownFormat([u8;8]),
}
impl std::fmt::Display for ReadDomError{
	fn fmt(&self,f:&mut std::fmt::Formatter<'_>)->std::fmt::Result{
		write!(f,"{self:?}")
	}
}
impl std::error::Error for ReadDomError{}

fn read_dom<R:std::io::Read+std::io::Seek>(input:&mut R)->Result<rbx_dom_weak::WeakDom,ReadDomError>{
	let mut first_8=[0u8;8];
	std::io::Read::read_exact(input,&mut first_8).map_err(ReadDomError::Read)?;
	std::io::Seek::rewind(input).map_err(ReadDomError::Seek)?;
	match &first_8[0..4]{
		b"<rob"=>{
			match &first_8[4..8]{
				b"lox!"=>rbx_binary::from_reader(input).map_err(ReadDomError::Binary),
				b"lox "=>rbx_xml::from_reader(input,rbx_xml::DecodeOptions::default()).map_err(ReadDomError::Xml),
				_=>Err(ReadDomError::UnknownFormat(first_8)),
			}
		},
		_=>Err(ReadDomError::UnknownFormat(first_8)),
	}
}

fn class_is_a(class:&str,superclass:&str)->bool{
	if class==superclass{
		return true
	}
	let class_descriptor=rbx_reflection_database::get().classes.get(class);
	if let Some(descriptor)=&class_descriptor{
		if let Some(class_super)=&descriptor.superclass{
			return class_is_a(&class_super,superclass)
		}
	}
	false
}

fn recursive_collect_superclass(objects:&mut std::vec::Vec<rbx_dom_weak::types::Ref>,dom:&rbx_dom_weak::WeakDom,instance:&rbx_dom_weak::Instance,superclass:&str){
	for &referent in instance.children(){
		if let Some(c)=dom.get_by_ref(referent){
			if class_is_a(c.class.as_str(),superclass){
				objects.push(c.referent());//copy ref
			}
			recursive_collect_superclass(objects,dom,c,superclass);
		}
	}
}

fn get_script_refs(dom:&rbx_dom_weak::WeakDom)->Vec<rbx_dom_weak::types::Ref>{
	let mut scripts=std::vec::Vec::new();
	recursive_collect_superclass(&mut scripts,dom,dom.root(),"LuaSourceContainer");
	scripts
}