localsend/transfer/
upload.rs

1use std::collections::HashMap;
2use std::net::SocketAddr;
3use std::path::PathBuf;
4use std::sync::Arc;
5
6use axum::body::Bytes;
7use axum::extract::{ConnectInfo, Query};
8use axum::Extension;
9use axum::{response::IntoResponse, Json};
10use axum::http::StatusCode;
11
12use native_dialog::MessageDialog;
13use serde::{Deserialize, Serialize};
14use tokio::sync::Mutex;
15use uuid::Uuid;
16use crate::error::{LocalSendError, Result};
17use crate::transfer::session::{Session, SessionStatus};
18use crate::{models::{device::DeviceInfo, file::FileMetadata}, Client};
19
20#[derive(Debug, Serialize, Deserialize)]
21#[serde(rename_all = "camelCase")]
22pub struct PrepareUploadResponse {
23    pub session_id: String,
24    pub files: HashMap<String, String>,
25}
26
27#[derive(Debug, Serialize, Deserialize)]
28#[serde(rename_all = "camelCase")]
29pub struct PrepareUploadRequest {
30    pub info: DeviceInfo,
31    pub files: HashMap<String, FileMetadata>,
32}
33
34impl Client {
35    pub async fn prepare_upload(&self, peer: String, files: HashMap<String, FileMetadata>) -> Result<PrepareUploadResponse> {
36        if !self.peers.lock().await.contains_key(&peer) {
37            return Err(LocalSendError::PeerNotFound);
38        }
39
40        let peer = self.peers.lock().await.get(&peer).unwrap().clone();
41        println!("Peer: {:?}", peer);
42
43        let response = self
44            .http_client
45            .post(&format!("{}://{}/api/localsend/v2/prepare-upload", peer.1.protocol, peer.0.clone()))
46            .json(&PrepareUploadRequest {
47                info: self.device.clone(),
48                files: files.clone(),
49            })
50            .send()
51            .await?;
52
53        println!("Response: {:?}", response);
54
55        let response: PrepareUploadResponse = response.json().await?;
56
57        let session = Session {
58            session_id: response.session_id.clone(),
59            files,
60            file_tokens: response.files.clone(),
61            receiver: peer.1,
62            sender: self.device.clone(),
63            status: SessionStatus::Active,
64            addr: peer.0,
65        };
66
67        self.sessions.lock().await.insert(response.session_id.clone(), session);
68
69        Ok(response)
70    }
71
72    pub async fn upload(&self, session_id: String, file_id: String, token: String, body: Bytes) -> Result<()> {
73        let sessions = self.sessions.lock().await;
74        let session = sessions.get(&session_id).unwrap();
75
76        if session.status != SessionStatus::Active {
77            return Err(LocalSendError::SessionInactive);
78        }
79
80        if session.file_tokens.get(&file_id) != Some(&token) {
81            return Err(LocalSendError::InvalidToken);
82        }
83
84        let request = self
85            .http_client
86            .post(&format!("{}://{}/api/localsend/v2/upload?sessionId={}&fileId={}&token={}", session.receiver.protocol, session.addr, session_id, file_id, token))
87            //.post(&format!("https://webhook.site/2f23a529-b687-4375-ad5f-54906ab26ac7?session_id={}&file_id={}&token={}", session_id, file_id, token))
88            .body(body);
89
90        println!("Uploading file: {:?}", request);
91        let response = request.send().await?;
92
93        if response.status() != 200 {
94            println!("Upload failed: {:?}", response);
95            return Err(LocalSendError::UploadFailed);
96        }
97
98        Ok(())
99    }
100
101    pub async fn send_file(&self, peer: String, file_path: PathBuf) -> Result<()> {
102        // Generate file metadata
103        let file_metadata = FileMetadata::from_path(&file_path)?;
104
105        // Prepare files map
106        let mut files = HashMap::new();
107        files.insert(file_metadata.id.clone(), file_metadata.clone());
108
109        // Prepare upload
110        let prepare_response = self.prepare_upload(peer, files).await?;
111
112        // Get file token
113        let token = prepare_response.files.get(&file_metadata.id)
114            .ok_or(LocalSendError::InvalidToken)?;
115
116        // Read file contents
117        let file_contents = tokio::fs::read(&file_path).await?;
118        let bytes = Bytes::from(file_contents);
119
120        // Upload file
121        self.upload(
122            prepare_response.session_id,
123            file_metadata.id,
124            token.clone(),
125            bytes
126        ).await?;
127
128        Ok(())
129    }
130
131    pub async fn cancel_upload(&self, session_id: String) -> Result<()> {
132        let sessions = self.sessions.lock().await;
133        let session = sessions.get(&session_id).unwrap();
134
135        let request = self
136            .http_client
137            .post(&format!("{}://{}/api/localsend/v2/cancel?sessionId={}", session.receiver.protocol, session.addr, session_id))
138            .send()
139            .await?;
140
141        if request.status() != 200 {
142            return Err(LocalSendError::CancelFailed);
143        }
144
145        Ok(())
146    }
147}
148
149pub async fn register_prepare_upload(
150    Extension(client): Extension<DeviceInfo>,
151    Extension(sessions): Extension<Arc<Mutex<HashMap<String, Session>>>>,
152    ConnectInfo(addr): ConnectInfo<SocketAddr>,
153    Json(req): Json<PrepareUploadRequest>,
154) -> impl IntoResponse {
155    println!("Received upload request from alias: {}", req.info.alias);
156
157    let result = MessageDialog::new()
158        .set_title(&req.info.alias)
159        .set_text("Do you want to receive files from this device?")
160        .show_confirm()
161        .unwrap();
162
163    if result {
164        let session_id = Uuid::new_v4().to_string();
165
166        let file_tokens: HashMap<String, String> = req.files.iter()
167                                                      .map(|(id, _)| (id.clone(), Uuid::new_v4().to_string())) // Replace with actual token logic
168                                                      .collect();
169
170        let session = Session {
171            session_id: session_id.clone(),
172            files: req.files.clone(),
173            file_tokens: file_tokens.clone(),
174            receiver: client.clone(),
175            sender: req.info.clone(),
176            status: SessionStatus::Active,
177            addr,
178        };
179
180        sessions.lock().await.insert(session_id.clone(), session);
181
182        return (StatusCode::OK,
183                Json(PrepareUploadResponse {
184                    session_id,
185                    files: file_tokens,
186                })).into_response();
187    } else {
188        return StatusCode::FORBIDDEN.into_response();
189    }
190}
191
192pub async fn register_upload(
193    Query(params): Query<UploadParams>,
194    Extension(sessions): Extension<Arc<Mutex<HashMap<String, Session>>>>,
195    Extension(download_dir): Extension<String>,
196    body: Bytes,
197) -> impl IntoResponse {
198    // Extract query parameters
199    let session_id = &params.session_id;
200    let file_id = &params.file_id;
201    let token = &params.token;
202
203    // Get session and validate
204    let mut sessions_lock = sessions.lock().await;
205    let session = match sessions_lock.get_mut(session_id) {
206        Some(session) => session,
207        None => return StatusCode::BAD_REQUEST.into_response(),
208    };
209
210
211    if session.status != SessionStatus::Active {
212        return StatusCode::BAD_REQUEST.into_response()
213    }
214
215    // Validate token
216    if session.file_tokens.get(file_id) != Some(&token.to_string()) {
217        return StatusCode::FORBIDDEN.into_response();
218    }
219
220    // Get file metadata
221    let file_metadata = match session.files.get(file_id) {
222        Some(metadata) => metadata,
223        None => return (
224            StatusCode::INTERNAL_SERVER_ERROR,
225            "File not found".to_string(),
226        )
227            .into_response(),
228    };
229
230    // Create directory if it doesn't exist
231    if let Err(e) = tokio::fs::create_dir_all(&*download_dir).await {
232        return (
233            StatusCode::INTERNAL_SERVER_ERROR,
234            format!("Failed to create directory: {}", e),
235        )
236            .into_response();
237    }
238
239    // Create file path
240    let file_path = format!("{}/{}", download_dir, file_metadata.file_name);
241
242    // Write file
243    if let Err(e) = tokio::fs::write(&file_path, body).await {
244        return (
245            StatusCode::INTERNAL_SERVER_ERROR,
246            format!("Failed to write file: {}", e),
247        )
248            .into_response();
249    }
250
251    StatusCode::OK.into_response()
252}
253
254// Query parameters struct
255#[derive(Deserialize)]
256#[serde(rename_all = "camelCase")]
257pub struct UploadParams {
258    session_id: String,
259    file_id: String,
260    token: String,
261}
262
263pub async fn register_cancel(
264    Query(params): Query<CancelParams>,
265    Extension(sessions): Extension<Arc<Mutex<HashMap<String, Session>>>>,
266) -> impl IntoResponse {
267    let mut sessions_lock = sessions.lock().await;
268    let session = match sessions_lock.get_mut(&params.session_id) {
269        Some(session) => session,
270        None => return StatusCode::BAD_REQUEST.into_response(),
271    };
272    session.status = SessionStatus::Cancelled;
273    StatusCode::OK.into_response()
274}
275
276// Cancel parameters struct
277#[derive(Deserialize)]
278#[serde(rename_all = "camelCase")]
279pub struct CancelParams {
280    session_id: String,
281}
localsend/transfer/upload.rs

localsend/transfer/
upload.rs
