From bd7bc81107a958304737613e1a6472ef1d07bea3 Mon Sep 17 00:00:00 2001 From: Boris-Chengbiao Zhou Date: Thu, 2 Mar 2017 19:20:57 +0100 Subject: [PATCH] Retry network related stuff if failing --- src/main.rs | 114 +++++++++++++++++++++++++++------------------------- 1 file changed, 59 insertions(+), 55 deletions(-) diff --git a/src/main.rs b/src/main.rs index c8321ce..d9da841 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,7 +14,7 @@ use std::collections::HashMap; use std::time::{Instant, Duration}; use json::JsonValue; -use reqwest::{Client, Response, RedirectPolicy}; +use reqwest::{Client, Response, RedirectPolicy, RequestBuilder}; use reqwest::header::{Cookie, Charset, SetCookie}; use reqwest::header::{ContentDisposition, DispositionType, DispositionParam}; use zip::ZipArchive; @@ -34,7 +34,8 @@ mod catalogs { // TODO: use clap for proper CLI // TODO: replace unwraps with descriptive expects -// TODO: retry some things if sensible + +const MAX_RETRIES: u8 = 10; lazy_static! { static ref CLIENT: Client = { @@ -107,10 +108,10 @@ fn get_auth(username: &str, password: &str) -> String { form_data.insert("UserName", username); form_data.insert("Password", password); - let res = CLIENT.post("https://streams.tum.de/Mediasite/Login") - .form(&form_data) - .send() - .unwrap(); + let res = try_to_get_response(|client| { + client.post("https://streams.tum.de/Mediasite/Login") + .form(&form_data) + }); let set_cookie: &SetCookie = res.headers().get().unwrap(); let cookie = cookie::Cookie::parse(set_cookie.0[0].to_string()).unwrap(); @@ -139,11 +140,9 @@ fn get_catalog_id(name: &str, auth: &str) -> String { println!("Fetching catalog id!"); let url = format!("https://streams.tum.de/Mediasite/Catalog/catalogs/{}", name); - let res = CLIENT.get(&url) - .header(construct_cookie(auth)) - .send(); + let mut res = try_to_get_response(|client| client.get(&url).header(construct_cookie(auth))); let mut string = String::new(); - res.unwrap().read_to_string(&mut string).unwrap(); + res.read_to_string(&mut string).unwrap(); let prefix = "CatalogId: '"; let idx = string.find(prefix).unwrap(); @@ -155,50 +154,28 @@ fn get_catalog_id(name: &str, auth: &str) -> String { } fn download_presentation(presentation: &Presentation, out_dir: &Path, auth: &str) { - fn try_download(presentation: &Presentation, auth: &str) -> Option<(Response, String)> { - let response = CLIENT.get(&presentation.download_url) - .header(construct_cookie(auth)) - .send() - .unwrap(); + let response = + try_to_get_valid_response(|client| { + client.get(&presentation.download_url) + .header(construct_cookie(auth)) + }, + |resp| resp.headers().get::().is_some()); - let filename = if let Some(content_disposition) = - response.headers().get::() { - assert_eq!(content_disposition.disposition, DispositionType::Attachment); - let mut name = None; - for param in &content_disposition.parameters { - match *param { - DispositionParam::Filename(ref charset, _, ref vec) => { - assert_eq!(&Charset::Ext("UTF-8".to_string()), charset); - name = Some(String::from_utf8(vec.to_vec()).unwrap()); - } - _ => continue, + let filename = { + let content_disposition = response.headers().get::().unwrap(); + assert_eq!(content_disposition.disposition, DispositionType::Attachment); + let mut name = None; + for param in &content_disposition.parameters { + match *param { + DispositionParam::Filename(ref charset, _, ref vec) => { + assert_eq!(&Charset::Ext("UTF-8".to_string()), charset); + name = Some(String::from_utf8(vec.to_vec()).unwrap()); } + _ => continue, } - name.expect("Missing filename in ContentDisposition!") - } else { - return None; - }; - - Some((response, filename)) - } - - let mut tries = 0; - let mut tuple = None; - while tries < 8 { - if tries > 0 { - println!("Failed to start download! Retrying!"); } - let try = try_download(presentation, auth); - if try.is_some() { - tuple = try; - break; - } - tries += 1; - } - if tuple.is_none() { - panic!("Failed to get a proper response while trying to download presentation!"); - } - let (response, filename) = tuple.unwrap(); + name.expect("Missing filename in ContentDisposition!") + }; let mut path = PathBuf::from(out_dir); path.push(fix_filename(&filename)); @@ -274,12 +251,13 @@ fn get_json(catalog_id: &str, auth: &str) -> String { data.insert("CurrentFolderId", catalog_id); data.insert("ItemsPerPage", "500"); - let res = CLIENT.post("https://streams.tum.de/Mediasite/Catalog/Data/GetPresentationsForFolder") - .header(construct_cookie(auth)) - .json(&data) - .send(); + let mut res = try_to_get_response(|client| { + client.post("https://streams.tum.de/Mediasite/Catalog/Data/GetPresentationsForFolder") + .header(construct_cookie(auth)) + .json(&data) + }); let mut string = String::new(); - res.unwrap().read_to_string(&mut string).unwrap(); + res.read_to_string(&mut string).unwrap(); string } @@ -293,3 +271,29 @@ fn zip_is_valid(path: &Path) -> bool { let file = File::open(path).unwrap(); ZipArchive::new(file).is_ok() } + +fn try_to_get_response(f: F) -> Response + where F: Fn(&Client) -> RequestBuilder +{ + try_to_get_valid_response(f, |_| true) +} + +fn try_to_get_valid_response(f1: F1, f2: F2) -> Response + where F1: Fn(&Client) -> RequestBuilder, + F2: Fn(&Response) -> bool +{ + let mut retries = 0; + while retries <= MAX_RETRIES { + if retries > 0 { + println!("Retrying request!"); + } + let response = f1(&*CLIENT).send(); + if let Ok(response) = response { + if f2(&response) { + return response; + } + } + retries += 1; + } + panic!("Reached maximum amout of retries!") +}