<commit_before>use flowgger::config::Config;
use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use flate2::{Decompress, Flush, Status};
use std::io::{stderr, Write};
use std::net::UdpSocket;
use std::str;
use std::sync::mpsc::SyncSender;
use super::Input;

const DEFAULT_LISTEN: &'static str = "0.0.0.0:514";
const MAX_UDP_PACKET_SIZE: usize = 65527;
const MAX_COMPRESSION_RATIO: usize = 5;

pub struct UdpInput {
    listen: String,
}

impl UdpInput {
    pub fn new(config: &Config) -> UdpInput {
        let listen = config.lookup("input.listen")
            .map_or(DEFAULT_LISTEN,
                    |x| x.as_str().expect("input.listen must be an ip:port string"))
            .to_owned();
        UdpInput { listen: listen }
    }
}

impl Input for UdpInput {
    fn accept(&self,
              tx: SyncSender<Vec<u8>>,
              decoder: Box<Decoder + Send>,
              encoder: Box<Encoder + Send>) {
        let socket = UdpSocket::bind(&self.listen as &str)
            .expect(&format!("Unable to listen to {}", self.listen));
        let tx = tx.clone();
        let (decoder, encoder): (Box<Decoder>, Box<Encoder>) = (decoder.clone_boxed(),
                                                                encoder.clone_boxed());
        let mut buf = [0; MAX_UDP_PACKET_SIZE];
        loop {
            let (length, _src) = match socket.recv_from(&mut buf) {
                Ok(res) => res,
                Err(_) => continue,
            };
            let line = &buf[..length];
            if let Err(e) = handle_record_maybe_compressed(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}", e);
            }
        }
    }
}

fn handle_record_maybe_compressed(line: &[u8],
                                  tx: &SyncSender<Vec<u8>>,
                                  decoder: &Box<Decoder>,
                                  encoder: &Box<Encoder>)
                                  -> Result<(), &'static str> {
    if line.len() > 2 && line[0] == 0x78 &&
       (line[1] == 0x01 || line[1] == 0x9c || line[1] == 0xda) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match Decompress::new(true).decompress_vec(line, &mut decompressed, Flush::Finish) {
            Ok(Status::Ok) |
            Ok(Status::StreamEnd) => handle_record(&decompressed, tx, decoder, encoder),
            _ => return Err("Corrupted compressed record"),
        }
    } else {
        handle_record(line, tx, decoder, encoder)
    }
}

fn handle_record(line: &[u8],
                 tx: &SyncSender<Vec<u8>>,
                 decoder: &Box<Decoder>,
                 encoder: &Box<Encoder>)
                 -> Result<(), &'static str> {
    let line = match str::from_utf8(line) {
        Err(_) => return Err("Invalid UTF-8 input"),
        Ok(line) => line,
    };
    let decoded = try!(decoder.decode(line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}<commit_msg>Attempt different ways to uncompress UDP+GELF payloads<commit_after>use flowgger::config::Config;
use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use flate2::FlateReadExt;
use std::io::{stderr, Read, Write};
use std::net::UdpSocket;
use std::str;
use std::sync::mpsc::SyncSender;
use super::Input;

const DEFAULT_LISTEN: &'static str = "0.0.0.0:514";
const MAX_UDP_PACKET_SIZE: usize = 65527;
const MAX_COMPRESSION_RATIO: usize = 5;

pub struct UdpInput {
    listen: String,
}

impl UdpInput {
    pub fn new(config: &Config) -> UdpInput {
        let listen = config.lookup("input.listen")
            .map_or(DEFAULT_LISTEN,
                    |x| x.as_str().expect("input.listen must be an ip:port string"))
            .to_owned();
        UdpInput { listen: listen }
    }
}

impl Input for UdpInput {
    fn accept(&self,
              tx: SyncSender<Vec<u8>>,
              decoder: Box<Decoder + Send>,
              encoder: Box<Encoder + Send>) {
        let socket = UdpSocket::bind(&self.listen as &str)
            .expect(&format!("Unable to listen to {}", self.listen));
        let tx = tx.clone();
        let (decoder, encoder): (Box<Decoder>, Box<Encoder>) = (decoder.clone_boxed(),
                                                                encoder.clone_boxed());
        let mut buf = [0; MAX_UDP_PACKET_SIZE];
        loop {
            let (length, _src) = match socket.recv_from(&mut buf) {
                Ok(res) => res,
                Err(_) => continue,
            };
            let line = &buf[..length];
            if let Err(e) = handle_record_maybe_compressed(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}", e);
            }
        }
    }
}

fn handle_record_maybe_compressed(line: &[u8],
                                  tx: &SyncSender<Vec<u8>>,
                                  decoder: &Box<Decoder>,
                                  encoder: &Box<Encoder>)
                                  -> Result<(), &'static str> {
    if line.len() > 2 &&
       (line[0] == 0x78 && (line[1] == 0x01 || line[1] == 0x9c || line[1] == 0xda)) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match line.zlib_decode().read_to_end(&mut decompressed) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => return Err("Corrupted compressed (zlib) record"),
        }
    } else if line.len() > 3 && (line[0] == 0x1f && line[1] == 0x8b && line[2] == 0x08) {
        let mut decompressed = Vec::with_capacity(MAX_UDP_PACKET_SIZE * MAX_COMPRESSION_RATIO);
        match line.gz_decode().and_then(|mut x| x.read_to_end(&mut decompressed)) {
            Ok(_) => handle_record(&decompressed, tx, decoder, encoder),
            Err(_) => return Err("Corrupted compressed (gzip) record"),
        }
    } else {
        handle_record(line, tx, decoder, encoder)
    }
}

fn handle_record(line: &[u8],
                 tx: &SyncSender<Vec<u8>>,
                 decoder: &Box<Decoder>,
                 encoder: &Box<Encoder>)
                 -> Result<(), &'static str> {
    let line = match str::from_utf8(line) {
        Err(_) => return Err("Invalid UTF-8 input"),
        Ok(line) => line,
    };
    let decoded = try!(decoder.decode(line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
{'fraction_non_alphanumeric': 0.14861304819463816, 'fraction_numerical': 0.012862234619556795, 'mean_word_length': 2.099903938520653, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 58, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1662812', 'n_tokens_mistral': 2080, 'n_tokens_neox': 1993, 'n_words': 441}