1use crate::*;
2use tokio::{
3 io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
4 sync::{mpsc, oneshot},
5};
6
7pub struct Worker(mpsc::Sender<Call>);
8struct Call {
9 method: String,
10 args: Value,
11 reply: oneshot::Sender<Result<Value>>,
12}
13
14impl Worker {
15 pub fn new() -> Self {
16 let (send, mut receive) = mpsc::channel::<Call>(64);
17 tokio::spawn(async move {
18 let mut child: Option<(
19 tokio::process::Child,
20 tokio::process::ChildStdin,
21 tokio::io::Lines<BufReader<tokio::process::ChildStdout>>,
22 )> = None;
23 let mut id = 0u64;
24 while let Some(call) = receive.recv().await {
25 if call.reply.is_closed() {
26 continue;
27 }
28 let response = tokio::time::timeout(Duration::from_secs(25), async {
29 if child.is_none() {
30 if std::env::var("STUDIO_YT_STATE").is_err() {
31 return Err(Error::new(
32 501,
33 "YouTube isn't connected to this home server. Configure its archive state, then retry.",
34 ));
35 }
36 let python = env("STUDIO_YT_PYTHON", "python3");
37 let script = env("STUDIO_YT_WORKER", "server/youtube-worker.py");
38 let mut command = tokio::process::Command::new(python);
39 for (key, name) in [
40 ("SMTP_HOST", "smtp_host"),
41 ("SMTP_USER", "smtp_user"),
42 ("SMTP_PASS", "smtp_pass"),
43 ] {
44 if let Ok(secret) = host::call(
45 json!({"operation":"deploy.secret.get","service":"ytdl","key":name}),
46 ).await {
47 if let Some(value) = secret.as_str() {
48 command.env(key, value);
49 }
50 }
51 }
52 let mut process = command
53 .arg(script)
54 .stdin(std::process::Stdio::piped())
55 .stdout(std::process::Stdio::piped())
56 .stderr(std::process::Stdio::inherit())
57 .kill_on_drop(true)
58 .spawn()?;
59 let input = process.stdin.take().unwrap();
60 let output = BufReader::new(process.stdout.take().unwrap()).lines();
61 child = Some((process, input, output));
62 }
63 id += 1;
64 let (_, input, output) = child.as_mut().unwrap();
65 input
66 .write_all(
67 format!(
68 "{}\n",
69 json!({"id":id,"method":call.method,"args":call.args})
70 )
71 .as_bytes(),
72 )
73 .await?;
74 input.flush().await?;
75 let line = output.next_line().await?.ok_or_else(|| {
76 Error::new(502, "YouTube processing stopped. Retry shortly.")
77 })?;
78 let response: Value = serde_json::from_str(&line)?;
79 if response["id"] != id {
80 return Err(Error::new(
81 502,
82 "YouTube processing stopped. Retry shortly.",
83 ));
84 }
85 if let Some(error) = response["error"].as_str() {
86 return Err(Error::new(
87 if error.contains("read-only") || error.contains("changed since") {
88 409
89 } else {
90 502
91 },
92 error,
93 ));
94 }
95 Ok(response["result"].clone())
96 })
97 .await
98 .unwrap_or_else(|_| {
99 Err(Error::new(
100 504,
101 "YouTube is taking too long. Retry shortly.",
102 ))
103 });
104 if response.as_ref().is_err_and(|e| {
105 e.status == 504
106 || e.status == 500
107 || (e.status == 502 && e.message.contains("stopped"))
108 }) {
109 child = None;
110 }
111 let _ = call.reply.send(response);
112 }
113 });
114 Self(send)
115 }
116 pub async fn call(&self, method: &str, args: Value) -> Result<Value> {
117 let (reply, result) = oneshot::channel();
118 tokio::time::timeout(Duration::from_secs(30), async {
119 self.0
120 .send(Call {
121 method: method.into(),
122 args,
123 reply,
124 })
125 .await
126 .map_err(|_| Error::new(502, "YouTube processing stopped. Retry shortly."))?;
127 result
128 .await
129 .map_err(|_| Error::new(502, "YouTube processing stopped. Retry shortly."))?
130 })
131 .await
132 .map_err(|_| Error::new(504, "YouTube is taking too long. Retry shortly."))?
133 }
134}
135
136pub async fn route(
137 app: Arc<App>,
138 method: &Method,
139 parts: &[&str],
140 body: Value,
141) -> Result<Response> {
142 if method == Method::GET {
143 let call = match parts {
144 [] => "snapshot",
145 ["channels"] => "channels",
146 ["configs"] => "configs",
147 ["library"] => "library",
148 _ => return Err(Error::new(404, "Not Found")),
149 };
150 let state = app.clone();
151 let document = app
152 .cache
153 .get(
154 format!("youtube:{call}"),
155 Duration::from_secs(if call == "library" { 30 } else { 2 }),
156 move || async move {
157 if call == "library" {
158 let _slot = state.heavy.acquire().await?;
159 let python = env("STUDIO_YT_PYTHON", "python3");
160 let script = env("STUDIO_YT_WORKER", "server/youtube-worker.py");
161 return Ok(serde_json::from_slice(
162 &command(&python, &[&script, "--library"], None).await?,
163 )?);
164 }
165 let worker = state.youtube.get_or_init(Worker::new);
166 let mut value = worker.call(call, json!({})).await?;
167 if call == "snapshot" {
168 let archive = core::read_json(
169 std::path::Path::new(&env(
170 "STUDIO_YT_ARCHIVE_STATUS",
171 "/srv/prod/ytdl-sub/config/archive-status.json",
172 )),
173 Value::Null,
174 )
175 .await
176 .unwrap_or(Value::Null);
177 if !archive.is_null() {
178 value["archive"] = archive;
179 }
180 for job in value["jobs"].as_array_mut().into_iter().flatten() {
181 job["folder"] = apps::file_link(string(&job["folder"]), false);
182 }
183 }
184 Ok(value)
185 },
186 )
187 .await?;
188 return Ok(document.response());
189 }
190 let worker = app.youtube.get_or_init(Worker::new);
191 let value = match parts {
192 ["configs", name] if method == Method::PUT => {
193 if !body["original"].is_string() || !body["body"].is_string() {
194 return Err(Error::new(400, "Invalid config."));
195 }
196 worker
197 .call(
198 "saveConfig",
199 json!({"name":name,"original":body["original"],"body":body["body"]}),
200 )
201 .await?;
202 Value::Null
203 }
204 ["library", "rename"] if method == Method::POST => {
205 if string(&body["path"]).is_empty() || string(&body["title"]).trim().is_empty() {
206 return Err(Error::new(400, "Enter a title."));
207 }
208 for key in ["season", "episode"] {
209 if !body[key].is_null()
210 && (number(&body[key]) < 1.0 || number(&body[key]).fract() != 0.0)
211 {
212 return Err(Error::new(400, "Invalid episode number."));
213 }
214 }
215 worker.call("rename", body).await?;
216 Value::Null
217 }
218 ["channels"] if method == Method::PUT => {
219 validate_channels(&body)?;
220 worker.call("setChannels", body).await?;
221 Value::Null
222 }
223 ["pending"] if method == Method::POST => {
224 let video = regex::Regex::new(
225 r"^https?://((www|m|music)\.)?(youtube\.com/(watch\?|shorts/|live/)|youtu\.be/)",
226 )
227 .unwrap();
228 if !body["urls"].is_array() || array(&body["urls"]).is_empty() {
229 return Err(Error::new(400, "Paste a YouTube video link."));
230 }
231 for url in array(&body["urls"]) {
232 if !video.is_match(string(url)) {
233 return Err(Error::new(
234 400,
235 format!("{} isn't a YouTube video link.", string(url)),
236 ));
237 }
238 }
239 worker.call("add", body).await?;
240 Value::Null
241 }
242 ["pending", key, "ingest"] if method == Method::POST => {
243 let dest = string(&body["dest"]);
244 if !["independent", "music", "indie"].contains(&dest) {
245 return Err(Error::new(400, "Choose where this video goes."));
246 }
247 if dest == "indie" {
248 if string(&body["show"]).trim().is_empty() || !body["title"].is_string() {
249 return Err(Error::new(400, "Enter a show name."));
250 }
251 for key in ["season", "episode"] {
252 let n = number(&body[key]);
253 if n < 1.0 || n.fract() != 0.0 {
254 return Err(Error::new(400, "Invalid episode number."));
255 }
256 }
257 }
258 worker
259 .call("ingest", json!({"key":key,"choice":body}))
260 .await?;
261 Value::Null
262 }
263 ["pending", key, "skip"] if method == Method::POST => {
264 worker.call("skip", json!({"key":key})).await?;
265 Value::Null
266 }
267 ["jobs", id, "retry"] if method == Method::POST => {
268 worker.call("retry", json!({"id":id})).await?;
269 Value::Null
270 }
271 ["upscaler"] if method == Method::PUT => {
272 if !body["enabled"].is_boolean() {
273 return Err(Error::new(400, "Choose whether upscaling is enabled."));
274 }
275 worker.call("setUpscaler", body).await?;
276 Value::Null
277 }
278 _ => return Err(Error::new(404, "Not Found")),
279 };
280 for call in ["snapshot", "channels", "configs", "library"] {
281 app.cache.invalidate(&format!("youtube:{call}"));
282 }
283 Ok(if value.is_null() {
284 StatusCode::NO_CONTENT.into_response()
285 } else {
286 Document::new(value).response()
287 })
288}
289fn validate_channels(body: &Value) -> Result<()> {
290 let channel = regex::Regex::new(
291 r"^https?://(www\.|m\.)?youtube\.com/(@[\w.-]+|channel/[\w-]+|c/[^/?#]+|user/[^/?#]+)/?$",
292 )
293 .unwrap();
294 let date_pattern = regex::Regex::new(r"^\d{8}$").unwrap();
295 for key in ["notify", "archive"] {
296 if !body[key].is_array() {
297 return Err(Error::new(400, "Enter channel lists."));
298 }
299 let mut names = std::collections::HashSet::new();
300 for item in array(&body[key]) {
301 let name = string(&item["name"]).trim();
302 if name.is_empty() {
303 return Err(Error::new(400, "Enter a name for the channel."));
304 }
305 if !names.insert(name) {
306 return Err(Error::new(
307 400,
308 format!("Two channels are named {name}. Rename one."),
309 ));
310 }
311 if !channel.is_match(string(&item["url"]).trim()) {
312 return Err(Error::new(
313 400,
314 "That isn't a YouTube channel link. Use one like https://www.youtube.com/@handle.",
315 ));
316 }
317 if key == "archive" {
318 let rules = &item["rules"];
319 if !rules.is_object() {
320 return Err(Error::new(400, "Enter channel rules."));
321 }
322 if let Some(date) = rules.get("download_after")
323 && !date_pattern.is_match(string(date))
324 {
325 return Err(Error::new(400, "Pick a date for the backlog."));
326 }
327 for rule in [
328 "title_include_keywords",
329 "title_exclude_keywords",
330 "description_include_keywords",
331 "description_exclude_keywords",
332 ] {
333 if let Some(words) = rules.get(rule)
334 && (!words.is_array()
335 || array(words).iter().any(|w| string(w).trim().is_empty()))
336 {
337 return Err(Error::new(400, "Enter keywords."));
338 }
339 }
340 }
341 }
342 }
343 Ok(())
344}