1use crate::*;
2use axum::extract::{
3 FromRequestParts,
4 ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade},
5};
6use futures::{SinkExt, StreamExt};
7use tokio::io::{AsyncReadExt, AsyncWriteExt};
8
9pub fn visible(me: &Value, owner: &Value) -> bool {
10 array(&me["groups"])
11 .iter()
12 .any(|group| group == "infra-admin")
13 || (owner.as_str().is_some_and(|id| !id.is_empty()) && owner == &me["id"])
14}
15
16pub async fn authorize(me: &Value, name: &str) -> Result<()> {
17 let owner = host::call(json!({"operation":"vm.owner","payload":{"name":name}})).await?;
18 if visible(me, &owner) {
19 Ok(())
20 } else {
21 Err(Error::new(
22 404,
23 "This VM isn't in your account. Open your VM home.",
24 ))
25 }
26}
27
28pub async fn console(request: Request, name: &str, serial: bool) -> Result<Response> {
29 // Browsers send Origin on a WebSocket handshake. Reject cross-site control.
30 same_origin(&request)?;
31 let (mut parts, _) = request.into_parts();
32 let upgrade = WebSocketUpgrade::from_request_parts(&mut parts, &())
33 .await
34 .map_err(|_| Error::new(400, "Open the screen from the VM page."))?
35 .max_message_size(1024 * 1024)
36 .max_frame_size(1024 * 1024);
37 let name = name.to_owned();
38 Ok(upgrade
39 .protocols(["binary"])
40 .on_upgrade(move |mut socket| async move {
41 let stream = async {
42 if serial {
43 return host::open(json!({"operation":"vm.serial","payload":{"name":name}}))
44 .await
45 .map(|(_, stream)| stream);
46 }
47 let mut stream = tokio::net::UnixStream::connect(env(
48 "STUDIO_VM_SCREEN_SOCKET",
49 "/run/studio-vm-screen/screen.sock",
50 ))
51 .await?;
52 stream
53 .write_all(format!("{}\n", json!({"name":name})).as_bytes())
54 .await?;
55 Ok(stream)
56 }
57 .await;
58 match stream {
59 Ok(stream) => bridge(socket, stream).await,
60 Err(error) => {
61 let reason = if error.message.len() <= 120 {
62 error.message
63 } else {
64 "The host operation couldn't finish. Check its logs, then retry.".into()
65 };
66 let _ = socket
67 .send(Message::Close(Some(CloseFrame {
68 code: 1011,
69 reason: reason.into(),
70 })))
71 .await;
72 }
73 }
74 })
75 .into_response())
76}
77
78fn same_origin(request: &Request) -> Result<()> {
79 let supplied = request
80 .headers()
81 .get("origin")
82 .and_then(|v| v.to_str().ok())
83 .and_then(|v| url::Url::parse(v).ok());
84 // The ingress strips forwarded headers and supplies the original Host.
85 let host = request
86 .headers()
87 .get("host")
88 .and_then(|v| v.to_str().ok())
89 .unwrap_or("");
90 if supplied.as_ref().is_none_or(|origin| {
91 let authority = match origin.port() {
92 Some(port) => format!("{}:{port}", origin.host_str().unwrap_or("")),
93 None => origin.host_str().unwrap_or("").to_owned(),
94 };
95 authority != host || !matches!(origin.scheme(), "http" | "https")
96 }) {
97 return Err(Error::new(403, "Open the VM page on this site to connect."));
98 }
99 Ok(())
100}
101
102async fn bridge(socket: WebSocket, stream: tokio::net::UnixStream) {
103 let (mut send, mut receive) = socket.split();
104 let (mut read, mut write) = stream.into_split();
105 let upstream = async {
106 while let Some(Ok(message)) = receive.next().await {
107 match message {
108 Message::Binary(data) => write.write_all(&data).await?,
109 Message::Text(data) => write.write_all(data.as_bytes()).await?,
110 Message::Close(_) => break,
111 Message::Ping(_) | Message::Pong(_) => {}
112 }
113 }
114 Ok::<_, std::io::Error>(())
115 };
116 let downstream = async {
117 let mut buffer = vec![0; 65536];
118 loop {
119 let count = read.read(&mut buffer).await?;
120 if count == 0 {
121 break;
122 }
123 send.send(Message::Binary(Bytes::copy_from_slice(&buffer[..count])))
124 .await
125 .map_err(std::io::Error::other)?;
126 }
127 Ok::<_, std::io::Error>(())
128 };
129 tokio::select! { _ = upstream => {}, _ = downstream => {} }
130 // Dropping either half disconnects the host tunnel and releases its slot.
131}
132
133pub async fn upload(
134 app: Arc<App>,
135 request: Request,
136 query: &HashMap<String, String>,
137) -> Result<Response> {
138 same_origin(&request)?;
139 let volume = query
140 .get("name")
141 .ok_or_else(|| Error::new(400, "Choose an ISO file to upload."))?;
142 let size = request
143 .headers()
144 .get("content-length")
145 .and_then(|v| v.to_str().ok())
146 .and_then(|v| v.parse::<u64>().ok())
147 .filter(|v| (32768..=32 * 2u64.pow(30)).contains(v))
148 .ok_or_else(|| Error::new(400, "Choose an ISO between 32 KiB and 32 GiB."))?;
149 let (_, mut stream) =
150 host::open(json!({"operation":"vm.upload","payload":{"volume":volume,"size":size}}))
151 .await?;
152 let mut incoming = request.into_body().into_data_stream();
153 let mut remaining = size;
154 while let Some(data) = tokio::time::timeout(Duration::from_secs(120), incoming.next())
155 .await
156 .map_err(|_| Error::new(408, "The upload paused too long. Upload the ISO again."))?
157 {
158 let data = data?;
159 if data.len() as u64 > remaining {
160 return Err(Error::new(
161 400,
162 "The ISO size changed. Upload the file again.",
163 ));
164 }
165 tokio::time::timeout(Duration::from_secs(120), stream.write_all(&data))
166 .await
167 .map_err(|_| {
168 Error::new(
169 504,
170 "The host stopped receiving the ISO. Try uploading again.",
171 )
172 })??;
173 remaining -= data.len() as u64;
174 }
175 if remaining != 0 {
176 return Err(Error::new(
177 400,
178 "The upload was interrupted. Upload the ISO again.",
179 ));
180 }
181 let value = tokio::time::timeout(Duration::from_secs(120), host::response(&mut stream))
182 .await
183 .map_err(|_| {
184 Error::new(
185 504,
186 "The host is still saving the ISO. Refresh the library before retrying.",
187 )
188 })??;
189 app.cache.invalidate("vms:library");
190 Ok(Document::new(value).response())
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use axum::body::Body;
197
198 #[test]
199 fn ownership_follows_account_id_and_effective_admin_role() {
200 let user = json!({"id":"account-1","name":"renamed","groups":["vm"]});
201 assert!(visible(&user, &json!("account-1")));
202 assert!(!visible(&user, &json!("account-2")));
203 assert!(!visible(&user, &Value::Null));
204 assert!(!visible(
205 &json!({"id":"snow","groups":["vm"]}),
206 &json!("other")
207 ));
208 assert!(visible(
209 &json!({"id":"snow","groups":["infra-admin"]}),
210 &Value::Null
211 ));
212 }
213
214 #[test]
215 fn screen_rejects_missing_and_cross_site_origins() {
216 for (origin, allowed) in [
217 (None, false),
218 (Some("https://other.test"), false),
219 (Some("null"), false),
220 (Some("https://vm.test"), true),
221 (Some("http://vm.test:5178"), false),
222 ] {
223 let mut request = Request::builder().header("host", "vm.test");
224 if let Some(origin) = origin {
225 request = request.header("origin", origin);
226 }
227 assert_eq!(
228 same_origin(&request.body(Body::empty()).unwrap()).is_ok(),
229 allowed
230 );
231 }
232 let request = Request::builder()
233 .header("host", "127.0.0.1:5178")
234 .header("origin", "http://127.0.0.1:5178")
235 .body(Body::empty())
236 .unwrap();
237 assert!(same_origin(&request).is_ok());
238 }
239}