Skip to content

Commit

Permalink
Merge pull request #16 from subhashdasyam/subhashdasyam-handle-response
Browse files Browse the repository at this point in the history
Update mod.rs add handle the response type properly
  • Loading branch information
evilsocket committed Jul 6, 2024
2 parents f19b901 + aa8f63f commit 9eda43f
Showing 1 changed file with 36 additions and 14 deletions.
50 changes: 36 additions & 14 deletions src/agent/namespaces/http/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -147,59 +147,81 @@ impl Action for Request {
let method = reqwest::Method::from_str(attrs.get("method").unwrap())?;
let parsed = Self::create_url_from(&state, payload.clone()).await?;
let query_str = parsed.query().unwrap_or("").to_string();

let mut client = reqwest::Client::new().request(method.clone(), parsed.clone());
let lock = state.lock().await;
let headers = lock.get_storage("http-headers")?;

for (key, value) in headers.iter() {
client = client.header(key, &value.data);
}

// if there're parameters and we're not in GET, set them as the body
if !query_str.is_empty() {
if !matches!(method, reqwest::Method::GET) {
client = client.body(query_str);
}
}

log::info!(
"{}.{} {} ...",
"http".bold(),
method.to_string().yellow(),
parsed.to_string(),
);

let start = Instant::now();
let res = client.send().await?;
let elaps = start.elapsed();

return if res.status().is_success() {
let reason = res.status().canonical_reason().unwrap();
let mut resp = format!("{} {}\n", res.status().as_u16(), &reason);

for (key, val) in res.headers() {
resp += &format!("{}: {}\n", key, val.to_str().unwrap());
}

// TODO: handle the response type properly
// Handle the response type properly
resp += "\n\n";
resp += &res.text().await?;

match res.headers().get("content-type") {
Some(content_type) if content_type.to_str().unwrap().starts_with("application/") || content_type.to_str().unwrap().starts_with("text/") => {
if content_type.to_str().unwrap() == "application/octet-stream" {
// Download the first few bytes to determine if it's binary or text
let partial_content = res.bytes().await?;
let is_binary = partial_content.iter().any(|&byte| byte == 0 || (byte < 32 && byte != 9 && byte != 10 && byte != 13));

if is_binary {
log::warn!("Ignoring binary data with content type: application/octet-stream");
} else {
resp += std::str::from_utf8(&partial_content).unwrap_or("");
}
} else {
resp += &res.text().await?;
}
},
Some(content_type) => {
log::warn!("Ignoring non-textual content type: {}", content_type.to_str().unwrap());
},
None => {
log::warn!("No content type specified in the response");
}
}

log::info!(
" {} {} -> {} bytes",
reason.green(),
format!("({:?})", elaps).dimmed(),
resp.len()
);

Ok(Some(resp))
} else {
let reason = res.status().canonical_reason().unwrap();
let resp = format!("{} {}", res.status().as_u16(), &reason);

log::error!(" {} {}", reason.red(), format!("({:?})", elaps).dimmed(),);

Err(anyhow!(resp))
};
}
Expand Down

0 comments on commit 9eda43f

Please sign in to comment.