inline very thin function

Cargo.lock

@@ -1144,6 +1144,8 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
 [[package]]
 name = "submissions-api"
 version = "0.3.0"
+source = "sparse+https://git.itzana.me/api/packages/strafesnet/cargo/"
+checksum = "72e67dbf479fc6a5e22514208d533534a2e0543eab7c6c1b8860ee3f6f0c6290"
 dependencies = [
  "reqwest",
  "serde",

src/main.rs

@@ -1,4 +1,8 @@
 use clap::{Args,Parser,Subcommand};
+use futures::{StreamExt,TryStreamExt};
+
+const READ_CONCURRENCY:usize=16;
+const REMOTE_CONCURRENCY:usize=16;
 
 #[derive(Parser)]
 #[command(author,version,about,long_about=None)]
@@ -11,12 +15,22 @@ struct Cli{
 #[derive(Subcommand)]
 enum Commands{
 	Review(ReviewCommand),
+	UploadScripts(UploadScriptsCommand),
 }
 
 #[derive(Args)]
 struct ReviewCommand{
 	#[arg(long)]
-	cookie:String,
+	session_id:String,
+	#[arg(long)]
+	api_url:String,
+}
+#[derive(Args)]
+struct UploadScriptsCommand{
+	#[arg(long)]
+	session_id:String,
+	#[arg(long)]
+	api_url:String,
 }
 
 #[tokio::main]
@@ -24,7 +38,12 @@ async fn main(){
 	let cli=Cli::parse();
 	match cli.command{
 		Commands::Review(command)=>review(ReviewConfig{
-			cookie:command.cookie,
+			session_id:command.session_id,
+			api_url:command.api_url,
+		}).await.unwrap(),
+		Commands::UploadScripts(command)=>upload_scripts(UploadConfig{
+			session_id:command.session_id,
+			api_url:command.api_url,
 		}).await.unwrap(),
 	}
 }
@@ -66,14 +85,15 @@ enum ReviewError{
 }
 
 struct ReviewConfig{
-	cookie:String,
+	session_id:String,
+	api_url:String,
 }
 
 async fn review(config:ReviewConfig)->Result<(),ReviewError>{
 	// download unreviewed policies
 	// review them
-	let cookie=submissions_api::Cookie::new(&config.cookie).map_err(ReviewError::Cookie)?;
-	let api=submissions_api::external::Context::new("http://localhost:8083".to_owned(),cookie).map_err(ReviewError::Reqwest)?;
+	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(ReviewError::Cookie)?;
+	let api=submissions_api::external::Context::new(config.api_url,cookie).map_err(ReviewError::Reqwest)?;
 
 	let unreviewed_policies=api.get_script_policies(submissions_api::types::GetScriptPoliciesRequest{
 		Page:1,
@@ -119,9 +139,7 @@ async fn review(config:ReviewConfig)->Result<(),ReviewError>{
 					submissions_api::types::Policy::Allowed
 				}else{
 					// compute hash
-					let mut hasher=siphasher::sip::SipHasher::new();
-					std::hash::Hasher::write(&mut hasher,source.as_bytes());
-					let hash=std::hash::Hasher::finish(&hasher);
+					let hash=hash_source(source.as_str());
 
 					// check if modified script already exists
 					let maybe_script_response=api.get_script_from_hash(submissions_api::types::HashRequest{
@@ -149,7 +167,7 @@ async fn review(config:ReviewConfig)->Result<(),ReviewError>{
 
 		// update policy
 		api.update_script_policy(submissions_api::types::UpdateScriptPolicyRequest{
-			ScriptPolicyID:unreviewed_policy.ID,
+			ID:unreviewed_policy.ID,
 			FromScriptID:None,
 			ToScriptID:to_script_id,
 			Policy:Some(reviewed_policy),
@@ -158,3 +176,233 @@ async fn review(config:ReviewConfig)->Result<(),ReviewError>{
 
 	Ok(())
 }
+
+#[allow(dead_code)]
+#[derive(Debug)]
+enum ScriptUploadError{
+	Cookie(submissions_api::CookieError),
+	Reqwest(submissions_api::ReqwestError),
+	AllowedSet(std::io::Error),
+	AllowedMap(GetMapError),
+	ReplaceMap(GetMapError),
+	BlockedSet(std::io::Error),
+	GOC(GOCError),
+	GOCPolicyReplace(GOCError),
+	GOCPolicyAllowed(GOCError),
+	GOCPolicyBlocked(GOCError),
+}
+
+fn read_dir_stream(dir:tokio::fs::ReadDir)->impl futures::stream::Stream<Item=std::io::Result<tokio::fs::DirEntry>>{
+    futures::stream::unfold(dir,|mut dir|async{
+        match dir.next_entry().await{
+            Ok(Some(entry))=>Some((Ok(entry),dir)),
+            Ok(None)=>None, // End of directory
+            Err(e)=>Some((Err(e),dir)), // Error encountered
+        }
+    })
+}
+
+async fn get_set_from_file(path:impl AsRef<std::path::Path>)->std::io::Result<std::collections::HashSet<String>>{
+	read_dir_stream(tokio::fs::read_dir(path).await?)
+	.map(|dir_entry|async{
+		tokio::fs::read_to_string(dir_entry?.path()).await
+	})
+	.buffer_unordered(READ_CONCURRENCY)
+	.try_collect().await
+}
+
+async fn get_allowed_set()->std::io::Result<std::collections::HashSet<String>>{
+	get_set_from_file("scripts/allowed").await
+}
+
+async fn get_blocked_set()->std::io::Result<std::collections::HashSet<String>>{
+	get_set_from_file("scripts/blocked").await
+}
+
+#[allow(dead_code)]
+#[derive(Debug)]
+enum GetMapError{
+	IO(std::io::Error),
+	FileStem,
+	ToStr,
+	ParseInt(std::num::ParseIntError),
+}
+
+async fn get_allowed_map()->Result<std::collections::HashMap::<u32,String>,GetMapError>{
+	read_dir_stream(tokio::fs::read_dir("scripts/allowed").await.map_err(GetMapError::IO)?)
+	.map(|dir_entry|async{
+		let path=dir_entry.map_err(GetMapError::IO)?.path();
+		let id:u32=path
+		.file_stem().ok_or(GetMapError::FileStem)?
+		.to_str().ok_or(GetMapError::ToStr)?
+		.parse().map_err(GetMapError::ParseInt)?;
+		let source=tokio::fs::read_to_string(path).await.map_err(GetMapError::IO)?;
+		Ok((id,source))
+	})
+	.buffer_unordered(READ_CONCURRENCY)
+	.try_collect().await
+}
+
+async fn get_replace_map()->Result<std::collections::HashMap::<String,u32>,GetMapError>{
+	read_dir_stream(tokio::fs::read_dir("scripts/replace").await.map_err(GetMapError::IO)?)
+	.map(|dir_entry|async{
+		let path=dir_entry.map_err(GetMapError::IO)?.path();
+		let id:u32=path
+		.file_stem().ok_or(GetMapError::FileStem)?
+		.to_str().ok_or(GetMapError::ToStr)?
+		.parse().map_err(GetMapError::ParseInt)?;
+		let source=tokio::fs::read_to_string(path).await.map_err(GetMapError::IO)?;
+		Ok((source,id))
+	})
+	.buffer_unordered(READ_CONCURRENCY)
+	.try_collect().await
+}
+
+fn hash_source(source:&str)->u64{
+	let mut hasher=siphasher::sip::SipHasher::new();
+	std::hash::Hasher::write(&mut hasher,source.as_bytes());
+	std::hash::Hasher::finish(&hasher)
+}
+
+fn hash_format(hash:u64)->String{
+	format!("{:016x}",hash)
+}
+
+type GOCError=submissions_api::types::SingleItemError;
+type GOCResult=Result<submissions_api::types::ScriptID,GOCError>;
+
+async fn get_or_create_script(api:&submissions_api::external::Context,source:&str)->GOCResult{
+	let script_response=api.get_script_from_hash(submissions_api::types::HashRequest{
+		hash:hash_format(hash_source(source)).as_str(),
+	}).await?;
+
+	Ok(match script_response{
+		Some(script_response)=>script_response.ID,
+		None=>api.create_script(submissions_api::types::CreateScriptRequest{
+			Name:"Script",
+			Source:source,
+			SubmissionID:None,
+		}).await.map_err(GOCError::Other)?.ID
+	})
+}
+
+async fn check_or_create_script_poicy(
+	api:&submissions_api::external::Context,
+	hash:&str,
+	script_policy:submissions_api::types::CreateScriptPolicyRequest,
+)->Result<(),GOCError>{
+	let script_policy_result=api.get_script_policy_from_hash(submissions_api::types::HashRequest{
+		hash,
+	}).await?;
+
+	match script_policy_result{
+		Some(script_policy_reponse)=>{
+			// check that everything matches the expectation
+			assert!(hash==script_policy_reponse.FromScriptHash);
+			assert!(script_policy.ToScriptID==script_policy_reponse.ToScriptID);
+			assert!(script_policy.Policy==script_policy_reponse.Policy);
+		},
+		None=>{
+			// create a new policy
+			api.create_script_policy(script_policy).await.map_err(GOCError::Other)?;
+		}
+	}
+
+	Ok(())
+}
+
+struct UploadConfig{
+	session_id:String,
+	api_url:String,
+}
+
+async fn upload_scripts(config:UploadConfig)->Result<(),ScriptUploadError>{
+	let cookie=submissions_api::Cookie::new(&config.session_id).map_err(ScriptUploadError::Cookie)?;
+	let api=&submissions_api::external::Context::new(config.api_url,cookie).map_err(ScriptUploadError::Reqwest)?;
+
+	// load all script files
+	let (
+		allowed_set_result,
+		allowed_map_result,
+		replace_map_result,
+		blocked_set_result,
+	)=tokio::join!(
+		get_allowed_set(),
+		get_allowed_map(),
+		get_replace_map(),
+		get_blocked_set(),
+	);
+
+	let allowed_set=allowed_set_result.map_err(ScriptUploadError::AllowedSet)?;
+	let allowed_map=allowed_map_result.map_err(ScriptUploadError::AllowedMap)?;
+	let replace_map=replace_map_result.map_err(ScriptUploadError::ReplaceMap)?;
+	let blocked_set=blocked_set_result.map_err(ScriptUploadError::BlockedSet)?;
+
+	// create a unified deduplicated set of all scripts
+	let script_set:std::collections::HashSet<&str>=allowed_set.iter()
+	.map(String::as_str)
+	.chain(
+		replace_map.keys().map(String::as_str)
+	).chain(
+		blocked_set.iter().map(String::as_str)
+	).collect();
+
+	// get or create every unique script
+	let script_ids:std::collections::HashMap<&str,submissions_api::types::ScriptID>=
+	futures::stream::iter(script_set)
+	.map(|source|async move{
+		let script_id=get_or_create_script(api,source).await?;
+		Ok((source,script_id))
+	})
+	.buffer_unordered(REMOTE_CONCURRENCY)
+	.try_collect().await.map_err(ScriptUploadError::GOC)?;
+
+	// get or create policy for each script in each category
+	//
+	// replace
+	let replace_fut=futures::stream::iter(replace_map.iter().map(Ok))
+	.try_for_each_concurrent(Some(REMOTE_CONCURRENCY),|(source,id)|async{
+		check_or_create_script_poicy(
+			api,
+			hash_format(hash_source(source)).as_str(),
+			submissions_api::types::CreateScriptPolicyRequest{
+				FromScriptID:script_ids[source.as_str()],
+				ToScriptID:script_ids[allowed_map[id].as_str()],
+				Policy:submissions_api::types::Policy::Replace,
+			}
+		).await.map_err(ScriptUploadError::GOCPolicyReplace)
+	});
+
+	// allowed
+	let allowed_fut=futures::stream::iter(allowed_set.iter().map(Ok))
+	.try_for_each_concurrent(Some(REMOTE_CONCURRENCY),|source|async{
+		check_or_create_script_poicy(
+			api,
+			hash_format(hash_source(source)).as_str(),
+			submissions_api::types::CreateScriptPolicyRequest{
+				FromScriptID:script_ids[source.as_str()],
+				ToScriptID:script_ids[source.as_str()],
+				Policy:submissions_api::types::Policy::Allowed,
+			}
+		).await.map_err(ScriptUploadError::GOCPolicyAllowed)
+	});
+
+	// blocked
+	let blocked_fut=futures::stream::iter(blocked_set.iter().map(Ok))
+	.try_for_each_concurrent(Some(REMOTE_CONCURRENCY),|source|async{
+		check_or_create_script_poicy(
+			api,
+			hash_format(hash_source(source)).as_str(),
+			submissions_api::types::CreateScriptPolicyRequest{
+				FromScriptID:script_ids[source.as_str()],
+				ToScriptID:script_ids[source.as_str()],
+				Policy:submissions_api::types::Policy::Blocked,
+			}
+		).await.map_err(ScriptUploadError::GOCPolicyBlocked)
+	});
+
+	// run futures
+	tokio::try_join!(replace_fut,allowed_fut,blocked_fut)?;
+
+	Ok(())
+}