1
\$\begingroup\$

I am using this code to forward my syslog logs (converted to json using syslog-ng) to Azure Log Analytics on an ARM board (since there is no official client for ARM). It uses the HTTP Data Collector API to submit the logs.

The idea is to read logs from stdin (each line is a JSON object) in separate thread with a buffer and send them when tokio worker thread is ready or the buffer gets full (body can't be more 30MB).

My questions:

  • Is my usage of async Rust correct?
  • Is there a way to improve performance?

src/main.rs

use futures::sync::mpsc::channel;
use futures::{Future, Stream};
use hyper::client::HttpConnector;
use hyper_tls::HttpsConnector;
use native_tls::{Protocol, TlsConnector};
use std::io::BufRead;
use std::{io, thread};

const AZURE_REQUEST_LIMIT: usize = 30 * 1000 * 1000; // 30 MB
const LINE_LENGTH_LIMIT: usize = AZURE_REQUEST_LIMIT - 2; // should fit in '[' + <30 mb of json> + ']'

fn aggregated_logs() -> impl Stream<Item = String, Error = ()> {
    let (mut tx, rx) = channel(0);
    thread::spawn(move || {
        let input = io::stdin();

        let mut lock = input.lock();

        let mut buf = String::new();
        loop {
            let empty = buf.is_empty();

            let c = if empty { '[' } else { ',' };
            buf.push(c);

            let mut read_bytes = lock.read_line(&mut buf).expect("failed to read line");
            if read_bytes == 0 && empty {
                break;
            }

            // remove newline
            read_bytes -= 1;
            buf.pop();

            assert!(
                read_bytes <= LINE_LENGTH_LIMIT,
                "line is too long: {} bytes",
                read_bytes
            );
            let under_limit = buf.len() <= AZURE_REQUEST_LIMIT - 1;

            let mut msg = if under_limit {
                // send the entire buffer
                std::mem::replace(&mut buf, String::new())
            } else {
                // buffer got too big, save next line for next iteration
                let prev_len = buf.len() - read_bytes;
                let mut msg = String::from("[");
                msg.push_str(&buf[prev_len..]);

                buf.truncate(prev_len);
                std::mem::replace(&mut buf, msg)
            };

            msg.push(']');

            while let Err(e) = tx.try_send(msg) {
                debug_assert!(e.is_full());
                msg = e.into_inner();
                if under_limit {
                    buf = msg;
                    buf.pop(); // remove ']'
                    break;
                }
            }
        }
    });
    rx
}

const RESOURCE: &'static str = "/api/logs";
const CONTENT_TYPE: &'static str = "application/json";
const X_MS_DATE: &'static str = "x-ms-date";

type HmacSha256 = hmac::Hmac<sha2::Sha256>;
use hmac::Mac;

fn signature(rfc1123date: &str, len: usize, customer_id: &str, mut hmac: HmacSha256) -> String {
    let string_to_hash = format!(
        "POST\n{}\n{}\n{}:{}\n{}",
        len, CONTENT_TYPE, X_MS_DATE, rfc1123date, RESOURCE
    );
    hmac.input(string_to_hash.as_bytes());
    let result = hmac.result();
    let base64 = base64::encode(&result.code());
    format!("SharedKey {}:{}", customer_id, base64)
}

fn main() {
    let customer_id = std::env::var("CUSTOMER_ID").expect("Missing CUSTOMER_ID");
    let shared_key = std::env::var("SHARED_KEY").expect("Missing SHARED_KEY");
    let log_type = std::env::var("LOG_TYPE").expect("Missing LOG_TYPE");
    let time_generated = std::env::var("TIME_GENERATED_FIELD");

    let decoded_key = base64::decode(&shared_key).unwrap();
    let hmac = HmacSha256::new_varkey(&decoded_key).unwrap();

    let mut http = HttpConnector::new(num_cpus::get());
    http.enforce_http(false);

    let connector = TlsConnector::builder()
        .min_protocol_version(Some(Protocol::Tlsv12))
        .build()
        .unwrap();

    let mut https: HttpsConnector<_> = (http, connector).into();
    https.https_only(true);
    let client = hyper::Client::builder().build::<_, hyper::Body>(https);
    let uri = format!(
        "https://{}.ods.opinsights.azure.com{}?api-version=2016-04-01",
        customer_id, RESOURCE
    );

    tokio::run(
        aggregated_logs()
            .then(move |data| {
                let data = data.unwrap();
                let rfc1123date = chrono::Utc::now()
                    .format("%a, %d %b %Y %H:%M:%S GMT")
                    .to_string();

                let len = data.len();
                let signature = signature(&rfc1123date, len, &customer_id, hmac.clone());

                let body: hyper::Body = data.into();

                let mut builder = hyper::Request::post(&uri);
                builder
                    .header("Content-Type", CONTENT_TYPE)
                    .header("Content-Length", len)
                    .header("Authorization", signature)
                    .header("Log-Type", log_type.as_str())
                    .header(X_MS_DATE, rfc1123date);

                if let Ok(time_generated) = &time_generated {
                    builder.header("time-generated-field", time_generated.as_str());
                }
                let request = builder.body(body).unwrap();
                client.request(request).and_then(
                    |response| -> Box<dyn Future<Item = (), Error = hyper::Error> + Send> {
                        let status = response.status();
                        if status.is_success() {
                            Box::new(futures::finished(()))
                        } else {
                            Box::new(response.into_body().concat2().and_then(move |chunk| {
                                let message = String::from_utf8(chunk.to_vec()).unwrap();
                                eprintln!("{} {}", status, message);
                                Ok(())
                            }))
                        }
                    },
                )
            })
            .for_each(|_| Ok(()))
            .then(|e| {
                e.unwrap();
                Ok(())
            }),
    )
}

Cargo.toml:

[package]
name = "log_analytics"
version = "0.1.0"
edition = "2018"

[dependencies]
tokio = "0.1.13"
hyper = "0.12.19"
futures = "0.1.25"
hmac = "0.7.0"
sha2 = "0.8.0"
hyper-tls = "0.3.1"
chrono = "0.4.6"
base64 = "0.10.0"
native-tls = "0.2.2"
num_cpus = "1.9.0"
\$\endgroup\$
  • 1
    \$\begingroup\$ let under_limit = buf.len() <= AZURE_REQUEST_LIMIT - 1; is either an error or over complication for let under_limit = buf.len() < AZURE_REQUEST_LIMIT; \$\endgroup\$ – hellow Jan 14 at 9:30

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.