<commit_before>use std;
use env_logger;
use time;
use portaudio;
use rgsl;
#[cfg(feature="gnuplot")]
use gnuplot;

use thread;
use note;
use mvar;

fn consume<T>(_: T) {}

fn to_fft(mut samples: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = samples.len();
  let r = rgsl::fft::real_radix2::transform(&mut samples, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(samples)
}

fn of_fft(mut fft: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = fft.len();
  let r = rgsl::fft::real_radix2::inverse(&mut fft, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(fft)
}

fn string_err<Ok, Err: std::fmt::Display>(r: Result<Ok, Err>) -> Result<Ok, String> {
  r.map_err(|err| format!("{}", err))
}

fn sine_wave(note: note::T, sample_frequency: f64, start: u32, end: u32) -> Vec<f32> {
  let mut buf = Vec::new();

  let f = note.to_frequency();

  for t in start .. end {
    let t = t as f64 / sample_frequency;
    let sample = (t * f * 2.0 * std::f64::consts::PI).sin();
    let amplitude = 1.0;
    buf.push((sample * amplitude) as f32);
  }

  buf
}

fn with_play_channel<T, F>(sample_rate: f64, buf_size: u32, f: F) -> Result<T, String> where
  F: FnOnce(&mut portaudio::pa::Stream<f32, f32>) -> Result<T, String>,
{
  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_rate,
      buf_size,
      0,
      1,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let r = try!(f(&mut stream));

  try!(string_err(stream.stop()));

  Ok(r)
}

fn play_note(note: note::T, secs: u64) -> Result<(), String> {
  info!("Writing..");

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      let start_time = time::precise_time_ns();
      let mut i = 0;
      while time::precise_time_ns() <= start_time + secs*1_000_000_000 {
        let buf = sine_wave(note, sample_rate, i*buf_size, (i+1)*buf_size);
        assert!(buf.len() == buf_size as usize);
        try!(string_err(stream.write(buf, buf_size)));

        i = i + 1;
      }

      Ok(())
    }
  ));

  Ok(())
}

fn record(sample_frequency: f64, delta_t_ns: u64) -> Result<Vec<f64>, String> {
  let buf_size = 1 << 10;
  assert!(buf_size & (buf_size - 1) == 0);

  info!("Reading..");

  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_frequency,
      buf_size,
      1,
      0,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let mut buf = Vec::new();
  let start_time = time::precise_time_ns();
  while time::precise_time_ns() <= start_time + delta_t_ns || (buf.len() & (buf.len() - 1)) != 0 {
    let new = try!(string_err(stream.read(buf_size)));
    assert!(new.len () == buf_size as usize);
    buf.extend(new.into_iter().map(|f| f as f64));
  }

  try!(string_err(stream.stop()));

  Ok(buf)
}

fn detect_frequency(samples: Vec<f64>, timestep: f64) -> Result<Option<f64>, String> {
  let num_samples = samples.len();
  let samples = try!(to_fft(samples));

  let buckets: Vec<_> =
    (0 .. 1 + num_samples / 2)
    .filter_map(|i| {
      let real = samples[i];
      let imag =
        if i % num_samples / 2 == 0 {
          0.0
        } else {
          samples[num_samples - i]
        };

      let noise_threshold = 0.0;
      let x = real*real + imag*imag - noise_threshold;
      if x > noise_threshold {
        let f = i as f64 / num_samples as f64 / timestep;
        Some((f, x))
      } else {
        None
      }
    })
    .collect();
  consume(samples);

  if buckets.is_empty() {
    return Ok(None)
  }

  #[cfg(feature="gnuplot")]
  {
    let x: Vec<_> = buckets.iter().map(|&(f, _)| f).collect();
    let y: Vec<_> = buckets.iter().map(|&(_, y)| y).collect();
    let mut fg = gnuplot::Figure::new();
    fg
      .axes2d()
      .boxes(
        &x,
        &y,
        &[],
      );
    fg.echo_to_file("/tmp/gnuplot.txt");
  }

  let (max_f, max_val) =
    buckets.iter()
      .fold((0.0, 0.0), |(max_f, max_val), &(f, amp)| {
        let val = amp;
        if val > max_val {
          (f, val)
        } else {
          (max_f, max_val)
        }
      });

  debug!("Max index is {}", max_f);
  debug!("Max value is {}", max_val);

  Ok(Some(max_f))
}

fn harmony(fft: &[f64], semitones_up: f64) -> Vec<f64> {
  let num_samples = fft.len();

  let shift_factor: f64 = (2.0 as f64).powf(-semitones_up / 12.0);

  let mut shifted_fft: Vec<_> = std::iter::repeat(0.0).take(num_samples).collect();
  for i in 1..(num_samples / 2) {
    let source_i = (i as f64 * shift_factor).round() as usize;
    shifted_fft[i] = shifted_fft[i] + fft[source_i];
    shifted_fft[num_samples - i] = shifted_fft[num_samples - i] + fft[num_samples - source_i];
  }

  // This might be broken for nonzero shifts?
  shifted_fft[0] = fft[0];
  shifted_fft[num_samples / 2] = fft[num_samples / 2];

  shifted_fft
}

fn detect_pitch_main() -> Result<(), String> {
  let sample_frequency = 44100.0;
  let samples = mvar::new();

  let _record_thread =
    unsafe {
      thread::scoped(|| {
        loop {
          let new_samples = record(sample_frequency, 100_000_000).unwrap();
          info!("Collected {} samples", new_samples.len());
          samples.overwrite(new_samples);
        }
      })
    };

  let _detect_thread =
    unsafe {
      thread::scoped(|| {
        info!("Detecting pitch..");

        loop {
          let samples = samples.take();
          match detect_frequency(samples, 1.0 / sample_frequency).unwrap() {
            None => {},
            Some(f) => {
              debug!("Estimated frequency is {}", f);
              let note = note::of_frequency(f);
              println!("Estimated note is {}", note.to_string_human());
            },
          }
        }
      })
    };

  Ok(())
}

fn harmony_main(semitones: Vec<f64>) -> Result<(), String> {
  let sample_frequency = 44100.0;

  let samples = record(sample_frequency, 1_000_000_000).unwrap();
  info!("Collected {} samples", samples.len());

  let len = samples.len();
  let original_fft = try!(to_fft(samples));
  let mut fft = original_fft.clone();

  for &semitones in &semitones {
    let harmony = harmony(&original_fft, semitones);
    for i in 0 .. len {
      fft[i] += harmony[i];
    }
  }

  let fft = fft.into_iter().map(|f| f / semitones.len() as f64).collect();

  let samples = try!(of_fft(fft));

  let mut samples = samples.into_iter().map(|f| f as f32);

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      loop {
        let mut snip = Vec::new();
        for _ in 0..buf_size {
          match samples.next() {
            None => break,
            Some(x) => snip.push(x),
          }
        }

        if snip.is_empty() {
          break
        }

        let len = snip.len();
        try!(string_err(stream.write(snip, len as u32)));
      }

      Ok(())
    }
  ));

  Ok(())
}

fn errorful_main() -> Result<(), String> {
  try!(string_err(env_logger::init()));
  try!(string_err(portaudio::pa::initialize()));

  let matches =
    clap_app!(pitchdetect =>
      (@subcommand play =>
        (about: "Play a note")
        (@arg note: +required +takes_value --note)
        (@arg time: +required +takes_value --time)
      )
      (@subcommand detect =>
        (about: "Detect pitch from the microphone")
      )
      (@subcommand harmony =>
        (about: "Add harmony to a recorded clip")
      )
    ).get_matches();

  if let Some(matches) = matches.subcommand_matches("play") {
    let note = try!(note::from_str(matches.value_of("note").unwrap()));
    let duration = try!(string_err(matches.value_of("time").unwrap().parse()));
    try!(play_note(note, duration));
  } else if let Some(_matches) = matches.subcommand_matches("detect") {
    try!(detect_pitch_main());
  } else if let Some(_harmony) = matches.subcommand_matches("harmony") {
    //let semitones = vec!(4, 7, 12);
    let semitones = vec!(4.0);
    try!(harmony_main(semitones));
  }

  Ok(())
}

#[main]
fn main() {
  errorful_main().unwrap();
}
<commit_msg>factoring<commit_after>use std;
use env_logger;
use time;
use portaudio;
use rgsl;
#[cfg(feature="gnuplot")]
use gnuplot;

use thread;
use note;
use mvar;

fn consume<T>(_: T) {}

fn to_fft(mut samples: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = samples.len();
  let r = rgsl::fft::real_radix2::transform(&mut samples, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(samples)
}

fn of_fft(mut fft: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = fft.len();
  let r = rgsl::fft::real_radix2::inverse(&mut fft, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(fft)
}

fn string_err<Ok, Err: std::fmt::Display>(r: Result<Ok, Err>) -> Result<Ok, String> {
  r.map_err(|err| format!("{}", err))
}

fn sine_wave(note: note::T, sample_frequency: f64, start: u32, end: u32) -> Vec<f32> {
  let mut buf = Vec::new();

  let f = note.to_frequency();

  for t in start .. end {
    let t = t as f64 / sample_frequency;
    let sample = (t * f * 2.0 * std::f64::consts::PI).sin();
    let amplitude = 1.0;
    buf.push((sample * amplitude) as f32);
  }

  buf
}

fn with_play_channel<T, F>(sample_rate: f64, buf_size: u32, f: F) -> Result<T, String> where
  F: FnOnce(&mut portaudio::pa::Stream<f32, f32>) -> Result<T, String>,
{
  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_rate,
      buf_size,
      0,
      1,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let r = try!(f(&mut stream));

  try!(string_err(stream.stop()));

  Ok(r)
}

fn play_note(note: note::T, secs: u64) -> Result<(), String> {
  info!("Writing..");

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      let start_time = time::precise_time_ns();
      let mut i = 0;
      while time::precise_time_ns() <= start_time + secs*1_000_000_000 {
        let buf = sine_wave(note, sample_rate, i*buf_size, (i+1)*buf_size);
        assert!(buf.len() == buf_size as usize);
        try!(string_err(stream.write(buf, buf_size)));

        i = i + 1;
      }

      Ok(())
    }
  ));

  Ok(())
}

fn record(sample_frequency: f64, delta_t_ns: u64) -> Result<Vec<f64>, String> {
  let buf_size = 1 << 10;
  assert!(buf_size & (buf_size - 1) == 0);

  info!("Reading..");

  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_frequency,
      buf_size,
      1,
      0,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let mut buf = Vec::new();
  let start_time = time::precise_time_ns();
  while time::precise_time_ns() <= start_time + delta_t_ns || (buf.len() & (buf.len() - 1)) != 0 {
    let new = try!(string_err(stream.read(buf_size)));
    assert!(new.len () == buf_size as usize);
    buf.extend(new.into_iter().map(|f| f as f64));
  }

  try!(string_err(stream.stop()));

  Ok(buf)
}

fn detect_frequency(samples: Vec<f64>, timestep: f64) -> Result<Option<f64>, String> {
  let num_samples = samples.len();
  let samples = try!(to_fft(samples));

  let buckets: Vec<_> =
    (0 .. 1 + num_samples / 2)
    .filter_map(|i| {
      let real = samples[i];
      let imag =
        if i % num_samples / 2 == 0 {
          0.0
        } else {
          samples[num_samples - i]
        };

      let noise_threshold = 0.0;
      let x = real*real + imag*imag - noise_threshold;
      if x > noise_threshold {
        let f = i as f64 / num_samples as f64 / timestep;
        Some((f, x))
      } else {
        None
      }
    })
    .collect();
  consume(samples);

  if buckets.is_empty() {
    return Ok(None)
  }

  #[cfg(feature="gnuplot")]
  {
    let x: Vec<_> = buckets.iter().map(|&(f, _)| f).collect();
    let y: Vec<_> = buckets.iter().map(|&(_, y)| y).collect();
    let mut fg = gnuplot::Figure::new();
    fg
      .axes2d()
      .boxes(
        &x,
        &y,
        &[],
      );
    fg.echo_to_file("/tmp/gnuplot.txt");
  }

  let (max_f, max_val) =
    buckets.iter()
      .fold((0.0, 0.0), |(max_f, max_val), &(f, amp)| {
        let val = amp;
        if val > max_val {
          (f, val)
        } else {
          (max_f, max_val)
        }
      });

  debug!("Max index is {}", max_f);
  debug!("Max value is {}", max_val);

  Ok(Some(max_f))
}

fn harmony(fft: &[f64], semitones_up: f64) -> Vec<f64> {
  let num_samples = fft.len();

  let shift_factor: f64 = (2.0 as f64).powf(-semitones_up / 12.0);

  let mut shifted_fft: Vec<_> = std::iter::repeat(0.0).take(num_samples).collect();
  for i in 1..(num_samples / 2) {
    let source_i = (i as f64 * shift_factor).round() as usize;
    shifted_fft[i] = shifted_fft[i] + fft[source_i];
    shifted_fft[num_samples - i] = shifted_fft[num_samples - i] + fft[num_samples - source_i];
  }

  // This might be broken for nonzero shifts?
  shifted_fft[0] = fft[0];
  shifted_fft[num_samples / 2] = fft[num_samples / 2];

  shifted_fft
}

fn detect_pitch_main() -> Result<(), String> {
  let sample_frequency = 44100.0;
  let samples = mvar::new();

  let _record_thread =
    unsafe {
      thread::scoped(|| {
        loop {
          let new_samples = record(sample_frequency, 100_000_000).unwrap();
          info!("Collected {} samples", new_samples.len());
          samples.overwrite(new_samples);
        }
      })
    };

  let _detect_thread =
    unsafe {
      thread::scoped(|| {
        info!("Detecting pitch..");

        loop {
          let samples = samples.take();
          match detect_frequency(samples, 1.0 / sample_frequency).unwrap() {
            None => {},
            Some(f) => {
              debug!("Estimated frequency is {}", f);
              let note = note::of_frequency(f);
              println!("Estimated note is {}", note.to_string_human());
            },
          }
        }
      })
    };

  Ok(())
}

fn add_harmonies(samples: Vec<f64>, semitones: &[f64]) -> Result<Vec<f64>, String> {
  let len = samples.len();
  let original_fft = try!(to_fft(samples));
  let mut fft = original_fft.clone();

  for &semitones in semitones {
    let harmony = harmony(&original_fft, semitones);
    for i in 0 .. len {
      fft[i] += harmony[i];
    }
  }

  let fft = fft.into_iter().map(|f| f / semitones.len() as f64).collect();

  let samples = try!(of_fft(fft));
  Ok(samples)
}

fn harmony_main(semitones: Vec<f64>) -> Result<(), String> {
  let sample_frequency = 44100.0;

  let samples = record(sample_frequency, 1_000_000_000).unwrap();
  info!("Collected {} samples", samples.len());

  let samples = try!(add_harmonies(samples, &semitones));
  let mut samples = samples.into_iter().map(|f| f as f32);

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      loop {
        let mut snip = Vec::new();
        for _ in 0..buf_size {
          match samples.next() {
            None => break,
            Some(x) => snip.push(x),
          }
        }

        if snip.is_empty() {
          break
        }

        let len = snip.len();
        try!(string_err(stream.write(snip, len as u32)));
      }

      Ok(())
    }
  ));

  Ok(())
}

fn errorful_main() -> Result<(), String> {
  try!(string_err(env_logger::init()));
  try!(string_err(portaudio::pa::initialize()));

  let matches =
    clap_app!(pitchdetect =>
      (@subcommand play =>
        (about: "Play a note")
        (@arg note: +required +takes_value --note)
        (@arg time: +required +takes_value --time)
      )
      (@subcommand detect =>
        (about: "Detect pitch from the microphone")
      )
      (@subcommand harmony =>
        (about: "Add harmony to a recorded clip")
      )
    ).get_matches();

  if let Some(matches) = matches.subcommand_matches("play") {
    let note = try!(note::from_str(matches.value_of("note").unwrap()));
    let duration = try!(string_err(matches.value_of("time").unwrap().parse()));
    try!(play_note(note, duration));
  } else if let Some(_matches) = matches.subcommand_matches("detect") {
    try!(detect_pitch_main());
  } else if let Some(_harmony) = matches.subcommand_matches("harmony") {
    //let semitones = vec!(4, 7, 12);
    let semitones = vec!(4.0);
    try!(harmony_main(semitones));
  }

  Ok(())
}

#[main]
fn main() {
  errorful_main().unwrap();
}
<|endoftext|>


<commit_before>use std::env;
use std::f32;
use std::thread::sleep;
use std::time::Duration;

extern crate getopts;
use getopts::Options;
extern crate rand;
use rand::Rng;
use rand::distributions::{IndependentSample,Range};
extern crate sacn;
use sacn::DmxSource;

const UNIVERSE_SIZE: usize = 510;

const GAMMA: f32 = 2.2;

#[derive(Debug)]
struct Params {
    decay:         f32,
    threshold:     f32,
    max_intensity: f32,
    sleep:         Duration
}

#[derive(Debug)]
#[derive(Clone)]
struct RGB {
    red:   u8,
    green: u8,
    blue:  u8
}

#[derive(Debug)]
#[derive(Clone)]
struct Pixel { intensity: f32, age: u32, temp: u16, rgb: RGB }

struct Zone  { head: u8, body: u8, tail: u8, name: String }

fn build_params () -> Params {
    // seed default params
    let mut params = Params { decay: 0.002, threshold: 0.001, max_intensity: 1_f32, sleep: Duration::new(0, 20_000_000) };

    // parse command line args and adjust params accordingly
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("d", "decay", "slow decay by this factor, defaults to 2", "DECAY");
    opts.optopt("t", "threshold", "probablity that a pixel lights up, default 0.10", "THRESHOLD");
    opts.optopt("m", "maxintensity", "maximum brightness, 1..255, default 75", "MAX");
    opts.optopt("s", "sleep", "sleep interval in seconds, default 1.5", "SECONDS");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    if matches.opt_present("d") {
        params.decay = matches.opt_str("d").unwrap().parse::<f32>().unwrap();
    }
    if matches.opt_present("t") {
        params.threshold = matches.opt_str("t").unwrap().parse::<f32>().unwrap();
    }
    if matches.opt_present("m") {
        let max: u8 = matches.opt_str("m").unwrap().parse::<u8>().unwrap();
        params.max_intensity = (max as f32)/255_f32
    }
    if matches.opt_present("s") {
        // take float seconds
        // convert to int seconds and nanoseconds to make Duration happy
        let seconds: f32 = matches.opt_str("s").unwrap().parse::<f32>().unwrap();
        let whole_seconds: u64 = seconds as u64;
        let nano_seconds: u32 = ((seconds - whole_seconds as f32) * 1_000_000_000_f32) as u32;
        params.sleep = Duration::new(whole_seconds, nano_seconds);
    }
    return params;
}

fn kelvin (mut temp: u16) -> RGB {
    // http://www.tannerhelland.com/4435/convert-temperature-rgb-algorithm-code/
    temp /= 100;

    let mut rgb: RGB = RGB { red: 0, green: 0, blue: 0 };
    // calculate red
    if temp <= 66 {
        rgb.red = 255;
    } else {
        let red: f32 = (temp - 60) as f32;
        rgb.red = (329.698727446 * red.powf(-0.1332047592)).round() as u8;
    }
    // calculate green
    if temp <= 66 {
        let green: f32 = temp as f32;
        rgb.green = (99.4708025861 * green.ln() - 161.1195681661).round() as u8;
    } else {
        let green: f32 = (temp - 60) as f32;
        rgb.green = (288.1221695283 * green.powf(-0.0755148492)).round() as u8;
    }
    // calculate blue
    if temp >= 66 {
        rgb.blue = 255;
    } else {
        if temp <= 19 {
            rgb.blue = 0;
        } else {
            let blue: f32 = (temp - 10) as f32;
            rgb.blue = (138.5177312231 * blue.ln() - 305.0447927307).round() as u8;
        }
    }
    return rgb;
}

fn gamma_correct(rgb: &RGB) -> RGB {
    let mut c: RGB = RGB {red: 0, green: 0, blue: 0 };
    c.red   = (255_f32 * (rgb.red   as f32 / 255_f32).powf(GAMMA)) as u8;
    c.green = (255_f32 * (rgb.green as f32 / 255_f32).powf(GAMMA)) as u8;
    c.blue  = (255_f32 * (rgb.blue  as f32 / 255_f32).powf(GAMMA)) as u8;
    return c;
}

fn scale_rgb(rgb: RGB, intensity: f32, params: &Params) -> RGB {
    let i: f32 = intensity * params.max_intensity;
    let scaled: RGB = RGB {
        red:   (rgb.red   as f32 * i).round() as u8,
        green: (rgb.green as f32 * i).round() as u8,
        blue:  (rgb.blue  as f32 * i).round() as u8
    };
    return scaled;
}

fn main() {
    let params = build_params();

    let dmx = DmxSource::new("Controller").unwrap();

    let zones: [Zone; 6] = [
        Zone { head: 0, body: 44, tail: 3, name: "10".to_string() },
        Zone { head: 2, body: 91, tail: 3, name: "11a".to_string() },
        Zone { head: 2, body: 92, tail: 2, name: "11b".to_string() },
        Zone { head: 2, body: 90, tail: 3, name: "12a".to_string() },
        Zone { head: 2, body: 91, tail: 3, name: "12b".to_string() },
        Zone { head: 2, body: 43, tail: 0, name: "13".to_string() }
    ];

    let mut lights: Vec<Pixel> = vec![];
    // TODO: probably a more idiomatic way to built the default state
    for zone in zones.iter() {
        for i in 1..zone.body {
            let pixel = Pixel {
                intensity: 0_f32,
                age: 0,
                temp: 0,
                rgb: RGB { red: 0, green: 0, blue: 0 },
            };
            lights.push(pixel);
        }
    }

    let mut rng = rand::thread_rng();
    let zero_to_one = Range::new(0_f32, 1_f32);
    let temp_range = Range::new(2700_f32, 5500_f32);

    loop {
        for light in lights.iter_mut() {
            if light.intensity == 0_f32 {
                // light is currently dark
                // test to see if we want to light it
                if zero_to_one.ind_sample(&mut rng) < params.threshold {
                    // we do, so pick a random
                    light.intensity = zero_to_one.ind_sample(&mut rng) as f32;
                    light.age += 1;
                    light.temp = temp_range.ind_sample(&mut rng).round() as u16;
                    light.rgb = scale_rgb(kelvin(light.temp), light.intensity, &params);
                }
            } else {
                // light is lit
                // test to see if we rise or fall
                // probability of falling should go up as light ages
                if zero_to_one.ind_sample(&mut rng) > 1.0/(light.age as f32) {
                    // falling
                    light.intensity -= (zero_to_one.ind_sample(&mut rng) * params.decay) as f32;
                    light.age += 1;
                    // test to see if we bottomed out
                    // check if intensity fell below zero
                    if light.intensity <= 0_f32 {
                        light.intensity = 0_f32;
                        light.age = 0;
                        light.temp = 0;
                        light.rgb.red = 0;
                        light.rgb.green = 0;
                        light.rgb.blue = 0;
                    }
                    // count zeroes, to avoid single-color output
                    //   when gamma-correct bottoms out, usually with red
                    let mut zeroes: u8 = 0;
                    if light.rgb.red == 0 { zeroes += 1 };
                    if light.rgb.green == 0 { zeroes += 1 };
                    if light.rgb.blue == 0 { zeroes += 1 };
                    if zeroes >= 2 {
                        light.intensity = 0_f32;
                        light.age = 0;
                        light.temp = 0;
                        light.rgb.red = 0;
                        light.rgb.green = 0;
                        light.rgb.blue = 0;
                    }
                } else {
                    // rising
                    light.intensity += zero_to_one.ind_sample(&mut rng) * params.decay * (params.max_intensity - light.intensity);

                    light.age += 1;
                    light.rgb = scale_rgb(kelvin(light.temp), light.intensity, &params);
                }
            }
        }
        render(&lights, &zones, &dmx);
        sleep(params.sleep);
    }
    dmx.terminate_stream(1);
}

// debug, rendor output to console
// fn render( lights: &[Pixel], zones: &[Zone], dmx: &DmxSource ) {
//    let mut out: Vec<i16> = vec![];
//    for light in lights {
//        out.push(light.level);
//   }
//    println!("{:?}", out)
//     let mut lit: u16 = 0;
//     let mut level: i16 = 0;
//     for light in lights {
//         if light.level > 0 {
//             lit += 1;
//             level += light.level;
//         }
//     }
//     println!("{} pixels, {} lit, {} avg level", lights.len(), lit, (level as f32)/(lit as f32));
// }

// output to lighting controller via sACN
fn render( lights: &[Pixel], zones: &[Zone], dmx: &DmxSource ) {
    let mut out: Vec<u8> = vec![];
    let mut copy: Vec<Pixel> = vec![];
    copy.extend_from_slice(lights);
    let mut idx: usize = 0;
    for zone in zones {
        // null pixels at head
        for i in 0..zone.head {
            out.push(0); // Red
            out.push(0); // Green
            out.push(0); // Blue
        }
        // set via light.level in the body
        for i in 0..zone.body {
            let ref rgb = copy[idx].rgb;
            let gc = gamma_correct(&rgb);
            out.push(gc.red);
            out.push(gc.green);
            out.push(gc.blue);
            idx += 1;
            // HACK: patching up an off-by-one somewhere
            if idx == copy.len() { break };
        }
        // null pixels at tail
        for i in 0..zone.tail {
            out.push(0); // Red
            out.push(0); // Green
            out.push(0); // Blue
        }
    }
    let mut universes = Vec::new();
    while out.len() > UNIVERSE_SIZE {
        let u = out.split_off(UNIVERSE_SIZE);
        universes.push(out);
        out = u;
    }
    universes.push(out);
    let mut universe: u16 = 1;
    for u in universes {
        dmx.send(universe, &u);
        universe += 1;
    }
}

<commit_msg>trying (so far in vain) to detect bottoming out to red<commit_after>use std::env;
use std::f32;
use std::thread::sleep;
use std::time::Duration;

extern crate getopts;
use getopts::Options;
extern crate rand;
use rand::Rng;
use rand::distributions::{IndependentSample,Range};
extern crate sacn;
use sacn::DmxSource;

const UNIVERSE_SIZE: usize = 510;

const GAMMA: f32 = 2.2;

#[derive(Debug)]
struct Params {
    decay:         f32,
    threshold:     f32,
    max_intensity: f32,
    sleep:         Duration
}

#[derive(Debug)]
#[derive(Clone)]
struct RGB {
    red:   u8,
    green: u8,
    blue:  u8
}

#[derive(Debug)]
#[derive(Clone)]
struct Pixel { intensity: f32, age: u32, temp: u16, rgb: RGB }

struct Zone  { head: u8, body: u8, tail: u8, name: String }

fn build_params () -> Params {
    // seed default params
    let mut params = Params { decay: 0.002, threshold: 0.001, max_intensity: 0.8_f32, sleep: Duration::new(0, 20_000_000) };

    // parse command line args and adjust params accordingly
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    let mut opts = Options::new();
    opts.optopt("d", "decay", "slow decay by this factor, defaults to 2", "DECAY");
    opts.optopt("t", "threshold", "probablity that a pixel lights up, default 0.10", "THRESHOLD");
    opts.optopt("m", "maxintensity", "maximum brightness, 1..255, default 75", "MAX");
    opts.optopt("s", "sleep", "sleep interval in seconds, default 1.5", "SECONDS");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    if matches.opt_present("d") {
        params.decay = matches.opt_str("d").unwrap().parse::<f32>().unwrap();
    }
    if matches.opt_present("t") {
        params.threshold = matches.opt_str("t").unwrap().parse::<f32>().unwrap();
    }
    if matches.opt_present("m") {
        let max: u8 = matches.opt_str("m").unwrap().parse::<u8>().unwrap();
        params.max_intensity = (max as f32)/255_f32
    }
    if matches.opt_present("s") {
        // take float seconds
        // convert to int seconds and nanoseconds to make Duration happy
        let seconds: f32 = matches.opt_str("s").unwrap().parse::<f32>().unwrap();
        let whole_seconds: u64 = seconds as u64;
        let nano_seconds: u32 = ((seconds - whole_seconds as f32) * 1_000_000_000_f32) as u32;
        params.sleep = Duration::new(whole_seconds, nano_seconds);
    }
    return params;
}

fn kelvin (mut temp: u16) -> RGB {
    // http://www.tannerhelland.com/4435/convert-temperature-rgb-algorithm-code/
    temp /= 100;

    let mut rgb: RGB = RGB { red: 0, green: 0, blue: 0 };
    // calculate red
    if temp <= 66 {
        rgb.red = 255;
    } else {
        let red: f32 = (temp - 60) as f32;
        rgb.red = (329.698727446 * red.powf(-0.1332047592)).round() as u8;
    }
    // calculate green
    if temp <= 66 {
        let green: f32 = temp as f32;
        rgb.green = (99.4708025861 * green.ln() - 161.1195681661).round() as u8;
    } else {
        let green: f32 = (temp - 60) as f32;
        rgb.green = (288.1221695283 * green.powf(-0.0755148492)).round() as u8;
    }
    // calculate blue
    if temp >= 66 {
        rgb.blue = 255;
    } else {
        if temp <= 19 {
            rgb.blue = 0;
        } else {
            let blue: f32 = (temp - 10) as f32;
            rgb.blue = (138.5177312231 * blue.ln() - 305.0447927307).round() as u8;
        }
    }
    return rgb;
}

fn gamma_correct(rgb: &RGB) -> RGB {
    let mut c: RGB = RGB {red: 0, green: 0, blue: 0 };
    c.red   = (255_f32 * (rgb.red   as f32 / 255_f32).powf(GAMMA)) as u8;
    c.green = (255_f32 * (rgb.green as f32 / 255_f32).powf(GAMMA)) as u8;
    c.blue  = (255_f32 * (rgb.blue  as f32 / 255_f32).powf(GAMMA)) as u8;
    return c;
}

fn scale_rgb(rgb: RGB, intensity: f32, params: &Params) -> RGB {
    let i: f32 = intensity * params.max_intensity;
    let scaled: RGB = RGB {
        red:   (rgb.red   as f32 * i).round() as u8,
        green: (rgb.green as f32 * i).round() as u8,
        blue:  (rgb.blue  as f32 * i).round() as u8
    };
    return scaled;
}

fn main() {
    let params = build_params();

    let dmx = DmxSource::new("Controller").unwrap();

    let zones: [Zone; 6] = [
        Zone { head: 0, body: 44, tail: 3, name: "10".to_string() },
        Zone { head: 2, body: 91, tail: 3, name: "11a".to_string() },
        Zone { head: 2, body: 92, tail: 2, name: "11b".to_string() },
        Zone { head: 2, body: 90, tail: 3, name: "12a".to_string() },
        Zone { head: 2, body: 91, tail: 3, name: "12b".to_string() },
        Zone { head: 2, body: 43, tail: 0, name: "13".to_string() }
    ];

    let mut lights: Vec<Pixel> = vec![];
    // TODO: probably a more idiomatic way to built the default state
    for zone in zones.iter() {
        for i in 1..zone.body {
            let pixel = Pixel {
                intensity: 0_f32,
                age: 0,
                temp: 0,
                rgb: RGB { red: 0, green: 0, blue: 0 },
            };
            lights.push(pixel);
        }
    }

    let mut rng = rand::thread_rng();
    let zero_to_one = Range::new(0_f32, 1_f32);
    let temp_range = Range::new(2700_f32, 5500_f32);

    loop {
        for light in lights.iter_mut() {
            if light.intensity == 0_f32 {
                // light is currently dark
                // test to see if we want to light it
                if zero_to_one.ind_sample(&mut rng) < params.threshold {
                    // we do, so pick a random
                    light.intensity = zero_to_one.ind_sample(&mut rng) as f32;
                    light.age += 1;
                    light.temp = temp_range.ind_sample(&mut rng).round() as u16;
                    light.rgb = scale_rgb(kelvin(light.temp), light.intensity, &params);
                }
            } else {
                // light is lit
                // test to see if we rise or fall
                // probability of falling should go up as light ages
                if zero_to_one.ind_sample(&mut rng) > 1.0/(light.age as f32) {
                    // falling
                    light.intensity -= (zero_to_one.ind_sample(&mut rng) * params.decay) as f32;
                    light.age += 1;
                    // test to see if we bottomed out
                    // check if intensity fell below zero
                    if light.intensity <= 0_f32 {
                        light.intensity = 0_f32;
                        light.age = 0;
                        light.temp = 0;
                        light.rgb.red = 0;
                        light.rgb.green = 0;
                        light.rgb.blue = 0;
                    }
                    // count zeroes, to avoid single-color output
                    //   when gamma-correct bottoms out, usually with red
                    let mut zeroes: u8 = 0;
                    if light.rgb.red == 0 { zeroes += 1 };
                    if light.rgb.green == 0 { zeroes += 1 };
                    if light.rgb.blue == 0 { zeroes += 1 };
                    if zeroes >= 2 {
                        light.intensity = 0_f32;
                        light.age = 0;
                        light.temp = 0;
                        light.rgb.red = 0;
                        light.rgb.green = 0;
                        light.rgb.blue = 0;
                    }
                } else {
                    // rising
                    light.intensity += zero_to_one.ind_sample(&mut rng) * params.decay * (params.max_intensity - light.intensity);

                    light.age += 1;
                    light.rgb = scale_rgb(kelvin(light.temp), light.intensity, &params);
                }
            }
        }
        render(&lights, &zones, &dmx);
        sleep(params.sleep);
    }
    dmx.terminate_stream(1);
}

// debug, rendor output to console
// fn render( lights: &[Pixel], zones: &[Zone], dmx: &DmxSource ) {
//    let mut out: Vec<i16> = vec![];
//    for light in lights {
//        out.push(light.level);
//   }
//    println!("{:?}", out)
//     let mut lit: u16 = 0;
//     let mut level: i16 = 0;
//     for light in lights {
//         if light.level > 0 {
//             lit += 1;
//             level += light.level;
//         }
//     }
//     println!("{} pixels, {} lit, {} avg level", lights.len(), lit, (level as f32)/(lit as f32));
// }

// output to lighting controller via sACN
fn render( lights: &[Pixel], zones: &[Zone], dmx: &DmxSource ) {
    let mut out: Vec<u8> = vec![];
    let mut copy: Vec<Pixel> = vec![];
    copy.extend_from_slice(lights);
    let mut idx: usize = 0;
    for zone in zones {
        // null pixels at head
        for i in 0..zone.head {
            out.push(0); // Red
            out.push(0); // Green
            out.push(0); // Blue
        }
        // set via light.level in the body
        for i in 0..zone.body {
            let ref rgb = copy[idx].rgb;
            let gc = gamma_correct(&rgb);
            out.push(gc.red);
            out.push(gc.green);
            out.push(gc.blue);
            idx += 1;
            // HACK: patching up an off-by-one somewhere
            if idx == copy.len() { break };
        }
        // null pixels at tail
        for i in 0..zone.tail {
            out.push(0); // Red
            out.push(0); // Green
            out.push(0); // Blue
        }
    }
    let mut universes = Vec::new();
    while out.len() > UNIVERSE_SIZE {
        let u = out.split_off(UNIVERSE_SIZE);
        universes.push(out);
        out = u;
    }
    universes.push(out);
    let mut universe: u16 = 1;
    for u in universes {
        dmx.send(universe, &u);
        universe += 1;
    }
}

<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, OpCode, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

struct DnsRelay<Remote: upstream::Upstream> {
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote: upstream::Upstream> DnsRelay<Remote> {
    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_upstream;
        // Start querying name servers
        debug!(
            "attempting lookup of {:?} {} with ns {:?} and {:?}",
            query.query_type(), query.name(), local, remote
        );

        let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
        let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

        match should_forward_by_query(acl, query) {
            Some(true) => {
                let remote_response = remote_response_fut.await;
                debug!("pick remote response (query): {:?}", remote_response);
                return (remote_response, true);
            }
            Some(false) => {
                let local_response = local_response_fut.await;
                debug!("pick local response (query): {:?}", local_response);
                return (local_response, false);
            }
            None => (),
        }

        let decider = async {
            let local_response = local_response_fut.await;
            if should_forward_by_response(acl, &local_response, query) {
                None
            } else {
                Some(local_response)
            }
        };
        tokio::pin!(remote_response_fut, decider);
        let mut use_remote = false;
        let mut remote_response = None;
        loop {
            tokio::select! {
                response = &mut remote_response_fut, if remote_response.is_none() => {
                    if use_remote {
                        debug!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    if let Some(local_response) = decision {
                        debug!("pick local response (response): {:?}", local_response);
                        return (local_response, false);
                    } else if let Some(remote_response) = remote_response {
                        debug!("pick remote response (response): {:?}", remote_response);
                        return (remote_response, true);
                    } else {
                        use_remote = true;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);
        if !request.recursion_desired() {
            message.set_recursion_desired(false);
            message.set_response_code(ResponseCode::NotImp);
        } else if request.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    debug!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                        RData::AAAA(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let relay = Arc::new(DnsRelay {
        context: context.clone(),
        remote_upstream: upstream::ProxyTcpUpstream {
            context: context.clone(),
            svr_cfg: move || balancer.pick_server().server_config().clone(),
            ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
        }
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let relay = relay.clone();
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let message = relay.resolve(request).await;
            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>UdpSocket::send_to should not block<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    sync::Arc,
    time::Duration,
};

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, OpCode, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

pub mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || match iter.next() {
        Some(label) => std::str::from_utf8(label).unwrap_or("*"),
        None => "0",    // zero fill the missing labels
    };
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

fn check_name_in_proxy_list(acl: &AccessControl, name: &Name) -> Option<bool> {
    if name.is_fqdn() {
        // remove the last dot from FQDN
        let mut name = name.to_ascii();
        name.pop();
        acl.check_host_in_proxy_list(&name)
    } else {
        // unconditionally use default for PQDNs
        Some(acl.is_default_in_proxy_list())
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = check_name_in_proxy_list(acl, query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => unreachable!(),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match check_name_in_proxy_list(acl, name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V4(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::V6(*ip)),
                    RData::PTR(_) => unreachable!(),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        unreachable!()
    }
}

struct DnsRelay<Remote: upstream::Upstream> {
    context: SharedContext,
    remote_upstream: Remote,
}

impl<Remote: upstream::Upstream> DnsRelay<Remote> {
    async fn acl_lookup(&self, query: &Query) -> (io::Result<Message>, bool) {
        let acl = self.context.acl();
        let local = self.context.local_dns();
        let remote = &self.remote_upstream;
        // Start querying name servers
        debug!(
            "attempting lookup of {:?} {} with ns {:?} and {:?}",
            query.query_type(), query.name(), local, remote
        );

        let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
        let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

        match should_forward_by_query(acl, query) {
            Some(true) => {
                let remote_response = remote_response_fut.await;
                debug!("pick remote response (query): {:?}", remote_response);
                return (remote_response, true);
            }
            Some(false) => {
                let local_response = local_response_fut.await;
                debug!("pick local response (query): {:?}", local_response);
                return (local_response, false);
            }
            None => (),
        }

        let decider = async {
            let local_response = local_response_fut.await;
            if should_forward_by_response(acl, &local_response, query) {
                None
            } else {
                Some(local_response)
            }
        };
        tokio::pin!(remote_response_fut, decider);
        let mut use_remote = false;
        let mut remote_response = None;
        loop {
            tokio::select! {
                response = &mut remote_response_fut, if remote_response.is_none() => {
                    if use_remote {
                        debug!("pick remote response (response): {:?}", response);
                        return (response, true);
                    } else {
                        remote_response = Some(response);
                    }
                }
                decision = &mut decider, if !use_remote => {
                    if let Some(local_response) = decision {
                        debug!("pick local response (response): {:?}", local_response);
                        return (local_response, false);
                    } else if let Some(remote_response) = remote_response {
                        debug!("pick remote response (response): {:?}", remote_response);
                        return (remote_response, true);
                    } else {
                        use_remote = true;
                    }
                }
                else => unreachable!(),
            }
        }
    }

    async fn resolve(&self, request: Message) -> Message {
        let mut message = Message::new();
        message.set_id(request.id());
        message.set_recursion_desired(true);
        message.set_recursion_available(true);
        message.set_message_type(MessageType::Response);
        if !request.recursion_desired() {
            message.set_recursion_desired(false);
            message.set_response_code(ResponseCode::NotImp);
        } else if request.op_code() != OpCode::Query || request.message_type() != MessageType::Query {
            message.set_response_code(ResponseCode::NotImp);
        } else if request.query_count() > 0 {
            let (r, forward) = self.acl_lookup(&request.queries()[0]).await;
            if let Ok(result) = r {
                for rec in result.answers() {
                    debug!("dns answer: {:?}", rec);
                    match rec.rdata() {
                        RData::A(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V4(*ip), forward),
                        RData::AAAA(ref ip) => self.context.add_to_reverse_lookup_cache(&IpAddr::V6(*ip), forward),
                        _ => (),
                    }
                }
                message = result;
                message.set_id(request.id());
            } else {
                message.set_response_code(ResponseCode::ServFail);
            }
        }
        message
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, tx) = socket.split();
    let tx = Arc::new(tokio::sync::Mutex::new(tx));

    let config = context.config();
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let relay = Arc::new(DnsRelay {
        context: context.clone(),
        remote_upstream: upstream::ProxyTcpUpstream {
            context: context.clone(),
            svr_cfg: move || balancer.pick_server().server_config().clone(),
            ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
        }
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let relay = relay.clone();
        let tx = tx.clone();

        tokio::spawn(async move {
            let message = relay.resolve(request).await;
            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = tx.lock().await.send_to(&res_buffer, &src).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>//! Bindings to the client library `libwayland-server.so`
//!
//! The generated handle is named `WAYLAND_SERVER_HANDLE`

#![cfg_attr(rustfmt, rustfmt_skip)]

use super::common::*;
use libc::{gid_t, pid_t, uid_t};
use std::os::raw::{c_char, c_int, c_void};

pub enum wl_client {}
pub enum wl_display {}
pub enum wl_event_loop {}
pub enum wl_event_source {}
pub enum wl_global {}
pub enum wl_resource {}
pub enum wl_shm_buffer {}

pub type wl_event_loop_fd_func_t = unsafe extern "C" fn(c_int, u32, *mut c_void) -> c_int;
pub type wl_event_loop_timer_func_t = unsafe extern "C" fn(*mut c_void) -> c_int;
pub type wl_event_loop_signal_func_t = unsafe extern "C" fn(c_int, *mut c_void) -> c_int;
pub type wl_event_loop_idle_func_t = unsafe extern "C" fn(*mut c_void) -> ();
pub type wl_global_bind_func_t = unsafe extern "C" fn(*mut wl_client, *mut c_void, u32, u32) -> ();
pub type wl_notify_func_t = unsafe extern "C" fn(*mut wl_listener, *mut c_void) -> ();
pub type wl_resource_destroy_func_t = unsafe extern "C" fn(*mut wl_resource) -> ();

#[repr(C)]
pub struct wl_listener {
    pub link: wl_list,
    pub notify: wl_notify_func_t,
}

#[repr(C)]
pub struct wl_signal {
    pub listener_list: wl_list,
}

external_library!(WaylandServer, "wayland-server",
    functions:
    // wl_client
        fn wl_client_flush(*mut wl_client) -> (),
        fn wl_client_destroy(*mut wl_client) -> (),
        fn wl_client_get_display(*mut wl_client) -> *mut wl_display,
        fn wl_client_get_credentials(*mut wl_client, *mut pid_t, *mut uid_t, *mut gid_t) -> (),
        fn wl_client_get_object(*mut wl_client, u32) -> *mut wl_resource,
        fn wl_client_add_destroy_listener(*mut wl_client, *mut wl_listener) -> (),
        fn wl_client_get_destroy_listener(*mut wl_client, wl_notify_func_t) -> *mut wl_listener,
        fn wl_client_post_no_memory(*mut wl_client) -> (),
        fn wl_resource_create(*mut wl_client, *const wl_interface, c_int, u32) -> *mut wl_resource,
    // wl_display
        fn wl_client_create(*mut wl_display, c_int) -> *mut wl_client,
        fn wl_display_create() -> *mut wl_display,
        fn wl_display_destroy(*mut wl_display) -> (),
        fn wl_display_get_serial(*mut wl_display) -> u32,
        fn wl_display_next_serial(*mut wl_display) -> u32,
        fn wl_display_add_socket(*mut wl_display, *const c_char) -> c_int,
        fn wl_display_add_socket_auto(*mut wl_display) -> *const c_char,
        fn wl_display_add_socket_fd(*mut wl_display, c_int) -> c_int,
        fn wl_display_add_shm_format(*mut wl_display, u32) -> *mut u32,
        fn wl_display_get_event_loop(*mut wl_display) -> *mut wl_event_loop,
        fn wl_display_terminate(*mut wl_display) -> (),
        fn wl_display_run(*mut wl_display) -> (),
        fn wl_display_flush_clients(*mut wl_display) -> (),
        fn wl_display_add_destroy_listener(*mut wl_display, *mut wl_listener) -> (),
        fn wl_display_get_destroy_listener(*mut wl_display, wl_notify_func_t) -> *mut wl_listener,
        fn wl_global_create(*mut wl_display, *const wl_interface, c_int, *mut c_void, wl_global_bind_func_t) -> *mut wl_global,
        fn wl_display_init_shm(*mut wl_display) -> c_int,
    // wl_event_loop
        fn wl_event_loop_create() -> *mut wl_event_loop,
        fn wl_event_loop_destroy(*mut wl_event_loop) -> (),
        fn wl_event_loop_add_fd(*mut wl_event_loop, c_int, u32, wl_event_loop_fd_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_add_timer(*mut wl_event_loop, wl_event_loop_timer_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_add_signal(*mut wl_event_loop, c_int, wl_event_loop_signal_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_dispatch(*mut wl_event_loop, c_int) -> c_int,
        fn wl_event_loop_dispatch_idle(*mut wl_event_loop) -> (),
        fn wl_event_loop_add_idle(*mut wl_event_loop, wl_event_loop_idle_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_get_fd(*mut wl_event_loop) -> c_int,
        fn wl_event_loop_add_destroy_listener(*mut wl_event_loop, *mut wl_listener) -> (),
        fn wl_event_loop_get_destroy_listener(*mut wl_event_loop, wl_notify_func_t) -> *mut wl_listener,
    // wl_event_source
        fn wl_event_source_fd_update(*mut wl_event_source, u32) -> c_int,
        fn wl_event_source_timer_update(*mut wl_event_source, c_int) -> c_int,
        fn wl_event_source_remove(*mut wl_event_source) -> c_int,
        fn wl_event_source_check(*mut wl_event_source) -> (),
    // wl_global
        fn wl_global_destroy(*mut wl_global) -> (),
    // wl_resource
        fn wl_resource_post_event_array(*mut wl_resource, u32, *mut wl_argument) -> (),
        fn wl_resource_queue_event_array(*mut wl_resource, u32, *mut wl_argument) -> (),
        fn wl_resource_post_no_memory(*mut wl_resource) -> (),
        fn wl_resource_set_implementation(*mut wl_resource, *const c_void, *mut c_void, Option<wl_resource_destroy_func_t>) -> (),
        fn wl_resource_set_dispatcher(*mut wl_resource, wl_dispatcher_func_t, *const c_void, *mut c_void, Option<wl_resource_destroy_func_t>) -> (),
        fn wl_resource_destroy(*mut wl_resource) -> (),
        fn wl_resource_get_client(*mut wl_resource) -> *mut wl_client,
        fn wl_resource_get_id(*mut wl_resource) -> u32,
        fn wl_resource_get_link(*mut wl_resource) -> *mut wl_list,
        fn wl_resource_from_link(*mut wl_list) -> *mut wl_resource,
        fn wl_resource_find_for_client(*mut wl_list, *mut wl_client) -> (),
        fn wl_resource_set_user_data(*mut wl_resource, *mut c_void) -> (),
        fn wl_resource_get_user_data(*mut wl_resource) -> *mut c_void,
        fn wl_resource_get_version(*mut wl_resource) -> c_int,
        fn wl_resource_set_destructor(*mut wl_resource, Option<wl_resource_destroy_func_t>) -> (),
        fn wl_resource_instance_of(*mut wl_resource, *const wl_interface, *const c_void) -> c_int,
        fn wl_resource_add_destroy_listener(*mut wl_resource, wl_notify_func_t) -> (),
        fn wl_resource_get_destroy_listener(*mut wl_resource,wl_notify_func_t) -> *mut wl_listener,
    // wl_shm
        fn wl_shm_buffer_begin_access(*mut wl_shm_buffer) -> (),
        fn wl_shm_buffer_end_access(*mut wl_shm_buffer) -> (),
        fn wl_shm_buffer_get(*mut wl_resource) -> *mut wl_shm_buffer,
        fn wl_shm_buffer_get_data(*mut wl_shm_buffer) -> *mut c_void,
        fn wl_shm_buffer_get_stride(*mut wl_shm_buffer) -> i32,
        fn wl_shm_buffer_get_format(*mut wl_shm_buffer) -> u32,
        fn wl_shm_buffer_get_width(*mut wl_shm_buffer) -> i32,
        fn wl_shm_buffer_get_height(*mut wl_shm_buffer) -> i32,
    // wl_log
        fn wl_log_set_handler_server(wl_log_func_t) -> (),
    // wl_list
        fn wl_list_init(*mut wl_list) -> (),
        fn wl_list_insert(*mut wl_list, *mut wl_list) -> (),
        fn wl_list_remove(*mut wl_list) -> (),
        fn wl_list_length(*const wl_list) -> c_int,
        fn wl_list_empty(*const wl_list) -> c_int,
        fn wl_list_insert_list(*mut wl_list,*mut wl_list) -> (),

    // arrays
        fn wl_array_init(*mut wl_array) -> (),
        fn wl_array_release(*mut wl_array) -> (),
        fn wl_array_add(*mut wl_array,usize) -> (),
        fn wl_array_copy(*mut wl_array, *mut wl_array) -> (),
    varargs:
        fn wl_resource_post_event(*mut wl_resource, u32) -> (),
        fn wl_resource_queue_event(*mut wl_resource, u32) -> (),
        fn wl_resource_post_error(*mut wl_resource, u32, *const c_char) -> (),
);

#[cfg(feature = "dlopen")]
lazy_static!(
    pub static ref WAYLAND_SERVER_OPTION: Option<WaylandServer> = {
        // This is a workaround for Ubuntu 17.04, which doesn't have a bare symlink
        // for libwayland-server.so but does have it with the version numbers for
        // whatever reason.
        //
        // We could do some trickery with str slices but that is more trouble
        // than its worth
        let versions = ["libwayland-server.so",
                        "libwayland-server.so.0"];
        for ver in &versions {
            match WaylandServer::open(ver) {
                Ok(h) => return Some(h),
                Err(::dlib::DlError::NotFound) => continue,
                Err(::dlib::DlError::MissingSymbol(s)) => {
                    if ::std::env::var_os("WAYLAND_RS_DEBUG").is_some() {
                        // only print debug messages if WAYLAND_RS_DEBUG is set
                        eprintln!("[wayland-server] Found library {} cannot be used: symbol {} is missing.", ver, s);
                    }
                    return None;
                }
            }
        }
        None
    };
    pub static ref WAYLAND_SERVER_HANDLE: &'static WaylandServer = {
        WAYLAND_SERVER_OPTION.as_ref().expect("Library libwayland-server.so could not be loaded.")
    };
);

#[cfg(not(feature = "dlopen"))]
pub fn is_lib_available() -> bool {
    true
}
#[cfg(feature = "dlopen")]
pub fn is_lib_available() -> bool {
    WAYLAND_SERVER_OPTION.is_some()
}

pub mod signal {
    #[cfg(not(feature = "dlopen"))]
    use super::{wl_list_init, wl_list_insert};
    use super::{wl_listener, wl_notify_func_t, wl_signal};
    #[cfg(feature = "dlopen")]
    use super::WAYLAND_SERVER_HANDLE as WSH;
    use common::wl_list;
    use std::os::raw::c_void;

    // TODO: Is this really not UB ?
    macro_rules! offset_of(
        ($ty:ty, $field:ident) => {
            &(*(0 as *const $ty)).$field as *const _ as usize
        }
    );

    macro_rules! container_of(
        ($ptr: expr, $container: ty, $field: ident) => {
            ($ptr as *mut u8).offset(-(offset_of!($container, $field) as isize)) as *mut $container
        }
    );

    macro_rules! list_for_each(
        ($pos: ident, $head:expr, $container: ty, $field: ident, $action: block) => {
            let mut $pos = container_of!((*$head).next, $container, $field);
            while &mut (*$pos).$field as *mut _ != $head {
                $action;
                $pos = container_of!((*$pos).$field.next, $container, $field);
            }
        }
    );

    macro_rules! list_for_each_safe(
        ($pos: ident, $head: expr, $container: ty, $field: ident, $action: block) => {
            let mut $pos = container_of!((*$head).next, $container, $field);
            let mut tmp = container_of!((*$pos).$field.next, $container, $field);
            while &mut (*$pos).$field as *mut _ != $head {
                $action;
                $pos = tmp;
                tmp = container_of!((*$pos).$field.next, $container, $field);
            }
        }
    );

    pub unsafe fn wl_signal_init(signal: *mut wl_signal) {
        ffi_dispatch!(WSH, wl_list_init, &mut (*signal).listener_list);
    }

    pub unsafe fn wl_signal_add(signal: *mut wl_signal, listener: *mut wl_listener) {
        ffi_dispatch!(
            WSH,
            wl_list_insert,
            (*signal).listener_list.prev,
            &mut (*listener).link
        )
    }

    pub unsafe fn wl_signal_get(signal: *mut wl_signal, notify: wl_notify_func_t) -> *mut wl_listener {
        list_for_each!(
            l,
            &mut (*signal).listener_list as *mut wl_list,
            wl_listener,
            link,
            {
                if (*l).notify == notify {
                    return l;
                }
            }
        );
        return ::std::ptr::null_mut();
    }

    pub unsafe fn wl_signal_emit(signal: *mut wl_signal, data: *mut c_void) {
        list_for_each_safe!(
            l,
            &mut (*signal).listener_list as *mut wl_list,
            wl_listener,
            link,
            {
                ((*l).notify)(l, data);
            }
        );
    }
}
<commit_msg>sys: ListernerWithUserData<commit_after>//! Bindings to the client library `libwayland-server.so`
//!
//! The generated handle is named `WAYLAND_SERVER_HANDLE`

#![cfg_attr(rustfmt, rustfmt_skip)]

use super::common::*;
use libc::{gid_t, pid_t, uid_t};
use std::os::raw::{c_char, c_int, c_void};

pub enum wl_client {}
pub enum wl_display {}
pub enum wl_event_loop {}
pub enum wl_event_source {}
pub enum wl_global {}
pub enum wl_resource {}
pub enum wl_shm_buffer {}

pub type wl_event_loop_fd_func_t = unsafe extern "C" fn(c_int, u32, *mut c_void) -> c_int;
pub type wl_event_loop_timer_func_t = unsafe extern "C" fn(*mut c_void) -> c_int;
pub type wl_event_loop_signal_func_t = unsafe extern "C" fn(c_int, *mut c_void) -> c_int;
pub type wl_event_loop_idle_func_t = unsafe extern "C" fn(*mut c_void) -> ();
pub type wl_global_bind_func_t = unsafe extern "C" fn(*mut wl_client, *mut c_void, u32, u32) -> ();
pub type wl_notify_func_t = unsafe extern "C" fn(*mut wl_listener, *mut c_void) -> ();
pub type wl_resource_destroy_func_t = unsafe extern "C" fn(*mut wl_resource) -> ();

#[repr(C)]
pub struct wl_listener {
    pub link: wl_list,
    pub notify: wl_notify_func_t,
}

#[repr(C)]
pub struct wl_signal {
    pub listener_list: wl_list,
}

external_library!(WaylandServer, "wayland-server",
    functions:
    // wl_client
        fn wl_client_flush(*mut wl_client) -> (),
        fn wl_client_destroy(*mut wl_client) -> (),
        fn wl_client_get_display(*mut wl_client) -> *mut wl_display,
        fn wl_client_get_credentials(*mut wl_client, *mut pid_t, *mut uid_t, *mut gid_t) -> (),
        fn wl_client_get_object(*mut wl_client, u32) -> *mut wl_resource,
        fn wl_client_add_destroy_listener(*mut wl_client, *mut wl_listener) -> (),
        fn wl_client_get_destroy_listener(*mut wl_client, wl_notify_func_t) -> *mut wl_listener,
        fn wl_client_post_no_memory(*mut wl_client) -> (),
        fn wl_resource_create(*mut wl_client, *const wl_interface, c_int, u32) -> *mut wl_resource,
    // wl_display
        fn wl_client_create(*mut wl_display, c_int) -> *mut wl_client,
        fn wl_display_create() -> *mut wl_display,
        fn wl_display_destroy(*mut wl_display) -> (),
        fn wl_display_get_serial(*mut wl_display) -> u32,
        fn wl_display_next_serial(*mut wl_display) -> u32,
        fn wl_display_add_socket(*mut wl_display, *const c_char) -> c_int,
        fn wl_display_add_socket_auto(*mut wl_display) -> *const c_char,
        fn wl_display_add_socket_fd(*mut wl_display, c_int) -> c_int,
        fn wl_display_add_shm_format(*mut wl_display, u32) -> *mut u32,
        fn wl_display_get_event_loop(*mut wl_display) -> *mut wl_event_loop,
        fn wl_display_terminate(*mut wl_display) -> (),
        fn wl_display_run(*mut wl_display) -> (),
        fn wl_display_flush_clients(*mut wl_display) -> (),
        fn wl_display_add_destroy_listener(*mut wl_display, *mut wl_listener) -> (),
        fn wl_display_get_destroy_listener(*mut wl_display, wl_notify_func_t) -> *mut wl_listener,
        fn wl_global_create(*mut wl_display, *const wl_interface, c_int, *mut c_void, wl_global_bind_func_t) -> *mut wl_global,
        fn wl_display_init_shm(*mut wl_display) -> c_int,
        fn wl_display_add_client_created_listener(*mut wl_display, *mut wl_listener) -> (),
    // wl_event_loop
        fn wl_event_loop_create() -> *mut wl_event_loop,
        fn wl_event_loop_destroy(*mut wl_event_loop) -> (),
        fn wl_event_loop_add_fd(*mut wl_event_loop, c_int, u32, wl_event_loop_fd_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_add_timer(*mut wl_event_loop, wl_event_loop_timer_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_add_signal(*mut wl_event_loop, c_int, wl_event_loop_signal_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_dispatch(*mut wl_event_loop, c_int) -> c_int,
        fn wl_event_loop_dispatch_idle(*mut wl_event_loop) -> (),
        fn wl_event_loop_add_idle(*mut wl_event_loop, wl_event_loop_idle_func_t, *mut c_void) -> *mut wl_event_source,
        fn wl_event_loop_get_fd(*mut wl_event_loop) -> c_int,
        fn wl_event_loop_add_destroy_listener(*mut wl_event_loop, *mut wl_listener) -> (),
        fn wl_event_loop_get_destroy_listener(*mut wl_event_loop, wl_notify_func_t) -> *mut wl_listener,
    // wl_event_source
        fn wl_event_source_fd_update(*mut wl_event_source, u32) -> c_int,
        fn wl_event_source_timer_update(*mut wl_event_source, c_int) -> c_int,
        fn wl_event_source_remove(*mut wl_event_source) -> c_int,
        fn wl_event_source_check(*mut wl_event_source) -> (),
    // wl_global
        fn wl_global_destroy(*mut wl_global) -> (),
    // wl_resource
        fn wl_resource_post_event_array(*mut wl_resource, u32, *mut wl_argument) -> (),
        fn wl_resource_queue_event_array(*mut wl_resource, u32, *mut wl_argument) -> (),
        fn wl_resource_post_no_memory(*mut wl_resource) -> (),
        fn wl_resource_set_implementation(*mut wl_resource, *const c_void, *mut c_void, Option<wl_resource_destroy_func_t>) -> (),
        fn wl_resource_set_dispatcher(*mut wl_resource, wl_dispatcher_func_t, *const c_void, *mut c_void, Option<wl_resource_destroy_func_t>) -> (),
        fn wl_resource_destroy(*mut wl_resource) -> (),
        fn wl_resource_get_client(*mut wl_resource) -> *mut wl_client,
        fn wl_resource_get_id(*mut wl_resource) -> u32,
        fn wl_resource_get_link(*mut wl_resource) -> *mut wl_list,
        fn wl_resource_from_link(*mut wl_list) -> *mut wl_resource,
        fn wl_resource_find_for_client(*mut wl_list, *mut wl_client) -> (),
        fn wl_resource_set_user_data(*mut wl_resource, *mut c_void) -> (),
        fn wl_resource_get_user_data(*mut wl_resource) -> *mut c_void,
        fn wl_resource_get_version(*mut wl_resource) -> c_int,
        fn wl_resource_set_destructor(*mut wl_resource, Option<wl_resource_destroy_func_t>) -> (),
        fn wl_resource_instance_of(*mut wl_resource, *const wl_interface, *const c_void) -> c_int,
        fn wl_resource_add_destroy_listener(*mut wl_resource, wl_notify_func_t) -> (),
        fn wl_resource_get_destroy_listener(*mut wl_resource,wl_notify_func_t) -> *mut wl_listener,
    // wl_shm
        fn wl_shm_buffer_begin_access(*mut wl_shm_buffer) -> (),
        fn wl_shm_buffer_end_access(*mut wl_shm_buffer) -> (),
        fn wl_shm_buffer_get(*mut wl_resource) -> *mut wl_shm_buffer,
        fn wl_shm_buffer_get_data(*mut wl_shm_buffer) -> *mut c_void,
        fn wl_shm_buffer_get_stride(*mut wl_shm_buffer) -> i32,
        fn wl_shm_buffer_get_format(*mut wl_shm_buffer) -> u32,
        fn wl_shm_buffer_get_width(*mut wl_shm_buffer) -> i32,
        fn wl_shm_buffer_get_height(*mut wl_shm_buffer) -> i32,
    // wl_log
        fn wl_log_set_handler_server(wl_log_func_t) -> (),
    // wl_list
        fn wl_list_init(*mut wl_list) -> (),
        fn wl_list_insert(*mut wl_list, *mut wl_list) -> (),
        fn wl_list_remove(*mut wl_list) -> (),
        fn wl_list_length(*const wl_list) -> c_int,
        fn wl_list_empty(*const wl_list) -> c_int,
        fn wl_list_insert_list(*mut wl_list,*mut wl_list) -> (),

    // arrays
        fn wl_array_init(*mut wl_array) -> (),
        fn wl_array_release(*mut wl_array) -> (),
        fn wl_array_add(*mut wl_array,usize) -> (),
        fn wl_array_copy(*mut wl_array, *mut wl_array) -> (),
    varargs:
        fn wl_resource_post_event(*mut wl_resource, u32) -> (),
        fn wl_resource_queue_event(*mut wl_resource, u32) -> (),
        fn wl_resource_post_error(*mut wl_resource, u32, *const c_char) -> (),
);

#[cfg(feature = "dlopen")]
lazy_static!(
    pub static ref WAYLAND_SERVER_OPTION: Option<WaylandServer> = {
        // This is a workaround for Ubuntu 17.04, which doesn't have a bare symlink
        // for libwayland-server.so but does have it with the version numbers for
        // whatever reason.
        //
        // We could do some trickery with str slices but that is more trouble
        // than its worth
        let versions = ["libwayland-server.so",
                        "libwayland-server.so.0"];
        for ver in &versions {
            match WaylandServer::open(ver) {
                Ok(h) => return Some(h),
                Err(::dlib::DlError::NotFound) => continue,
                Err(::dlib::DlError::MissingSymbol(s)) => {
                    if ::std::env::var_os("WAYLAND_RS_DEBUG").is_some() {
                        // only print debug messages if WAYLAND_RS_DEBUG is set
                        eprintln!("[wayland-server] Found library {} cannot be used: symbol {} is missing.", ver, s);
                    }
                    return None;
                }
            }
        }
        None
    };
    pub static ref WAYLAND_SERVER_HANDLE: &'static WaylandServer = {
        WAYLAND_SERVER_OPTION.as_ref().expect("Library libwayland-server.so could not be loaded.")
    };
);

#[cfg(not(feature = "dlopen"))]
pub fn is_lib_available() -> bool {
    true
}
#[cfg(feature = "dlopen")]
pub fn is_lib_available() -> bool {
    WAYLAND_SERVER_OPTION.is_some()
}

pub mod signal {
    #[cfg(not(feature = "dlopen"))]
    use super::{wl_list_init, wl_list_insert};
    use super::{wl_listener, wl_notify_func_t, wl_signal};
    #[cfg(feature = "dlopen")]
    use super::WAYLAND_SERVER_HANDLE as WSH;
    use common::wl_list;
    use std::os::raw::c_void;
    use std::ptr;

    // TODO: Is this really not UB ?
    macro_rules! offset_of(
        ($ty:ty, $field:ident) => {
            &(*(0 as *const $ty)).$field as *const _ as usize
        }
    );

    macro_rules! container_of(
        ($ptr: expr, $container: ty, $field: ident) => {
            ($ptr as *mut u8).offset(-(offset_of!($container, $field) as isize)) as *mut $container
        }
    );

    macro_rules! list_for_each(
        ($pos: ident, $head:expr, $container: ty, $field: ident, $action: block) => {
            let mut $pos = container_of!((*$head).next, $container, $field);
            while &mut (*$pos).$field as *mut _ != $head {
                $action;
                $pos = container_of!((*$pos).$field.next, $container, $field);
            }
        }
    );

    macro_rules! list_for_each_safe(
        ($pos: ident, $head: expr, $container: ty, $field: ident, $action: block) => {
            let mut $pos = container_of!((*$head).next, $container, $field);
            let mut tmp = container_of!((*$pos).$field.next, $container, $field);
            while &mut (*$pos).$field as *mut _ != $head {
                $action;
                $pos = tmp;
                tmp = container_of!((*$pos).$field.next, $container, $field);
            }
        }
    );

    pub unsafe fn wl_signal_init(signal: *mut wl_signal) {
        ffi_dispatch!(WSH, wl_list_init, &mut (*signal).listener_list);
    }

    pub unsafe fn wl_signal_add(signal: *mut wl_signal, listener: *mut wl_listener) {
        ffi_dispatch!(
            WSH,
            wl_list_insert,
            (*signal).listener_list.prev,
            &mut (*listener).link
        )
    }

    pub unsafe fn wl_signal_get(signal: *mut wl_signal, notify: wl_notify_func_t) -> *mut wl_listener {
        list_for_each!(
            l,
            &mut (*signal).listener_list as *mut wl_list,
            wl_listener,
            link,
            {
                if (*l).notify == notify {
                    return l;
                }
            }
        );
        return ptr::null_mut();
    }

    pub unsafe fn wl_signal_emit(signal: *mut wl_signal, data: *mut c_void) {
        list_for_each_safe!(
            l,
            &mut (*signal).listener_list as *mut wl_list,
            wl_listener,
            link,
            {
                ((*l).notify)(l, data);
            }
        );
    }

    #[repr(C)]
    struct ListenerWithUserData {
        listener: wl_listener,
        user_data: *mut c_void
    }

    pub fn rust_listener_create(notify: wl_notify_func_t) -> *mut wl_listener {
        let data = Box::into_raw(Box::new(ListenerWithUserData {
            listener: wl_listener {
                link: wl_list {
                    prev: ptr::null_mut(),
                    next: ptr::null_mut()
                },
                notify: notify
            },
            user_data: ptr::null_mut()
        }));
        return unsafe { &mut (*data).listener as *mut wl_listener }
    }

    pub unsafe fn rust_listener_get_user_data(listener: *mut wl_listener) -> *mut c_void {
        let data = container_of!(listener, ListenerWithUserData, listener);
        return (*data).user_data
    }

    pub unsafe fn rust_listener_set_user_data(listener: *mut wl_listener, user_data: *mut c_void) {
        let data = container_of!(listener, ListenerWithUserData, listener);
        (*data).user_data = user_data
    }

    pub unsafe fn rust_listener_destroy(listener: *mut wl_listener) {
        let data = container_of!(listener, ListenerWithUserData, listener);
        let _ = Box::from_raw(data);
    }
}
<|endoftext|>


<commit_before>
#![feature(plugin)]
#![plugin(clippy)]

#[macro_use]
extern crate kafka;
extern crate log;
extern crate openssl;
extern crate toml;

mod flowgger;

const DEFAULT_CONFIG_FILE: &'static str = "flowgger.toml";

fn main() {
    println!("Starting flowgger 0.1.2");
    let config_file = std::env::args().skip(1).next().unwrap_or(DEFAULT_CONFIG_FILE.to_owned());
    flowgger::start(&config_file);
}
<commit_msg>Make the version string a variable, at least.<commit_after>
#![feature(plugin)]
#![plugin(clippy)]

#[macro_use]
extern crate kafka;
extern crate log;
extern crate openssl;
extern crate toml;

mod flowgger;

const DEFAULT_CONFIG_FILE: &'static str = "flowgger.toml";
const FLOWGGER_VERSION_STRING: &'static str = "0.1.2";

fn main() {
    println!("Flowgger v{}", FLOWGGER_VERSION_STRING);
    let config_file = std::env::args().skip(1).next().unwrap_or(DEFAULT_CONFIG_FILE.to_owned());
    flowgger::start(&config_file);
}
<|endoftext|>