grain.social is a photo sharing platform built on atproto.
grain.social
1use anyhow::{Result, anyhow};
2use fantoccini::ClientBuilder;
3use std::net::TcpStream;
4use std::process::{Child, Command};
5use std::{thread, time};
6use tracing::info;
7
8pub async fn capture_screenshot(preview_url: &str) -> Result<Vec<u8>> {
9 capture_screenshot_with_size(preview_url, "1200,769").await
10}
11
12async fn capture_screenshot_with_size(preview_url: &str, window_size: &str) -> Result<Vec<u8>> {
13 info!("Starting screenshot capture for: {}", preview_url);
14
15 // Check if ChromeDriver is running on port 9515
16 let chromedriver_addr = "127.0.0.1:9515";
17 let mut chromedriver_child: Option<Child> = None;
18 let chromedriver_running = TcpStream::connect(chromedriver_addr).is_ok();
19 if !chromedriver_running {
20 let chromedriver_path = std::env::var("CHROMEDRIVER_PATH")
21 .unwrap_or_else(|_| "/usr/bin/chromedriver".to_string());
22 info!("Starting ChromeDriver at {}...", chromedriver_path);
23 let child = Command::new(chromedriver_path)
24 .arg("--port=9515")
25 .spawn()
26 .map_err(|e| anyhow!("Failed to start ChromeDriver: {}", e))?;
27 chromedriver_child = Some(child);
28 // Wait for ChromeDriver to become available
29 let max_tries = 10;
30 let delay = time::Duration::from_millis(300);
31 let mut started = false;
32 for _ in 0..max_tries {
33 if TcpStream::connect(chromedriver_addr).is_ok() {
34 started = true;
35 break;
36 }
37 thread::sleep(delay);
38 }
39 if !started {
40 return Err(anyhow!("ChromeDriver did not start on port 9515"));
41 }
42 }
43
44 let mut caps = serde_json::map::Map::new();
45 let opts = serde_json::json!({
46 "binary": std::env::var("CHROME_PATH")
47 .unwrap_or_else(|_| "/usr/bin/chromium".to_string()),
48 "args": [
49 "--headless",
50 "--no-sandbox",
51 "--disable-gpu",
52 "--disable-dev-shm-usage",
53 &format!("--window-size={}", window_size),
54 "--font-render-hinting=medium",
55 "--enable-font-antialiasing",
56 ],
57 });
58 caps.insert("goog:chromeOptions".to_string(), opts);
59
60 // Create WebDriver client
61 let client = ClientBuilder::native()
62 .capabilities(caps)
63 .connect("http://localhost:9515")
64 .await
65 .map_err(|e| anyhow!("Failed to connect to ChromeDriver: {}", e))?;
66
67 let result = async {
68 info!("Navigating to URL: {}", preview_url);
69 client
70 .goto(preview_url)
71 .await
72 .map_err(|e| anyhow!("Failed to navigate to {}: {}", preview_url, e))?;
73
74 // Wait for body to be present
75 use fantoccini::wd::Locator;
76 client.wait().for_element(Locator::Css("body")).await?;
77
78 // Wait for all web fonts to be loaded
79 client
80 .execute(
81 "return document.fonts.status === 'loaded' ? true : await document.fonts.ready.then(() => true);",
82 vec![],
83 )
84 .await
85 .map_err(|e| anyhow!("Failed to wait for fonts: {}", e))?;
86
87 // Wait for screenshot ready signal (for dynamic content) or fallback to delay
88 info!("Waiting for screenshot ready signal...");
89 let ready = client
90 .execute(
91 r#"
92 return new Promise((resolve) => {
93 // Check if already ready
94 if (document.body.dataset.screenshotReady === 'true') {
95 resolve(true);
96 return;
97 }
98
99 // Wait for the ready signal with timeout
100 const timeout = setTimeout(() => {
101 console.log('Screenshot ready timeout - proceeding anyway');
102 resolve(false);
103 }, 15000); // 15 second timeout
104
105 document.addEventListener('screenshotReady', () => {
106 clearTimeout(timeout);
107 resolve(true);
108 });
109 });
110 "#,
111 vec![],
112 )
113 .await
114 .map_err(|e| anyhow!("Failed to wait for screenshot ready: {}", e))?;
115
116 if ready.as_bool().unwrap_or(false) {
117 info!("Screenshot ready signal received");
118 } else {
119 info!("Screenshot ready timeout - using fallback delay");
120 tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
121 }
122
123 info!("Taking screenshot...");
124 let screenshot_data = client
125 .screenshot()
126 .await
127 .map_err(|e| anyhow!("Failed to capture screenshot: {}", e))?;
128
129 info!(
130 "Screenshot captured successfully, size: {} bytes",
131 screenshot_data.len()
132 );
133 Ok(screenshot_data)
134 }
135 .await;
136
137 // Clean up
138 client.close().await.ok();
139
140 // If we started ChromeDriver, kill it
141 if let Some(mut child) = chromedriver_child {
142 let _ = child.kill();
143 }
144
145 result
146}