From 7ae1d02c02ef6e56c18ed20986c45b76c53e61d6 Mon Sep 17 00:00:00 2001
From: alemi <me@alemi.dev>
Date: Tue, 16 Jul 2024 01:19:47 +0200
Subject: [PATCH] fix: also cloak object intrinsic image

---
 upub/cli/src/cloak.rs             | 51 ++++++++++++++++++++++---------
 upub/cli/src/lib.rs               | 12 +++++---
 upub/core/src/traits/normalize.rs |  4 +++
 3 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/upub/cli/src/cloak.rs b/upub/cli/src/cloak.rs
index 4a4f60b..3eebcc5 100644
--- a/upub/cli/src/cloak.rs
+++ b/upub/cli/src/cloak.rs
@@ -2,7 +2,7 @@ use futures::TryStreamExt;
 use sea_orm::{ActiveModelTrait, ActiveValue::{NotSet, Set, Unchanged}, ColumnTrait, Condition, EntityTrait, IntoActiveModel, QueryFilter, QuerySelect, SelectColumns};
 use upub::traits::{fetch::RequestError, Cloaker};
 
-pub async fn cloak(ctx: upub::Context, post_contents: bool, actors: bool) -> Result<(), RequestError> {
+pub async fn cloak(ctx: upub::Context, post_contents: bool, objects: bool, actors: bool) -> Result<(), RequestError> {
 	let local_base = format!("{}%", ctx.base());
 	{
 		let mut stream = upub::model::attachment::Entity::find()
@@ -21,25 +21,23 @@ pub async fn cloak(ctx: upub::Context, post_contents: bool, actors: bool) -> Res
 
 	if post_contents {
 		let mut stream = upub::model::object::Entity::find()
-			.filter(upub::model::object::Column::Content.like("%<img%"))
+			.filter(upub::model::object::Column::Image.is_not_null())
+			.filter(upub::model::object::Column::Image.not_like(&local_base))
 			.select_only()
 			.select_column(upub::model::object::Column::Internal)
-			.select_column(upub::model::object::Column::Content)
+			.select_column(upub::model::object::Column::Image)
 			.into_tuple::<(i64, String)>()
 			.stream(ctx.db())
 			.await?;
 
-		while let Some((internal, content)) = stream.try_next().await? {
-			let sanitized = ctx.sanitize(&content);
-			if sanitized != content {
-				tracing::info!("sanitizing object #{internal}");
-				let model = upub::model::object::ActiveModel {
-					internal: Unchanged(internal),
-					content: Set(Some(sanitized)),
-					..Default::default()
-				};
-				model.update(ctx.db()).await?;
-			}
+		while let Some((internal, image)) = stream.try_next().await? {
+			tracing::info!("cloaking object image {image}");
+			let model = upub::model::object::ActiveModel {
+				internal: Unchanged(internal),
+				image: Set(Some(ctx.cloaked(&image))),
+				..Default::default()
+			};
+			model.update(ctx.db()).await?;
 		}
 	}
 
@@ -59,6 +57,7 @@ pub async fn cloak(ctx: upub::Context, post_contents: bool, actors: bool) -> Res
 			.await?;
 
 		while let Some((internal, image, icon)) = stream.try_next().await? {
+			tracing::info!("cloaking user #{internal}");
 			if image.is_none() && icon.is_none() { continue }
 			// TODO can this if/else/else be made nicer??
 			let image = if let Some(img) = image {
@@ -88,5 +87,29 @@ pub async fn cloak(ctx: upub::Context, post_contents: bool, actors: bool) -> Res
 		}
 	}
 
+	if post_contents {
+		let mut stream = upub::model::object::Entity::find()
+			.filter(upub::model::object::Column::Content.like("%<img%"))
+			.select_only()
+			.select_column(upub::model::object::Column::Internal)
+			.select_column(upub::model::object::Column::Content)
+			.into_tuple::<(i64, String)>()
+			.stream(ctx.db())
+			.await?;
+
+		while let Some((internal, content)) = stream.try_next().await? {
+			let sanitized = ctx.sanitize(&content);
+			if sanitized != content {
+				tracing::info!("sanitizing object #{internal}");
+				let model = upub::model::object::ActiveModel {
+					internal: Unchanged(internal),
+					content: Set(Some(sanitized)),
+					..Default::default()
+				};
+				model.update(ctx.db()).await?;
+			}
+		}
+	}
+
 	Ok(())
 }
diff --git a/upub/cli/src/lib.rs b/upub/cli/src/lib.rs
index 84f3ed4..d96c29a 100644
--- a/upub/cli/src/lib.rs
+++ b/upub/cli/src/lib.rs
@@ -116,13 +116,17 @@ pub enum CliCommand {
 
 	/// replaces all attachment urls with proxied local versions (only useful for old instances)
 	Cloak {
-		/// also replace urls inside post contents
+		/// also cloak objects image urls
 		#[arg(long, default_value_t = false)]
-		post_contents: bool,
+		objects: bool,
 
 		/// also cloak actor images
 		#[arg(long, default_value_t = false)]
 		actors: bool,
+
+		/// also replace urls inside post contents
+		#[arg(long, default_value_t = false)]
+		contents: bool,
 	},
 }
 
@@ -145,7 +149,7 @@ pub async fn run(ctx: upub::Context, command: CliCommand) -> Result<(), Box<dyn
 			Ok(nuke(ctx, for_real, delete_objects).await?),
 		CliCommand::Thread { } =>
 			Ok(thread(ctx).await?),
-		CliCommand::Cloak { post_contents, actors } =>
-			Ok(cloak(ctx, post_contents, actors).await?),
+		CliCommand::Cloak { objects, actors, contents } =>
+			Ok(cloak(ctx, contents, objects, actors).await?),
 	}
 }
diff --git a/upub/core/src/traits/normalize.rs b/upub/core/src/traits/normalize.rs
index 07278e8..e214230 100644
--- a/upub/core/src/traits/normalize.rs
+++ b/upub/core/src/traits/normalize.rs
@@ -31,6 +31,10 @@ impl Normalizer for crate::Context {
 			object_model.content = Some(self.sanitize(&content));
 		}
 
+		if let Some(image) = object_model.image {
+			object_model.image = Some(self.cloaked(&image));
+		}
+
 		// fix context for remote posts
 		// > if any link is broken or we get rate limited, the whole insertion fails which is
 		// > kind of dumb. there should be a job system so this can be done in waves. or maybe there's