<commit_before>use futures::channel::mpsc as async_mpsc;

use gettextrs::gettext;

use gstreamer as gst;
use gstreamer::{prelude::*, ClockTime};

use glib;

use log::{debug, error, info, warn};

use std::{
    path::Path,
    sync::{Arc, Mutex},
};

use metadata::Format;

use super::{MediaEvent, Timestamp};

pub struct SplitterPipeline {
    pipeline: gst::Pipeline,
    format: Format,
    chapter: gst::TocEntry,
}

impl SplitterPipeline {
    pub fn check_requirements(format: Format) -> Result<(), String> {
        match format {
            Format::Flac => {
                gst::ElementFactory::make("flacenc", None)
                    .map(drop)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "flacenc", 1)
                    })
            }
            Format::Wave => {
                gst::ElementFactory::make("wavenc", None)
                    .map(drop)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "wavenc", 1)
                    })
            }
            Format::Opus => {
                gst::ElementFactory::make("opusenc", None)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "opusenc", 1)
                    })
                    .and_then(|_| {
                        gst::ElementFactory::make("oggmux", None)
                            .map(drop)
                            .map_err(|_| {
                                gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                                    .replacen("{element}", "oggmux", 1)
                            })
                    })
            }
            Format::Vorbis => {
                gst::ElementFactory::make("vorbisenc", None)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "vorbisenc", 1)
                    })
                    .and_then(|_| {
                        gst::ElementFactory::make("oggmux", None)
                            .map(drop)
                            .map_err(|_| {
                                gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                                    .replacen("{element}", "oggmux", 1)
                            })
                    })
            }
            Format::MP3 => {
                gst::ElementFactory::make("lamemp3enc", None)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "lamemp3enc", 1)
                    })
                    .and_then(|_| {
                        gst::ElementFactory::make("id3v2mux", None)
                            .map(drop)
                            .map_err(|_| {
                                gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                                    .replacen("{element}", "id3v2mux", 1)
                            })
                    })
            }
            _ => panic!(
                "SplitterPipeline::check_requirements unsupported format: {:?}",
                format
            ),
        }
    }

    pub fn try_new(
        input_path: &Path,
        output_path: &Path,
        stream_id: &str,
        format: Format,
        chapter: gst::TocEntry,
        sender: async_mpsc::Sender<MediaEvent>,
    ) -> Result<SplitterPipeline, String> {
        info!(
            "{}",
            gettext("Splitting {}...").replacen("{}", output_path.to_str().unwrap(), 1)
        );
        debug!("stream id {}", &stream_id);

        let mut this = SplitterPipeline {
            pipeline: gst::Pipeline::new(Some("splitter_pipeline")),
            format,
            chapter,
        };

        this.build_pipeline(input_path, output_path, stream_id);
        this.register_bus_inspector(sender);

        this.pipeline
            .set_state(gst::State::Paused)
            .map(|_| this)
            .map_err(|_| gettext("do you have permission to write the file?"))
    }

    pub fn get_current_ts(&self) -> Option<Timestamp> {
        // `get_current_ts` might be called when the pipeline is not ready yet
        // because we build a new pipeline for each split file.
        // When this happen, the ts position is negative so don't force it to 0
        // because we are most likely not at the begining of the source file.
        let mut position_query = gst::Query::new_position(gst::Format::Time);
        self.pipeline.query(&mut position_query);
        let position = position_query.get_result().get_value();
        if position >= 0 {
            Some(position.into())
        } else {
            None
        }
    }

    fn build_pipeline(&mut self, input_path: &Path, output_path: &Path, stream_id: &str) {
        /* There are multiple showstoppers to implementing something ideal
         * to export splitted chapters with audio and video (and subtitles):
         * 1. matroska-mux drops seek events explicitly (a message states: "discard for now").
         * This means that it is currently not possible to build a pipeline that would allow
         * seeking in a matroska media (using demux to interprete timestamps) and mux back to
         * to matroska. One solution would be to export streams to files and mux back
         * crossing fingers everything remains in sync.
         * 2. nlesrc (from gstreamer-editing-services) allows extracting frames from a starting
         * position for a given duration. However, it is designed to work with single stream
         * medias and decodes to raw formats.
         * 3. filesink can't change the file location without setting the pipeline to Null
         * which also unlinks elements.
         * 4. wavenc doesn't send header after a second seek in segment mode.
         * 5. flacenc can't handle a seek.
         *
         * Issues 3 & 4 lead to building a new pipeline for each chapter.
         *
         * Until I design a GUI for the user to select which stream to export to which codec,
         * current solution is to keep only the audio track which matches the initial purpose
         * of this application */

        // Input
        let filesrc = gst::ElementFactory::make("filesrc", None).unwrap();
        filesrc
            .set_property("location", &glib::Value::from(input_path.to_str().unwrap()))
            .unwrap();
        let decodebin = gst::ElementFactory::make("decodebin", None).unwrap();

        self.pipeline.add_many(&[&filesrc, &decodebin]).unwrap();

        filesrc.link(&decodebin).unwrap();
        decodebin.sync_state_with_parent().unwrap();

        // Audio encoder
        let audio_enc = match self.format {
            Format::Flac => gst::ElementFactory::make("flacenc", None).unwrap(),
            Format::Wave => gst::ElementFactory::make("wavenc", None).unwrap(),
            Format::Opus => gst::ElementFactory::make("opusenc", None).unwrap(),
            Format::Vorbis => gst::ElementFactory::make("vorbisenc", None).unwrap(),
            Format::MP3 => gst::ElementFactory::make("lamemp3enc", None).unwrap(),
            _ => panic!(
                "SplitterPipeline::build_pipeline unsupported format: {:?}",
                self.format
            ),
        };

        // Catch events and drop the upstream Tags & TOC
        let audio_enc_sink_pad = audio_enc.get_static_pad("sink").unwrap();
        audio_enc_sink_pad.add_probe(gst::PadProbeType::EVENT_DOWNSTREAM, |_pad, probe_info| {
            if let Some(ref data) = probe_info.data {
                if let gst::PadProbeData::Event(ref event) = *data {
                    match event.view() {
                        gst::EventView::Tag(ref _tag) => return gst::PadProbeReturn::Drop,
                        gst::EventView::Toc(ref _toc) => return gst::PadProbeReturn::Drop,
                        _ => (),
                    }
                }
            }
            gst::PadProbeReturn::Ok
        });

        // Some encoders (such as flacenc) starts encoding the preroll buffers when switching to
        // paused mode. When the seek is performed buffers from the new segments are appended
        // to the ones from the preroll. Moreover, flacenc doesn't handle discontinuities.
        // As a workaround, we will drop buffers before the seek. The first buffer with the Discont
        // flag shows that it is possible to seek. The next buffer with the Discont flag corresponds
        // to the first buffer from the target segment

        let (start, end) = self.chapter.get_start_stop_times().unwrap();

        // Note: can't use AtomicBool here as pad probes are multithreaded so the function is Fn
        // not FnMut. See: https://github.com/sdroege/gstreamer-rs/pull/71
        #[allow(clippy::mutex_atomic)]
        let seek_done = Arc::new(Mutex::new(false));
        let pipeline = self.pipeline.clone();
        audio_enc_sink_pad.add_probe(gst::PadProbeType::BUFFER, move |_pad, probe_info| {
            if let Some(ref data) = probe_info.data {
                if let gst::PadProbeData::Buffer(ref buffer) = *data {
                    if buffer.get_flags() & gst::BufferFlags::DISCONT == gst::BufferFlags::DISCONT {
                        let mut seek_done_grp = seek_done.lock().unwrap();
                        if !*seek_done_grp {
                            // First buffer before seek
                            // let's seek and drop buffers until seek start sending new segment
                            let _res = pipeline
                                .seek(
                                    1f64,
                                    gst::SeekFlags::FLUSH | gst::SeekFlags::ACCURATE,
                                    gst::SeekType::Set,
                                    ClockTime::from(start as u64),
                                    gst::SeekType::Set,
                                    ClockTime::from(end as u64),
                                )
                                .map_err(|_| {
                                    // FIXME: feedback to the user using the UI channel
                                    error!("{}", gettext("Failed to intialize the split"));
                                });
                            *seek_done_grp = true;
                        } else {
                            // First Discont buffer after seek => stop dropping buffers
                            return gst::PadProbeReturn::Remove;
                        }
                    }
                }
            }
            gst::PadProbeReturn::Drop
        });

        self.pipeline.add(&audio_enc).unwrap();

        // add a muxer when required
        let (tag_setter, audio_muxer) = match self.format {
            Format::Flac | Format::Wave => (audio_enc.clone(), audio_enc.clone()),
            Format::Opus | Format::Vorbis => {
                let ogg_muxer = gst::ElementFactory::make("oggmux", None).unwrap();
                self.pipeline.add(&ogg_muxer).unwrap();
                audio_enc.link(&ogg_muxer).unwrap();
                (audio_enc.clone(), ogg_muxer)
            }
            Format::MP3 => {
                let id3v2_muxer = gst::ElementFactory::make("id3v2mux", None).unwrap();
                self.pipeline.add(&id3v2_muxer).unwrap();
                audio_enc.link(&id3v2_muxer).unwrap();
                (id3v2_muxer.clone(), id3v2_muxer)
            }
            _ => unimplemented!(
                "SplitterPipeline::build_pipeline for format: {:?}",
                self.format
            ),
        };

        if let Some(tags) = self.chapter.get_tags() {
            let tag_setter = tag_setter.dynamic_cast::<gst::TagSetter>().unwrap();
            tag_setter.merge_tags(&tags, gst::TagMergeMode::ReplaceAll);
        }

        // Output sink
        let outsink = gst::ElementFactory::make("filesink", Some("filesink")).unwrap();
        outsink
            .set_property(
                "location",
                &glib::Value::from(output_path.to_str().unwrap()),
            )
            .unwrap();

        self.pipeline.add(&outsink).unwrap();
        audio_muxer.link(&outsink).unwrap();
        outsink.sync_state_with_parent().unwrap();

        let pipeline_cb = self.pipeline.clone();
        let stream_id = stream_id.to_owned();
        decodebin.connect_pad_added(move |_element, pad| {
            let caps = pad.get_current_caps().unwrap();
            let structure = caps.get_structure(0).unwrap();
            let name = structure.get_name();

            let queue = gst::ElementFactory::make("queue", None).unwrap();
            pipeline_cb.add(&queue).unwrap();
            let queue_sink_pad = queue.get_static_pad("sink").unwrap();
            pad.link(&queue_sink_pad).unwrap();
            queue.sync_state_with_parent().unwrap();
            let queue_src_pad = queue.get_static_pad("src").unwrap();

            if name.starts_with("audio/")
                && pipeline_cb.get_by_name("audioconvert").is_none()
                && stream_id
                    == pad
                        .get_stream_id()
                        .expect("SplitterPipeline::build_pipeline no stream_id for audio src pad")
            {
                let audio_conv =
                    gst::ElementFactory::make("audioconvert", Some("audioconvert")).unwrap();
                pipeline_cb.add(&audio_conv).unwrap();
                gst::Element::link_many(&[&queue, &audio_conv, &audio_enc]).unwrap();
                audio_conv.sync_state_with_parent().unwrap();
                audio_enc.sync_state_with_parent().unwrap();
            } else {
                let fakesink = gst::ElementFactory::make("fakesink", None).unwrap();
                pipeline_cb.add(&fakesink).unwrap();
                let fakesink_sink_pad = fakesink.get_static_pad("sink").unwrap();
                queue_src_pad.link(&fakesink_sink_pad).unwrap();
                fakesink.sync_state_with_parent().unwrap();
            }
        });
    }

    pub fn cancel(&self) {
        if self.pipeline.set_state(gst::State::Null).is_err() {
            warn!("could not stop the media");
        }
    }

    // Uses sender to notify the UI controllers
    fn register_bus_inspector(&self, mut sender: async_mpsc::Sender<MediaEvent>) {
        let pipeline = self.pipeline.clone();
        self.pipeline
            .get_bus()
            .unwrap()
            .add_watch(move |_, msg| {
                match msg.view() {
                    gst::MessageView::Eos(..) => {
                        if pipeline.set_state(gst::State::Null).is_err() {
                            sender
                                .try_send(MediaEvent::FailedToExport(gettext(
                                    "Failed to terminate properly. Check the resulting file.",
                                )))
                                .unwrap();
                        }
                        sender.try_send(MediaEvent::Eos).unwrap();
                        return glib::Continue(false);
                    }
                    gst::MessageView::Error(err) => {
                        let _ = sender
                            .try_send(MediaEvent::FailedToExport(err.get_error().to_string()));
                        return glib::Continue(false);
                    }
                    gst::MessageView::AsyncDone(_) => {
                        // Start splitting
                        if pipeline.set_state(gst::State::Playing).is_err() {
                            sender
                                .try_send(MediaEvent::FailedToExport(gettext(
                                    "Failed to start splitting.",
                                )))
                                .unwrap();
                        }
                    }
                    _ => (),
                }

                glib::Continue(true)
            })
            .unwrap();
    }
}
<commit_msg>Split pipeline: remove unneeded queue<commit_after>use futures::channel::mpsc as async_mpsc;

use gettextrs::gettext;

use gstreamer as gst;
use gstreamer::{prelude::*, ClockTime};

use glib;

use log::{debug, error, info, warn};

use std::{
    path::Path,
    sync::{Arc, Mutex},
};

use metadata::Format;

use super::{MediaEvent, Timestamp};

pub struct SplitterPipeline {
    pipeline: gst::Pipeline,
    format: Format,
    chapter: gst::TocEntry,
}

impl SplitterPipeline {
    pub fn check_requirements(format: Format) -> Result<(), String> {
        match format {
            Format::Flac => {
                gst::ElementFactory::make("flacenc", None)
                    .map(drop)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "flacenc", 1)
                    })
            }
            Format::Wave => {
                gst::ElementFactory::make("wavenc", None)
                    .map(drop)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "wavenc", 1)
                    })
            }
            Format::Opus => {
                gst::ElementFactory::make("opusenc", None)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "opusenc", 1)
                    })
                    .and_then(|_| {
                        gst::ElementFactory::make("oggmux", None)
                            .map(drop)
                            .map_err(|_| {
                                gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                                    .replacen("{element}", "oggmux", 1)
                            })
                    })
            }
            Format::Vorbis => {
                gst::ElementFactory::make("vorbisenc", None)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "vorbisenc", 1)
                    })
                    .and_then(|_| {
                        gst::ElementFactory::make("oggmux", None)
                            .map(drop)
                            .map_err(|_| {
                                gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                                    .replacen("{element}", "oggmux", 1)
                            })
                    })
            }
            Format::MP3 => {
                gst::ElementFactory::make("lamemp3enc", None)
                    .map_err(|_| {
                        gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                            .replacen("{element}", "lamemp3enc", 1)
                    })
                    .and_then(|_| {
                        gst::ElementFactory::make("id3v2mux", None)
                            .map(drop)
                            .map_err(|_| {
                                gettext("Missing `{element}`\ncheck your gst-plugins-good install")
                                    .replacen("{element}", "id3v2mux", 1)
                            })
                    })
            }
            _ => panic!(
                "SplitterPipeline::check_requirements unsupported format: {:?}",
                format
            ),
        }
    }

    pub fn try_new(
        input_path: &Path,
        output_path: &Path,
        stream_id: &str,
        format: Format,
        chapter: gst::TocEntry,
        sender: async_mpsc::Sender<MediaEvent>,
    ) -> Result<SplitterPipeline, String> {
        info!(
            "{}",
            gettext("Splitting {}...").replacen("{}", output_path.to_str().unwrap(), 1)
        );
        debug!("stream id {}", &stream_id);

        let mut this = SplitterPipeline {
            pipeline: gst::Pipeline::new(Some("splitter_pipeline")),
            format,
            chapter,
        };

        this.build_pipeline(input_path, output_path, stream_id);
        this.register_bus_inspector(sender);

        this.pipeline
            .set_state(gst::State::Paused)
            .map(|_| this)
            .map_err(|_| gettext("do you have permission to write the file?"))
    }

    pub fn get_current_ts(&self) -> Option<Timestamp> {
        // `get_current_ts` might be called when the pipeline is not ready yet
        // because we build a new pipeline for each split file.
        // When this happen, the ts position is negative so don't force it to 0
        // because we are most likely not at the begining of the source file.
        let mut position_query = gst::Query::new_position(gst::Format::Time);
        self.pipeline.query(&mut position_query);
        let position = position_query.get_result().get_value();
        if position >= 0 {
            Some(position.into())
        } else {
            None
        }
    }

    fn build_pipeline(&mut self, input_path: &Path, output_path: &Path, stream_id: &str) {
        /* There are multiple showstoppers to implementing something ideal
         * to export splitted chapters with audio and video (and subtitles):
         * 1. matroska-mux drops seek events explicitly (a message states: "discard for now").
         * This means that it is currently not possible to build a pipeline that would allow
         * seeking in a matroska media (using demux to interprete timestamps) and mux back to
         * to matroska. One solution would be to export streams to files and mux back
         * crossing fingers everything remains in sync.
         * 2. nlesrc (from gstreamer-editing-services) allows extracting frames from a starting
         * position for a given duration. However, it is designed to work with single stream
         * medias and decodes to raw formats.
         * 3. filesink can't change the file location without setting the pipeline to Null
         * which also unlinks elements.
         * 4. wavenc doesn't send header after a second seek in segment mode.
         * 5. flacenc can't handle a seek.
         *
         * Issues 3 & 4 lead to building a new pipeline for each chapter.
         *
         * Until I design a GUI for the user to select which stream to export to which codec,
         * current solution is to keep only the audio track which matches the initial purpose
         * of this application */

        // Input
        let filesrc = gst::ElementFactory::make("filesrc", None).unwrap();
        filesrc
            .set_property("location", &glib::Value::from(input_path.to_str().unwrap()))
            .unwrap();
        let decodebin = gst::ElementFactory::make("decodebin", None).unwrap();

        self.pipeline.add_many(&[&filesrc, &decodebin]).unwrap();

        filesrc.link(&decodebin).unwrap();
        decodebin.sync_state_with_parent().unwrap();

        // Audio encoder
        let audio_enc = match self.format {
            Format::Flac => gst::ElementFactory::make("flacenc", None).unwrap(),
            Format::Wave => gst::ElementFactory::make("wavenc", None).unwrap(),
            Format::Opus => gst::ElementFactory::make("opusenc", None).unwrap(),
            Format::Vorbis => gst::ElementFactory::make("vorbisenc", None).unwrap(),
            Format::MP3 => gst::ElementFactory::make("lamemp3enc", None).unwrap(),
            _ => panic!(
                "SplitterPipeline::build_pipeline unsupported format: {:?}",
                self.format
            ),
        };

        // Catch events and drop the upstream Tags & TOC
        let audio_enc_sink_pad = audio_enc.get_static_pad("sink").unwrap();
        audio_enc_sink_pad.add_probe(gst::PadProbeType::EVENT_DOWNSTREAM, |_pad, probe_info| {
            if let Some(ref data) = probe_info.data {
                if let gst::PadProbeData::Event(ref event) = *data {
                    match event.view() {
                        gst::EventView::Tag(ref _tag) => return gst::PadProbeReturn::Drop,
                        gst::EventView::Toc(ref _toc) => return gst::PadProbeReturn::Drop,
                        _ => (),
                    }
                }
            }
            gst::PadProbeReturn::Ok
        });

        // Some encoders (such as flacenc) starts encoding the preroll buffers when switching to
        // paused mode. When the seek is performed buffers from the new segments are appended
        // to the ones from the preroll. Moreover, flacenc doesn't handle discontinuities.
        // As a workaround, we will drop buffers before the seek. The first buffer with the Discont
        // flag shows that it is possible to seek. The next buffer with the Discont flag corresponds
        // to the first buffer from the target segment

        let (start, end) = self.chapter.get_start_stop_times().unwrap();

        // Note: can't use AtomicBool here as pad probes are multithreaded so the function is Fn
        // not FnMut. See: https://github.com/sdroege/gstreamer-rs/pull/71
        #[allow(clippy::mutex_atomic)]
        let seek_done = Arc::new(Mutex::new(false));
        let pipeline = self.pipeline.clone();
        audio_enc_sink_pad.add_probe(gst::PadProbeType::BUFFER, move |_pad, probe_info| {
            if let Some(ref data) = probe_info.data {
                if let gst::PadProbeData::Buffer(ref buffer) = *data {
                    if buffer.get_flags() & gst::BufferFlags::DISCONT == gst::BufferFlags::DISCONT {
                        let mut seek_done_grp = seek_done.lock().unwrap();
                        if !*seek_done_grp {
                            // First buffer before seek
                            // let's seek and drop buffers until seek start sending new segment
                            let _res = pipeline
                                .seek(
                                    1f64,
                                    gst::SeekFlags::FLUSH | gst::SeekFlags::ACCURATE,
                                    gst::SeekType::Set,
                                    ClockTime::from(start as u64),
                                    gst::SeekType::Set,
                                    ClockTime::from(end as u64),
                                )
                                .map_err(|_| {
                                    // FIXME: feedback to the user using the UI channel
                                    error!("{}", gettext("Failed to intialize the split"));
                                });
                            *seek_done_grp = true;
                        } else {
                            // First Discont buffer after seek => stop dropping buffers
                            return gst::PadProbeReturn::Remove;
                        }
                    }
                }
            }
            gst::PadProbeReturn::Drop
        });

        self.pipeline.add(&audio_enc).unwrap();

        // add a muxer when required
        let (tag_setter, audio_muxer) = match self.format {
            Format::Flac | Format::Wave => (audio_enc.clone(), audio_enc.clone()),
            Format::Opus | Format::Vorbis => {
                let ogg_muxer = gst::ElementFactory::make("oggmux", None).unwrap();
                self.pipeline.add(&ogg_muxer).unwrap();
                audio_enc.link(&ogg_muxer).unwrap();
                (audio_enc.clone(), ogg_muxer)
            }
            Format::MP3 => {
                let id3v2_muxer = gst::ElementFactory::make("id3v2mux", None).unwrap();
                self.pipeline.add(&id3v2_muxer).unwrap();
                audio_enc.link(&id3v2_muxer).unwrap();
                (id3v2_muxer.clone(), id3v2_muxer)
            }
            _ => unimplemented!(
                "SplitterPipeline::build_pipeline for format: {:?}",
                self.format
            ),
        };

        if let Some(tags) = self.chapter.get_tags() {
            let tag_setter = tag_setter.dynamic_cast::<gst::TagSetter>().unwrap();
            tag_setter.merge_tags(&tags, gst::TagMergeMode::ReplaceAll);
        }

        // Output sink
        let outsink = gst::ElementFactory::make("filesink", Some("filesink")).unwrap();
        outsink
            .set_property(
                "location",
                &glib::Value::from(output_path.to_str().unwrap()),
            )
            .unwrap();

        self.pipeline.add(&outsink).unwrap();
        audio_muxer.link(&outsink).unwrap();
        outsink.sync_state_with_parent().unwrap();

        let pipeline_cb = self.pipeline.clone();
        let stream_id = stream_id.to_owned();
        decodebin.connect_pad_added(move |_element, pad| {
            let caps = pad.get_current_caps().unwrap();
            let structure = caps.get_structure(0).unwrap();
            let name = structure.get_name();

            if name.starts_with("audio/")
                && pipeline_cb.get_by_name("audioconvert").is_none()
                && stream_id
                    == pad
                        .get_stream_id()
                        .expect("SplitterPipeline::build_pipeline no stream_id for audio src pad")
            {
                let audio_conv =
                    gst::ElementFactory::make("audioconvert", Some("audioconvert")).unwrap();
                pipeline_cb.add(&audio_conv).unwrap();
                let audio_conv_sink_pad = audio_conv.get_static_pad("sink").unwrap();
                pad.link(&audio_conv_sink_pad).unwrap();
                gst::Element::link_many(&[&audio_conv, &audio_enc]).unwrap();
                audio_conv.sync_state_with_parent().unwrap();
                audio_enc.sync_state_with_parent().unwrap();
            } else {
                let fakesink = gst::ElementFactory::make("fakesink", None).unwrap();
                pipeline_cb.add(&fakesink).unwrap();
                let fakesink_sink_pad = fakesink.get_static_pad("sink").unwrap();
                pad.link(&fakesink_sink_pad).unwrap();
                fakesink.sync_state_with_parent().unwrap();
            }
        });
    }

    pub fn cancel(&self) {
        if self.pipeline.set_state(gst::State::Null).is_err() {
            warn!("could not stop the media");
        }
    }

    // Uses sender to notify the UI controllers
    fn register_bus_inspector(&self, mut sender: async_mpsc::Sender<MediaEvent>) {
        let pipeline = self.pipeline.clone();
        self.pipeline
            .get_bus()
            .unwrap()
            .add_watch(move |_, msg| {
                match msg.view() {
                    gst::MessageView::Eos(..) => {
                        if pipeline.set_state(gst::State::Null).is_err() {
                            sender
                                .try_send(MediaEvent::FailedToExport(gettext(
                                    "Failed to terminate properly. Check the resulting file.",
                                )))
                                .unwrap();
                        }
                        sender.try_send(MediaEvent::Eos).unwrap();
                        return glib::Continue(false);
                    }
                    gst::MessageView::Error(err) => {
                        let _ = sender
                            .try_send(MediaEvent::FailedToExport(err.get_error().to_string()));
                        return glib::Continue(false);
                    }
                    gst::MessageView::AsyncDone(_) => {
                        // Start splitting
                        if pipeline.set_state(gst::State::Playing).is_err() {
                            sender
                                .try_send(MediaEvent::FailedToExport(gettext(
                                    "Failed to start splitting.",
                                )))
                                .unwrap();
                        }
                    }
                    _ => (),
                }

                glib::Continue(true)
            })
            .unwrap();
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2020 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/// Common utilities for testing
use std::{
    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
    str::from_utf8,
    sync::Arc,
};

use tokio::net::UdpSocket;
use tokio::sync::{mpsc, oneshot, watch};

use crate::{
    config::{Builder as ConfigBuilder, Config},
    endpoint::{Endpoint, EndpointAddress, Endpoints},
    filters::{prelude::*, FilterRegistry},
    metadata::Value,
};

// TestFilter is useful for testing that commands are executing filters appropriately.
pub struct TestFilter;

impl Filter for TestFilter {
    fn read(&self, mut ctx: ReadContext) -> Option<ReadResponse> {
        // append values on each run
        ctx.metadata
            .entry(Arc::new("downstream".into()))
            .and_modify(|e| e.as_mut_string().unwrap().push_str(":receive"))
            .or_insert_with(|| Value::String("receive".into()));

        ctx.contents
            .append(&mut format!(":odr:{}", ctx.source).into_bytes());
        Some(ctx.into())
    }

    fn write(&self, mut ctx: WriteContext) -> Option<WriteResponse> {
        // append values on each run
        ctx.metadata
            .entry("upstream".to_string().into())
            .and_modify(|e| e.as_mut_string().unwrap().push_str(":receive"))
            .or_insert_with(|| Value::String("receive".to_string()));

        ctx.contents
            .append(&mut format!(":our:{}:{}", ctx.source, ctx.dest).into_bytes());
        Some(ctx.into())
    }
}

impl StaticFilter for TestFilter {
    const NAME: &'static str = "TestFilter";
    type Configuration = ();
    type BinaryConfiguration = ();

    fn try_from_config(_: Option<Self::Configuration>) -> Result<Self, Error> {
        Ok(Self)
    }
}

#[derive(Default)]
pub struct TestHelper {
    /// Channel to subscribe to, and trigger the shutdown of created resources.
    shutdown_ch: Option<(watch::Sender<()>, watch::Receiver<()>)>,
    server_shutdown_tx: Vec<Option<watch::Sender<()>>>,
}

/// Returned from [creating a socket](TestHelper::open_socket_and_recv_single_packet)
pub struct OpenSocketRecvPacket {
    /// The opened socket
    pub socket: Arc<UdpSocket>,
    /// A channel on which the received packet will be forwarded.
    pub packet_rx: oneshot::Receiver<String>,
}

impl Drop for TestHelper {
    fn drop(&mut self) {
        for shutdown_tx in self.server_shutdown_tx.iter_mut().flat_map(|tx| tx.take()) {
            shutdown_tx
                .send(())
                .map_err(|error| {
                    tracing::warn!(
                        %error,
                        "Failed to send server shutdown over channel"
                    )
                })
                .ok();
        }

        if let Some((shutdown_tx, _)) = self.shutdown_ch.take() {
            shutdown_tx.send(()).unwrap();
        }
    }
}

impl TestHelper {
    /// Opens a new socket bound to an ephemeral port
    pub async fn create_socket(&self) -> Arc<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        Arc::new(UdpSocket::bind(addr).await.unwrap())
    }

    /// Opens a socket, listening for a packet. Once a packet is received, it
    /// is forwarded over the returned channel.
    pub async fn open_socket_and_recv_single_packet(&self) -> OpenSocketRecvPacket {
        let socket = self.create_socket().await;
        let (packet_tx, packet_rx) = oneshot::channel::<String>();
        let socket_recv = socket.clone();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = socket_recv.recv(&mut buf).await.unwrap();
            packet_tx
                .send(from_utf8(&buf[..size]).unwrap().to_string())
                .unwrap();
        });
        OpenSocketRecvPacket { socket, packet_rx }
    }

    /// Opens a socket, listening for packets. Received packets are forwarded over the
    /// returned channel.
    pub async fn open_socket_and_recv_multiple_packets(
        &mut self,
    ) -> (mpsc::Receiver<String>, Arc<UdpSocket>) {
        let (packet_tx, packet_rx) = mpsc::channel::<String>(10);
        let socket = self.create_socket().await;
        let mut shutdown_rx = self.get_shutdown_subscriber().await;
        let socket_recv = socket.clone();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            loop {
                tokio::select! {
                    received = socket_recv.recv_from(&mut buf) => {
                        let (size, _) = received.unwrap();
                        let str = from_utf8(&buf[..size]).unwrap().to_string();
                        match packet_tx.send(str).await {
                            Ok(_) => {}
                            Err(error) => {
                                tracing::warn!(target: "recv_multiple_packets", %error, "recv_chan dropped");
                                return;
                            }
                        };
                    },
                    _ = shutdown_rx.changed() => {
                        return;
                    }
                }
            }
        });
        (packet_rx, socket)
    }

    /// Runs a simple UDP server that echos back payloads.
    /// Returns the server's address.
    pub async fn run_echo_server(&mut self) -> EndpointAddress {
        self.run_echo_server_with_tap(|_, _, _| {}).await
    }

    /// Runs a simple UDP server that echos back payloads.
    /// The provided function is invoked for each received payload.
    /// Returns the server's address.
    pub async fn run_echo_server_with_tap<F>(&mut self, tap: F) -> EndpointAddress
    where
        F: Fn(SocketAddr, &[u8], SocketAddr) + Send + 'static,
    {
        let socket = self.create_socket().await;
        let addr = socket.local_addr().unwrap();
        let mut shutdown = self.get_shutdown_subscriber().await;
        let local_addr = addr;
        tokio::spawn(async move {
            loop {
                let mut buf = vec![0; 1024];
                tokio::select! {
                    recvd = socket.recv_from(&mut buf) => {
                        let (size, sender) = recvd.unwrap();
                        let packet = &buf[..size];
                        tap(sender, packet, local_addr);
                        socket.send_to(packet, sender).await.unwrap();
                    },
                    _ = shutdown.changed() => {
                        return;
                    }
                }
            }
        });
        addr.into()
    }

    /// Run a proxy server with a supplied config.
    /// Admin is disabled for this method, as the majority of tests will not need it, and it makes it
    /// easier to avoid issues with port collisions.
    pub fn run_server_with_config(&mut self, mut config: Config) {
        config.admin = None;

        self.run_server(<_>::try_from(config).unwrap());
    }

    pub fn run_server(&mut self, server: crate::Server) {
        let (shutdown_tx, shutdown_rx) = watch::channel::<()>(());
        self.server_shutdown_tx.push(Some(shutdown_tx));
        tokio::spawn(async move {
            server.run(shutdown_rx).await.unwrap();
        });
    }

    /// Returns a receiver subscribed to the helper's shutdown event.
    async fn get_shutdown_subscriber(&mut self) -> watch::Receiver<()> {
        // If this is the first call, then we set up the channel first.
        match self.shutdown_ch {
            Some((_, ref rx)) => rx.clone(),
            None => {
                let ch = watch::channel(());
                let recv = ch.1.clone();
                self.shutdown_ch = Some(ch);
                recv
            }
        }
    }
}

/// assert that read makes no changes
pub fn assert_filter_read_no_change<F>(filter: &F)
where
    F: Filter,
{
    let endpoints = vec!["127.0.0.1:80".parse::<Endpoint>().unwrap()];
    let source = "127.0.0.1:90".parse().unwrap();
    let contents = "hello".to_string().into_bytes();

    match filter.read(ReadContext::new(
        Endpoints::new(endpoints.clone()).into(),
        source,
        contents.clone(),
    )) {
        None => unreachable!("should return a result"),
        Some(response) => {
            assert_eq!(
                endpoints,
                response.endpoints.iter().cloned().collect::<Vec<_>>()
            );
            assert_eq!(contents, response.contents);
        }
    }
}

/// assert that write makes no changes
pub fn assert_write_no_change<F>(filter: &F)
where
    F: Filter,
{
    let endpoint = "127.0.0.1:90".parse::<Endpoint>().unwrap();
    let contents = "hello".to_string().into_bytes();

    match filter.write(WriteContext::new(
        &endpoint,
        endpoint.address.clone(),
        "127.0.0.1:70".parse().unwrap(),
        contents.clone(),
    )) {
        None => unreachable!("should return a result"),
        Some(response) => assert_eq!(contents, response.contents),
    }
}

pub fn config_with_dummy_endpoint() -> ConfigBuilder {
    ConfigBuilder::default().endpoints(vec![Endpoint {
        address: "127.0.0.1:8080".parse().unwrap(),
        ..<_>::default()
    }])
}
/// Creates a dummy endpoint with `id` as a suffix.
pub fn ep(id: u8) -> Endpoint {
    Endpoint {
        address: ([127, 0, 0, id], 8080).into(),
        ..<_>::default()
    }
}

pub fn new_test_chain() -> crate::filters::SharedFilterChain {
    <_>::try_from([crate::config::Filter {
        name: "TestFilter".into(),
        config: None,
    }])
    .unwrap()
}

pub fn load_test_filters() {
    FilterRegistry::register([TestFilter::factory()]);
}

#[cfg(test)]
mod tests {
    use crate::test_utils::TestHelper;

    #[tokio::test]
    async fn test_echo_server() {
        let mut t = TestHelper::default();
        let echo_addr = t.run_echo_server().await;
        let endpoint = t.open_socket_and_recv_single_packet().await;
        let msg = "hello";
        endpoint
            .socket
            .send_to(msg.as_bytes(), &echo_addr.to_socket_addr().unwrap())
            .await
            .unwrap();
        assert_eq!(msg, endpoint.packet_rx.await.unwrap());
    }
}
<commit_msg>Test Utility to enable tracing logging. (#537)<commit_after>/*
 * Copyright 2020 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/// Common utilities for testing
use std::{
    net::{Ipv4Addr, SocketAddr, SocketAddrV4},
    str::from_utf8,
    sync::Arc,
};

use once_cell::sync::Lazy;
use tokio::net::UdpSocket;
use tokio::sync::{mpsc, oneshot, watch};

use crate::{
    config::{Builder as ConfigBuilder, Config},
    endpoint::{Endpoint, EndpointAddress, Endpoints},
    filters::{prelude::*, FilterRegistry},
    metadata::Value,
};

static ENABLE_LOG: Lazy<()> = Lazy::new(|| {
    tracing_subscriber::fmt()
        .pretty()
        .with_max_level(tracing::Level::TRACE)
        .init()
});

/// Call to safely enable `tracing` logging calls, at the TRACE level.
/// This can be very useful when attempting to debug unit and integration tests.
pub fn enable_tracing_log() {
    Lazy::force(&ENABLE_LOG);
}

// TestFilter is useful for testing that commands are executing filters appropriately.
pub struct TestFilter;

impl Filter for TestFilter {
    fn read(&self, mut ctx: ReadContext) -> Option<ReadResponse> {
        // append values on each run
        ctx.metadata
            .entry(Arc::new("downstream".into()))
            .and_modify(|e| e.as_mut_string().unwrap().push_str(":receive"))
            .or_insert_with(|| Value::String("receive".into()));

        ctx.contents
            .append(&mut format!(":odr:{}", ctx.source).into_bytes());
        Some(ctx.into())
    }

    fn write(&self, mut ctx: WriteContext) -> Option<WriteResponse> {
        // append values on each run
        ctx.metadata
            .entry("upstream".to_string().into())
            .and_modify(|e| e.as_mut_string().unwrap().push_str(":receive"))
            .or_insert_with(|| Value::String("receive".to_string()));

        ctx.contents
            .append(&mut format!(":our:{}:{}", ctx.source, ctx.dest).into_bytes());
        Some(ctx.into())
    }
}

impl StaticFilter for TestFilter {
    const NAME: &'static str = "TestFilter";
    type Configuration = ();
    type BinaryConfiguration = ();

    fn try_from_config(_: Option<Self::Configuration>) -> Result<Self, Error> {
        Ok(Self)
    }
}

#[derive(Default)]
pub struct TestHelper {
    /// Channel to subscribe to, and trigger the shutdown of created resources.
    shutdown_ch: Option<(watch::Sender<()>, watch::Receiver<()>)>,
    server_shutdown_tx: Vec<Option<watch::Sender<()>>>,
}

/// Returned from [creating a socket](TestHelper::open_socket_and_recv_single_packet)
pub struct OpenSocketRecvPacket {
    /// The opened socket
    pub socket: Arc<UdpSocket>,
    /// A channel on which the received packet will be forwarded.
    pub packet_rx: oneshot::Receiver<String>,
}

impl Drop for TestHelper {
    fn drop(&mut self) {
        for shutdown_tx in self.server_shutdown_tx.iter_mut().flat_map(|tx| tx.take()) {
            shutdown_tx
                .send(())
                .map_err(|error| {
                    tracing::warn!(
                        %error,
                        "Failed to send server shutdown over channel"
                    )
                })
                .ok();
        }

        if let Some((shutdown_tx, _)) = self.shutdown_ch.take() {
            shutdown_tx.send(()).unwrap();
        }
    }
}

impl TestHelper {
    /// Opens a new socket bound to an ephemeral port
    pub async fn create_socket(&self) -> Arc<UdpSocket> {
        let addr = SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0);
        Arc::new(UdpSocket::bind(addr).await.unwrap())
    }

    /// Opens a socket, listening for a packet. Once a packet is received, it
    /// is forwarded over the returned channel.
    pub async fn open_socket_and_recv_single_packet(&self) -> OpenSocketRecvPacket {
        let socket = self.create_socket().await;
        let (packet_tx, packet_rx) = oneshot::channel::<String>();
        let socket_recv = socket.clone();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            let size = socket_recv.recv(&mut buf).await.unwrap();
            packet_tx
                .send(from_utf8(&buf[..size]).unwrap().to_string())
                .unwrap();
        });
        OpenSocketRecvPacket { socket, packet_rx }
    }

    /// Opens a socket, listening for packets. Received packets are forwarded over the
    /// returned channel.
    pub async fn open_socket_and_recv_multiple_packets(
        &mut self,
    ) -> (mpsc::Receiver<String>, Arc<UdpSocket>) {
        let (packet_tx, packet_rx) = mpsc::channel::<String>(10);
        let socket = self.create_socket().await;
        let mut shutdown_rx = self.get_shutdown_subscriber().await;
        let socket_recv = socket.clone();
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];
            loop {
                tokio::select! {
                    received = socket_recv.recv_from(&mut buf) => {
                        let (size, _) = received.unwrap();
                        let str = from_utf8(&buf[..size]).unwrap().to_string();
                        match packet_tx.send(str).await {
                            Ok(_) => {}
                            Err(error) => {
                                tracing::warn!(target: "recv_multiple_packets", %error, "recv_chan dropped");
                                return;
                            }
                        };
                    },
                    _ = shutdown_rx.changed() => {
                        return;
                    }
                }
            }
        });
        (packet_rx, socket)
    }

    /// Runs a simple UDP server that echos back payloads.
    /// Returns the server's address.
    pub async fn run_echo_server(&mut self) -> EndpointAddress {
        self.run_echo_server_with_tap(|_, _, _| {}).await
    }

    /// Runs a simple UDP server that echos back payloads.
    /// The provided function is invoked for each received payload.
    /// Returns the server's address.
    pub async fn run_echo_server_with_tap<F>(&mut self, tap: F) -> EndpointAddress
    where
        F: Fn(SocketAddr, &[u8], SocketAddr) + Send + 'static,
    {
        let socket = self.create_socket().await;
        let addr = socket.local_addr().unwrap();
        let mut shutdown = self.get_shutdown_subscriber().await;
        let local_addr = addr;
        tokio::spawn(async move {
            loop {
                let mut buf = vec![0; 1024];
                tokio::select! {
                    recvd = socket.recv_from(&mut buf) => {
                        let (size, sender) = recvd.unwrap();
                        let packet = &buf[..size];
                        tap(sender, packet, local_addr);
                        socket.send_to(packet, sender).await.unwrap();
                    },
                    _ = shutdown.changed() => {
                        return;
                    }
                }
            }
        });
        addr.into()
    }

    /// Run a proxy server with a supplied config.
    /// Admin is disabled for this method, as the majority of tests will not need it, and it makes it
    /// easier to avoid issues with port collisions.
    pub fn run_server_with_config(&mut self, mut config: Config) {
        config.admin = None;

        self.run_server(<_>::try_from(config).unwrap());
    }

    pub fn run_server(&mut self, server: crate::Server) {
        let (shutdown_tx, shutdown_rx) = watch::channel::<()>(());
        self.server_shutdown_tx.push(Some(shutdown_tx));
        tokio::spawn(async move {
            server.run(shutdown_rx).await.unwrap();
        });
    }

    /// Returns a receiver subscribed to the helper's shutdown event.
    async fn get_shutdown_subscriber(&mut self) -> watch::Receiver<()> {
        // If this is the first call, then we set up the channel first.
        match self.shutdown_ch {
            Some((_, ref rx)) => rx.clone(),
            None => {
                let ch = watch::channel(());
                let recv = ch.1.clone();
                self.shutdown_ch = Some(ch);
                recv
            }
        }
    }
}

/// assert that read makes no changes
pub fn assert_filter_read_no_change<F>(filter: &F)
where
    F: Filter,
{
    let endpoints = vec!["127.0.0.1:80".parse::<Endpoint>().unwrap()];
    let source = "127.0.0.1:90".parse().unwrap();
    let contents = "hello".to_string().into_bytes();

    match filter.read(ReadContext::new(
        Endpoints::new(endpoints.clone()).into(),
        source,
        contents.clone(),
    )) {
        None => unreachable!("should return a result"),
        Some(response) => {
            assert_eq!(
                endpoints,
                response.endpoints.iter().cloned().collect::<Vec<_>>()
            );
            assert_eq!(contents, response.contents);
        }
    }
}

/// assert that write makes no changes
pub fn assert_write_no_change<F>(filter: &F)
where
    F: Filter,
{
    let endpoint = "127.0.0.1:90".parse::<Endpoint>().unwrap();
    let contents = "hello".to_string().into_bytes();

    match filter.write(WriteContext::new(
        &endpoint,
        endpoint.address.clone(),
        "127.0.0.1:70".parse().unwrap(),
        contents.clone(),
    )) {
        None => unreachable!("should return a result"),
        Some(response) => assert_eq!(contents, response.contents),
    }
}

pub fn config_with_dummy_endpoint() -> ConfigBuilder {
    ConfigBuilder::default().endpoints(vec![Endpoint {
        address: "127.0.0.1:8080".parse().unwrap(),
        ..<_>::default()
    }])
}
/// Creates a dummy endpoint with `id` as a suffix.
pub fn ep(id: u8) -> Endpoint {
    Endpoint {
        address: ([127, 0, 0, id], 8080).into(),
        ..<_>::default()
    }
}

pub fn new_test_chain() -> crate::filters::SharedFilterChain {
    <_>::try_from([crate::config::Filter {
        name: "TestFilter".into(),
        config: None,
    }])
    .unwrap()
}

pub fn load_test_filters() {
    FilterRegistry::register([TestFilter::factory()]);
}

#[cfg(test)]
mod tests {
    use crate::test_utils::TestHelper;

    #[tokio::test]
    async fn test_echo_server() {
        let mut t = TestHelper::default();
        let echo_addr = t.run_echo_server().await;
        let endpoint = t.open_socket_and_recv_single_packet().await;
        let msg = "hello";
        endpoint
            .socket
            .send_to(msg.as_bytes(), &echo_addr.to_socket_addr().unwrap())
            .await
            .unwrap();
        assert_eq!(msg, endpoint.packet_rx.await.unwrap());
    }
}
<|endoftext|>


<commit_before>#![crate_type = "lib"]
#![crate_name = "coinaddress"]

//! Functions for validating the base58 hash checksums, including specifically
//! the bitcoin and litecoin addresses.

extern crate num;
extern crate rustc;
extern crate test;

use num::bigint::BigUint;
use num::Zero;
use std::num::{from_uint, from_int, ToPrimitive};
use rustc::util::sha2::{Sha256, Digest};

#[derive(PartialEq, Debug, Copy)]
pub enum ValidationError {
    /// Given address is too short to be valid
    TooShort,
    /// Encoding is not a valid base58
    InvalidEncoding,
    /// Computed hash does not match the embedded one
    HashMismatch,

    // currency specific
    /// This address is not a bitcoin address (testnet or real).
    /// May happen when attempting to validate btc address
    NotBitcoin,
    /// This address is not a litecoin address.
    /// May happen when attempting to validate ltc address
    NotLitecoin,
}

static BASE58_CHARS: &'static str = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

fn decode_base58(bc: &str) -> Option<BigUint> {
    let mut res: BigUint = Zero::zero();
    let b58: BigUint = from_int(58).unwrap();

    for c in bc.chars() {
        match BASE58_CHARS.find(c) {
            None => {
                return None
            },
            Some(x) => {
                res = res * &b58 + from_uint::<BigUint>(x).unwrap();
            }
        }
    }
    Some(res)
}

fn bigint_to_bytes(n: BigUint) -> Vec<u8> {
    let mut res: Vec<u8> = Vec::new();
    let mut tmp = n;
    let b256: BigUint = from_int(256).unwrap();

    while !tmp.is_zero() {
        res.insert(0, (&tmp % &b256).to_u8().unwrap());
        tmp = tmp / &b256;
    }

    if res.len() == 0 {
        return vec!(0);
    }

    res
}

fn pad_to(v: Vec<u8>, len: usize) -> Vec<u8> {
    let mut tmp: Vec<u8> = v;
    while tmp.len() < len {
        tmp.insert(0, 0)
    }
    tmp
}

fn double_sha256(chunk: &[u8]) -> Vec<u8> {
    let mut hash = Sha256::new();
    hash.input(chunk);
    let mut hash2 = Sha256::new();
    hash2.input(hash.result_bytes().as_slice());
    hash2.result_bytes()
}

/// Validate provided generic base58 hash.
/// Returns the hash version/type if correct and an error otherwise.
pub fn validate_base58_hash(addr: &str) -> Result<usize, ValidationError> {
    if addr.len() == 0 {
        return Err(ValidationError::TooShort);
    }

    let big = match decode_base58(addr) {
        None => return Err(ValidationError::InvalidEncoding),
        Some(x) => x
    };
    let bytes = bigint_to_bytes(big);
    let padded = pad_to(bytes, 25);

    let hash = double_sha256(padded.slice(0, padded.len() - 4));
    let short_hash = hash.slice(0, 4);
    let known = padded.slice(padded.len()-4, padded.len());
    if short_hash.as_slice() == known {
        Ok(padded[0] as usize)
    } else {
        Err(ValidationError::HashMismatch)
    }
}

/// Validate bitcoin address checksum.
/// Returns the hash version/type if correct and an error otherwise.
pub fn validate_btc_address(addr: &str) -> Result<usize, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(0) => Ok(0),      // real address
        Ok(5) => Ok(5),      // script hash
        Ok(111) => Ok(111),  // testnet address
        Ok(_) => Err(ValidationError::NotBitcoin),
        Err(x) => Err(x)
    }
}

/// Validate litecoin address checksum.
/// Returns the hash version/type if correct and an error otherwise.
pub fn validate_ltc_address(addr: &str) -> Result<usize, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(48) => Ok(48),      // real address
        Ok(111) => Ok(111),    // testnet address
        Ok(_) => Err(ValidationError::NotLitecoin),
        Err(x) => Err(x)
    }
}

#[test]
fn test_decoding() {
    use std::num::from_str_radix;

    assert_eq!(decode_base58("1"), from_int(0));
    assert_eq!(decode_base58("1AaZz"), from_int(0x1C8485));
    assert_eq!(decode_base58("AO"), None);
    assert_eq!(decode_base58("1AGNa15ZQXAZUgFiqJ2i7Z2DPU2J6hW62i"), from_str_radix("65A16059864A2FDBC7C99A4723A8395BC6F188EBC046B2FF", 16).ok());
}

#[test]
fn test_to_bytes() {
    assert_eq!(bigint_to_bytes(Zero::zero()), vec!(0));
    assert_eq!(bigint_to_bytes(from_int(1).unwrap()), vec!(1));
    assert_eq!(bigint_to_bytes(from_int(256).unwrap()), vec!(1, 0));
}

#[test]
fn test_validate_hash() {
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_base58_hash("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYheX"), Err(ValidationError::HashMismatch));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYh  "), Err(ValidationError::InvalidEncoding));
    assert_eq!(validate_base58_hash("1"), Err(ValidationError::HashMismatch));
    assert_eq!(validate_base58_hash(""), Err(ValidationError::TooShort));
}

#[test]
fn test_validate_btc_address() {
    assert_eq!(validate_btc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_btc_address("3EktnHQD7RiAE6uzMj2ZifT9YgRrkSgzQX"), Ok(5));
    assert_eq!(validate_btc_address("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_btc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Err(ValidationError::NotBitcoin));
}

#[test]
fn test_validate_ltc_address() {
    assert_eq!(validate_ltc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Ok(48));
    assert_eq!(validate_ltc_address("muen9zszN6rVwXaFw48xh6YkdUSjJcfzek"), Ok(111));
    assert_eq!(validate_ltc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Err(ValidationError::NotLitecoin));
}

#[bench]
fn bench_validate_hash(b: &mut test::Bencher) {
    b.iter(|| {
        validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem")
    })
}
<commit_msg>Use new slice syntax<commit_after>#![crate_type = "lib"]
#![crate_name = "coinaddress"]

//! Functions for validating the base58 hash checksums, including specifically
//! the bitcoin and litecoin addresses.

extern crate num;
extern crate rustc;
extern crate test;

use num::bigint::BigUint;
use num::Zero;
use std::num::{from_uint, from_int, ToPrimitive};
use rustc::util::sha2::{Sha256, Digest};

#[derive(PartialEq, Debug, Copy)]
pub enum ValidationError {
    /// Given address is too short to be valid
    TooShort,
    /// Encoding is not a valid base58
    InvalidEncoding,
    /// Computed hash does not match the embedded one
    HashMismatch,

    // currency specific
    /// This address is not a bitcoin address (testnet or real).
    /// May happen when attempting to validate btc address
    NotBitcoin,
    /// This address is not a litecoin address.
    /// May happen when attempting to validate ltc address
    NotLitecoin,
}

static BASE58_CHARS: &'static str = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";

fn decode_base58(bc: &str) -> Option<BigUint> {
    let mut res: BigUint = Zero::zero();
    let b58: BigUint = from_int(58).unwrap();

    for c in bc.chars() {
        match BASE58_CHARS.find(c) {
            None => {
                return None
            },
            Some(x) => {
                res = res * &b58 + from_uint::<BigUint>(x).unwrap();
            }
        }
    }
    Some(res)
}

fn bigint_to_bytes(n: BigUint) -> Vec<u8> {
    let mut res: Vec<u8> = Vec::new();
    let mut tmp = n;
    let b256: BigUint = from_int(256).unwrap();

    while !tmp.is_zero() {
        res.insert(0, (&tmp % &b256).to_u8().unwrap());
        tmp = tmp / &b256;
    }

    if res.len() == 0 {
        return vec!(0);
    }

    res
}

fn pad_to(v: Vec<u8>, len: usize) -> Vec<u8> {
    let mut tmp: Vec<u8> = v;
    while tmp.len() < len {
        tmp.insert(0, 0)
    }
    tmp
}

fn double_sha256(chunk: &[u8]) -> Vec<u8> {
    let mut hash = Sha256::new();
    hash.input(chunk);
    let mut hash2 = Sha256::new();
    hash2.input(hash.result_bytes().as_slice());
    hash2.result_bytes()
}

/// Validate provided generic base58 hash.
/// Returns the hash version/type if correct and an error otherwise.
pub fn validate_base58_hash(addr: &str) -> Result<usize, ValidationError> {
    if addr.len() == 0 {
        return Err(ValidationError::TooShort);
    }

    let big = match decode_base58(addr) {
        None => return Err(ValidationError::InvalidEncoding),
        Some(x) => x
    };
    let bytes = bigint_to_bytes(big);
    let padded = pad_to(bytes, 25);

    let hash = double_sha256(&padded[0 .. padded.len() - 4]);
    let short_hash = &hash[0..4];
    let known = &padded[padded.len()-4 .. padded.len()];
    if short_hash.as_slice() == known {
        Ok(padded[0] as usize)
    } else {
        Err(ValidationError::HashMismatch)
    }
}

/// Validate bitcoin address checksum.
/// Returns the hash version/type if correct and an error otherwise.
pub fn validate_btc_address(addr: &str) -> Result<usize, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(0) => Ok(0),      // real address
        Ok(5) => Ok(5),      // script hash
        Ok(111) => Ok(111),  // testnet address
        Ok(_) => Err(ValidationError::NotBitcoin),
        Err(x) => Err(x)
    }
}

/// Validate litecoin address checksum.
/// Returns the hash version/type if correct and an error otherwise.
pub fn validate_ltc_address(addr: &str) -> Result<usize, ValidationError> {
    match validate_base58_hash(addr) {
        Ok(48) => Ok(48),      // real address
        Ok(111) => Ok(111),    // testnet address
        Ok(_) => Err(ValidationError::NotLitecoin),
        Err(x) => Err(x)
    }
}

#[test]
fn test_decoding() {
    use std::num::from_str_radix;

    assert_eq!(decode_base58("1"), from_int(0));
    assert_eq!(decode_base58("1AaZz"), from_int(0x1C8485));
    assert_eq!(decode_base58("AO"), None);
    assert_eq!(decode_base58("1AGNa15ZQXAZUgFiqJ2i7Z2DPU2J6hW62i"), from_str_radix("65A16059864A2FDBC7C99A4723A8395BC6F188EBC046B2FF", 16).ok());
}

#[test]
fn test_to_bytes() {
    assert_eq!(bigint_to_bytes(Zero::zero()), vec!(0));
    assert_eq!(bigint_to_bytes(from_int(1).unwrap()), vec!(1));
    assert_eq!(bigint_to_bytes(from_int(256).unwrap()), vec!(1, 0));
}

#[test]
fn test_validate_hash() {
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_base58_hash("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYheX"), Err(ValidationError::HashMismatch));
    assert_eq!(validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYh  "), Err(ValidationError::InvalidEncoding));
    assert_eq!(validate_base58_hash("1"), Err(ValidationError::HashMismatch));
    assert_eq!(validate_base58_hash(""), Err(ValidationError::TooShort));
}

#[test]
fn test_validate_btc_address() {
    assert_eq!(validate_btc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Ok(0));
    assert_eq!(validate_btc_address("3EktnHQD7RiAE6uzMj2ZifT9YgRrkSgzQX"), Ok(5));
    assert_eq!(validate_btc_address("mipcBbFg9gMiCh81Kj8tqqdgoZub1ZJRfn"), Ok(111));
    assert_eq!(validate_btc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Err(ValidationError::NotBitcoin));
}

#[test]
fn test_validate_ltc_address() {
    assert_eq!(validate_ltc_address("LRELGDJyeCPRDXz4Dh1kWorMN9hTBB7CEz"), Ok(48));
    assert_eq!(validate_ltc_address("muen9zszN6rVwXaFw48xh6YkdUSjJcfzek"), Ok(111));
    assert_eq!(validate_ltc_address("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem"), Err(ValidationError::NotLitecoin));
}

#[bench]
fn bench_validate_hash(b: &mut test::Bencher) {
    b.iter(|| {
        validate_base58_hash("17VZNX1SN5NtKa8UQFxwQbFeFc3iqRYhem")
    })
}
<|endoftext|>


<commit_before>//! On Linux we try to use the more accurate `utimensat` syscall but this isn't
//! always available so we also fall back to `utimes` if we couldn't find
//! `utimensat` at runtime.

use std::fs;
use std::io;
use std::mem;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT};

use FileTime;
use super::libc::{self, c_int, c_char, timespec};

pub fn set_file_times(p: &Path, atime: Option<FileTime>, mtime: Option<FileTime>) -> io::Result<()> {
    set_times(p, atime, mtime, false)
}

pub fn set_file_handle_times(
    f: &mut fs::File,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
) -> io::Result<()> {
    if cfg!(feature = "utimensat") {
         match futimens() {
            Some(func) => {
                super::futimens(f, atime, mtime, func)
            },
            None => {
                Err(io::Error::new(
                    io::ErrorKind::NotFound,
                    "Version of libc does not support the futimens syscall"
                ))
            }
        }
    } else {
        // Try to use the more-accurate `utimensat` when possible.
        static INVALID: AtomicBool = ATOMIC_BOOL_INIT;
        if !INVALID.load(Ordering::SeqCst) {
            if let Some(func) = futimens() {
                // Even when libc has `futimens`, the kernel may return `ENOSYS`,
                // and then we'll need to use the `utimes` fallback instead.
                match super::futimens(f, atime, mtime, func) {
                    Err(ref e) if e.raw_os_error() == Some(libc::ENOSYS) => {
                        INVALID.store(true, Ordering::SeqCst);
                    }
                    valid => return valid,
                }
            }
        }

        // Safe to unwrap atime and mtime -- without the utimensat feature enabled, the public API
        // only exposes methods that intialize atime and mtime as Some.
        super::futimes(f, atime.unwrap(), mtime.unwrap(), libc::futimes)
    }

}

pub fn set_symlink_file_times(p: &Path, atime: Option<FileTime>, mtime: Option<FileTime>) -> io::Result<()> {
    set_times(p, atime, mtime, true)
}

fn set_times(p: &Path, atime: Option<FileTime>, mtime: Option<FileTime>, symlink: bool) -> io::Result<()> {
    let flags = if symlink { libc::AT_SYMLINK_NOFOLLOW } else { 0 };

    if cfg!(feature = "utimensat") {
         match utimensat() {
            Some(f) => {
                super::utimensat(p, atime, mtime, f, flags)
            },
            None => {
                Err(io::Error::new(
                    io::ErrorKind::NotFound,
                    "Version of libc does not support the utimensat syscall"
                ))
            }
        }
    } else {
        // Try to use the more-accurate `utimensat` when possible.
        static INVALID: AtomicBool = ATOMIC_BOOL_INIT;
        if !INVALID.load(Ordering::SeqCst) {
            if let Some(f) = utimensat() {
                // Even when libc has `utimensat`, the kernel may return `ENOSYS`,
                // and then we'll need to use the `utimes` fallback instead.
                match super::utimensat(p, atime, mtime, f, flags) {
                    Err(ref e) if e.raw_os_error() == Some(libc::ENOSYS) => {
                        INVALID.store(true, Ordering::SeqCst);
                    }
                    valid => return valid,
                }
            }
        }

        let utimes = if symlink { libc::lutimes } else { libc::utimes };

        // Safe to unwrap atime and mtime -- without the utimensat feature enabled, the public API
        // only exposes methods that intialize atime and mtime as Some.
        super::utimes(p, atime.unwrap(), mtime.unwrap(), utimes)
    }
}

fn utimensat() -> Option<unsafe extern fn(c_int, *const c_char, *const timespec, c_int) -> c_int> {
    static ADDR: AtomicUsize = ATOMIC_USIZE_INIT;
    unsafe {
        match ADDR.load(Ordering::SeqCst) {
            0 => {}
            1 => return None,
            n => return Some(mem::transmute(n)),
        }
        let name = b"utimensat\0";
        let sym = libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr() as *const _);
        let (val, ret) = if sym.is_null() {
            (1, None)
        } else {
            (sym as usize, Some(mem::transmute(sym)))
        };
        ADDR.store(val, Ordering::SeqCst);
        return ret
    }
}

fn futimens() -> Option<unsafe extern fn(c_int, *const timespec) -> c_int> {
    static ADDR: AtomicUsize = ATOMIC_USIZE_INIT;
    unsafe {
        match ADDR.load(Ordering::SeqCst) {
            0 => {}
            1 => return None,
            n => return Some(mem::transmute(n)),
        }
        let name = b"futimens\0";
        let sym = libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr() as *const _);
        let (val, ret) = if sym.is_null() {
            (1, None)
        } else {
            (sym as usize, Some(mem::transmute(sym)))
        };
        ADDR.store(val, Ordering::SeqCst);
        return ret
    }
}
<commit_msg>Remove usage of deprecated constants<commit_after>//! On Linux we try to use the more accurate `utimensat` syscall but this isn't
//! always available so we also fall back to `utimes` if we couldn't find
//! `utimensat` at runtime.

use std::fs;
use std::io;
use std::mem;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst, AtomicBool};

use FileTime;
use super::libc::{self, c_int, c_char, timespec};

pub fn set_file_times(p: &Path, atime: Option<FileTime>, mtime: Option<FileTime>) -> io::Result<()> {
    set_times(p, atime, mtime, false)
}

pub fn set_file_handle_times(
    f: &mut fs::File,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
) -> io::Result<()> {
    if cfg!(feature = "utimensat") {
         match futimens() {
            Some(func) => {
                super::futimens(f, atime, mtime, func)
            },
            None => {
                Err(io::Error::new(
                    io::ErrorKind::NotFound,
                    "Version of libc does not support the futimens syscall"
                ))
            }
        }
    } else {
        // Try to use the more-accurate `utimensat` when possible.
        static INVALID: AtomicBool = AtomicBool::new(false);
        if !INVALID.load(SeqCst) {
            if let Some(func) = futimens() {
                // Even when libc has `futimens`, the kernel may return `ENOSYS`,
                // and then we'll need to use the `utimes` fallback instead.
                match super::futimens(f, atime, mtime, func) {
                    Err(ref e) if e.raw_os_error() == Some(libc::ENOSYS) => {
                        INVALID.store(true, SeqCst);
                    }
                    valid => return valid,
                }
            }
        }

        // Safe to unwrap atime and mtime -- without the utimensat feature enabled, the public API
        // only exposes methods that intialize atime and mtime as Some.
        super::futimes(f, atime.unwrap(), mtime.unwrap(), libc::futimes)
    }

}

pub fn set_symlink_file_times(p: &Path, atime: Option<FileTime>, mtime: Option<FileTime>) -> io::Result<()> {
    set_times(p, atime, mtime, true)
}

fn set_times(p: &Path, atime: Option<FileTime>, mtime: Option<FileTime>, symlink: bool) -> io::Result<()> {
    let flags = if symlink { libc::AT_SYMLINK_NOFOLLOW } else { 0 };

    if cfg!(feature = "utimensat") {
         match utimensat() {
            Some(f) => {
                super::utimensat(p, atime, mtime, f, flags)
            },
            None => {
                Err(io::Error::new(
                    io::ErrorKind::NotFound,
                    "Version of libc does not support the utimensat syscall"
                ))
            }
        }
    } else {
        // Try to use the more-accurate `utimensat` when possible.
        static INVALID: AtomicBool = AtomicBool::new(false);
        if !INVALID.load(SeqCst) {
            if let Some(f) = utimensat() {
                // Even when libc has `utimensat`, the kernel may return `ENOSYS`,
                // and then we'll need to use the `utimes` fallback instead.
                match super::utimensat(p, atime, mtime, f, flags) {
                    Err(ref e) if e.raw_os_error() == Some(libc::ENOSYS) => {
                        INVALID.store(true, SeqCst);
                    }
                    valid => return valid,
                }
            }
        }

        let utimes = if symlink { libc::lutimes } else { libc::utimes };

        // Safe to unwrap atime and mtime -- without the utimensat feature enabled, the public API
        // only exposes methods that intialize atime and mtime as Some.
        super::utimes(p, atime.unwrap(), mtime.unwrap(), utimes)
    }
}

fn utimensat() -> Option<unsafe extern fn(c_int, *const c_char, *const timespec, c_int) -> c_int> {
    static ADDR: AtomicUsize = AtomicUsize::new(0);
    unsafe {
        match ADDR.load(SeqCst) {
            0 => {}
            1 => return None,
            n => return Some(mem::transmute(n)),
        }
        let name = b"utimensat\0";
        let sym = libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr() as *const _);
        let (val, ret) = if sym.is_null() {
            (1, None)
        } else {
            (sym as usize, Some(mem::transmute(sym)))
        };
        ADDR.store(val, SeqCst);
        return ret
    }
}

fn futimens() -> Option<unsafe extern fn(c_int, *const timespec) -> c_int> {
    static ADDR: AtomicUsize = AtomicUsize::new(0);
    unsafe {
        match ADDR.load(SeqCst) {
            0 => {}
            1 => return None,
            n => return Some(mem::transmute(n)),
        }
        let name = b"futimens\0";
        let sym = libc::dlsym(libc::RTLD_DEFAULT, name.as_ptr() as *const _);
        let (val, ret) = if sym.is_null() {
            (1, None)
        } else {
            (sym as usize, Some(mem::transmute(sym)))
        };
        ADDR.store(val, SeqCst);
        return ret
    }
}
<|endoftext|>


<commit_before>//! The futures-rs `select! macro implementation.

#![recursion_limit="128"]

extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::Span;
use proc_macro_hack::proc_macro_hack;
use quote::quote;
use syn::{parenthesized, parse_quote, Expr, Ident, Pat, Token};
use syn::parse::{Parse, ParseStream};

mod kw {
    syn::custom_keyword!(complete);
    syn::custom_keyword!(futures_crate_path);
}

struct Select {
    futures_crate_path: Option<syn::Path>,
    // span of `complete`, then expression after `=> ...`
    complete: Option<Expr>,
    default: Option<Expr>,
    normal_fut_exprs: Vec<Expr>,
    normal_fut_handlers: Vec<(Pat, Expr)>,
}

#[allow(clippy::large_enum_variant)]
enum CaseKind {
    Complete,
    Default,
    Normal(Pat, Expr),
}

impl Parse for Select {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut select = Select {
            futures_crate_path: None,
            complete: None,
            default: None,
            normal_fut_exprs: vec![],
            normal_fut_handlers: vec![],
        };

        // When `futures_crate_path(::path::to::futures::lib)` is provided,
        // it sets the path through which futures library functions will be
        // accessed.
        if input.peek(kw::futures_crate_path) {
            input.parse::<kw::futures_crate_path>()?;
            let content;
            parenthesized!(content in input);
            select.futures_crate_path = Some(content.parse()?);
        }

        while !input.is_empty() {
            let case_kind = if input.peek(kw::complete) {
                // `complete`
                if select.complete.is_some() {
                    return Err(input.error("multiple `complete` cases found, only one allowed"));
                }
                input.parse::<kw::complete>()?;
                CaseKind::Complete
            } else if input.peek(Token![default]) {
                // `default`
                if select.default.is_some() {
                    return Err(input.error("multiple `default` cases found, only one allowed"));
                }
                input.parse::<Ident>()?;
                CaseKind::Default
            } else {
                // `<pat> = <expr>`
                let pat = input.parse()?;
                input.parse::<Token![=]>()?;
                let expr = input.parse()?;
                CaseKind::Normal(pat, expr)
            };

            // `=> <expr>`
            input.parse::<Token![=>]>()?;
            let expr = input.parse::<Expr>()?;

            // Commas after the expression are only optional if it's a `Block`,
            // otherwise they are required.
            match &expr {
                Expr::Block(_) => { input.parse::<Option<Token![,]>>()?; }
                _ => { input.parse::<Token![,]>()?; }
            }

            match case_kind {
                CaseKind::Complete => select.complete = Some(expr),
                CaseKind::Default => select.default = Some(expr),
                CaseKind::Normal(pat, fut_expr) => {
                    select.normal_fut_exprs.push(fut_expr);
                    select.normal_fut_handlers.push((pat, expr));
                },
            }
        }

        Ok(select)
    }
}

// Enum over all the cases in which the `select!` waiting has completed and the result
// can be processed.
//
// `enum __PrivResult<_1, _2, ...> { _1(_1), _2(_2), ..., Complete }`
fn declare_result_enum(
    result_ident: Ident,
    variants: usize,
    complete: bool,
    span: Span
) -> (Vec<Ident>, syn::ItemEnum) {
    // "_0", "_1", "_2"
    let variant_names: Vec<Ident> =
        (0..variants)
            .map(|num| Ident::new(&format!("_{}", num), span))
            .collect();

    // `<_1, _2, ...>`
    let enum_generics = variant_names.iter()
        .map(|name| {
            syn::GenericParam::Type(parse_quote!( #name ))
        })
        .collect();

    let unit_variant = |name| syn::Variant {
        attrs: vec![],
        ident: Ident::new(name, span),
        fields: syn::Fields::Unit,
        discriminant: None,
    };

    // `_1(_1), _2(_2), ..., Complete`
    let enum_variants = variant_names.iter()
        .map(|name| parse_quote!( #name ( #name ) ))
        .chain(if complete { Some(unit_variant("Complete")) } else { None })
        .collect();

    (
        variant_names,
        syn::ItemEnum {
            attrs: vec![],
            vis: syn::Visibility::Inherited,
            enum_token: syn::token::Enum::default(),
            ident: result_ident,
            generics: syn::Generics {
                lt_token: None,
                params: enum_generics,
                gt_token: None,
                where_clause: None,
            },
            brace_token: syn::token::Brace { span },
            variants: enum_variants,
        },
    )
}

/// The `select!` macro.
#[proc_macro_hack]
pub fn select(input: TokenStream) -> TokenStream {
    let parsed = syn::parse_macro_input!(input as Select);

    let futures_crate: syn::Path = parsed.futures_crate_path.unwrap_or_else(|| parse_quote!(::futures_util));
    let rand_crate: syn::Path = parse_quote!(#futures_crate::rand_reexport);

    // should be def_site, but that's unstable
    let span = Span::call_site();

    let enum_ident = Ident::new("__PrivResult", span);

    let (variant_names, enum_item) = declare_result_enum(
        enum_ident.clone(),
        parsed.normal_fut_exprs.len(),
        parsed.complete.is_some(),
        span,
    );

    // bind non-`Ident` future exprs w/ `let`
    let mut future_let_bindings = Vec::with_capacity(parsed.normal_fut_exprs.len());
    let bound_future_names: Vec<_> = parsed.normal_fut_exprs.into_iter()
        .zip(variant_names.iter())
        .map(|(expr, variant_name)| {
            match expr {
                // Don't bind futures that are already a path.
                // This prevents creating redundant stack space
                // for them.
                syn::Expr::Path(path) => path,
                x => {
                    future_let_bindings.push(quote! {
                        let mut #variant_name = #x;
                    });
                    parse_quote! { #variant_name }
                }
            }
        })
        .collect();

    // For each future, make an `&mut dyn FnMut(&LocalWaker) -> Option<Poll<__PrivResult<...>>`
    // to use for polling that individual future. These will then be put in an array.
    let poll_functions = bound_future_names.iter().zip(variant_names.iter())
        .map(|(bound_future_name, variant_name)| {
            quote! {
                let mut #variant_name = |__lw: &_| {
                    if #futures_crate::future::FusedFuture::is_terminated(&#bound_future_name) {
                        None
                    } else {
                        Some(#futures_crate::future::FutureExt::poll_unpin(
                            &mut #bound_future_name,
                            __lw,
                        ).map(#enum_ident::#variant_name))
                    }
                };
                let #variant_name: &mut dyn FnMut(
                    &#futures_crate::task::LocalWaker
                ) -> Option<#futures_crate::task::Poll<_>> = &mut #variant_name;
            }
        });

    let none_polled = if parsed.complete.is_some() {
        quote! {
            #futures_crate::task::Poll::Ready(#enum_ident::Complete)
        }
    } else {
        quote! {
            panic!("all futures in select! were completed,\
                    but no `complete =>` handler was provided")
        }
    };

    let branches = parsed.normal_fut_handlers.into_iter()
        .zip(variant_names.iter())
        .map(|((pat, expr), variant_name)| {
            quote! {
                #enum_ident::#variant_name(#pat) => { #expr },
            }
        });
    let branches = quote! { #( #branches )* };

    let complete_branch = parsed.complete.map(|complete_expr| {
        quote! {
            #enum_ident::Complete => { #complete_expr },
        }
    });

    let branches = quote! {
        #branches
        #complete_branch
    };

    let await_and_select = if let Some(default_expr) = parsed.default {
        quote! {
            if let #futures_crate::task::Poll::Ready(x) =
                __poll_fn(#futures_crate::task::noop_local_waker_ref())
            {
                match x { #branches }
            } else {
                #default_expr
            };
        }
    } else {
        quote! {
            match r#await!(#futures_crate::future::poll_fn(__poll_fn)) {
                #branches
            }
        }
    };

    TokenStream::from(quote! { {
        #enum_item
        #( #future_let_bindings )*

        let mut __poll_fn = |__lw: &#futures_crate::task::LocalWaker| {
            let mut __any_polled = false;

            #( #poll_functions )*

            let mut __select_arr = [#( #variant_names ),*];
            #rand_crate::Rng::shuffle(
                &mut #rand_crate::thread_rng(),
                &mut __select_arr,
            );
            for poller in &mut __select_arr {
                let poller: &mut &mut dyn FnMut(
                    &#futures_crate::task::LocalWaker
                ) -> Option<#futures_crate::task::Poll<_>> = poller;
                match poller(__lw) {
                    Some(x @ #futures_crate::task::Poll::Ready(_)) =>
                        return x,
                    Some(#futures_crate::task::Poll::Pending) => {
                        __any_polled = true;
                    }
                    None => {}
                }
            }

            if !__any_polled {
                #none_polled
            } else {
                #futures_crate::task::Poll::Pending
            }
        };

        #await_and_select
    } })
}
<commit_msg>Address misc. review suggestions<commit_after>//! The futures-rs `select! macro implementation.

#![recursion_limit="128"]

extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::Span;
use proc_macro_hack::proc_macro_hack;
use quote::quote;
use syn::{parenthesized, parse_quote, Expr, Ident, Pat, Token};
use syn::parse::{Parse, ParseStream};

mod kw {
    syn::custom_keyword!(complete);
    syn::custom_keyword!(futures_crate_path);
}

struct Select {
    futures_crate_path: Option<syn::Path>,
    // span of `complete`, then expression after `=> ...`
    complete: Option<Expr>,
    default: Option<Expr>,
    normal_fut_exprs: Vec<Expr>,
    normal_fut_handlers: Vec<(Pat, Expr)>,
}

#[allow(clippy::large_enum_variant)]
enum CaseKind {
    Complete,
    Default,
    Normal(Pat, Expr),
}

impl Parse for Select {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut select = Select {
            futures_crate_path: None,
            complete: None,
            default: None,
            normal_fut_exprs: vec![],
            normal_fut_handlers: vec![],
        };

        // When `futures_crate_path(::path::to::futures::lib)` is provided,
        // it sets the path through which futures library functions will be
        // accessed.
        if input.peek(kw::futures_crate_path) {
            input.parse::<kw::futures_crate_path>()?;
            let content;
            parenthesized!(content in input);
            select.futures_crate_path = Some(content.parse()?);
        }

        while !input.is_empty() {
            let case_kind = if input.peek(kw::complete) {
                // `complete`
                if select.complete.is_some() {
                    return Err(input.error("multiple `complete` cases found, only one allowed"));
                }
                input.parse::<kw::complete>()?;
                CaseKind::Complete
            } else if input.peek(Token![default]) {
                // `default`
                if select.default.is_some() {
                    return Err(input.error("multiple `default` cases found, only one allowed"));
                }
                input.parse::<Ident>()?;
                CaseKind::Default
            } else {
                // `<pat> = <expr>`
                let pat = input.parse()?;
                input.parse::<Token![=]>()?;
                let expr = input.parse()?;
                CaseKind::Normal(pat, expr)
            };

            // `=> <expr>`
            input.parse::<Token![=>]>()?;
            let expr = input.parse::<Expr>()?;

            // Commas after the expression are only optional if it's a `Block`
            // or it is the last branch in the `match`.
            let is_block = match expr { Expr::Block(_) => true, _ => false };
            if is_block || input.is_empty() {
                input.parse::<Option<Token![,]>>()?;
            } else {
                input.parse::<Token![,]>()?;
            }

            match case_kind {
                CaseKind::Complete => select.complete = Some(expr),
                CaseKind::Default => select.default = Some(expr),
                CaseKind::Normal(pat, fut_expr) => {
                    select.normal_fut_exprs.push(fut_expr);
                    select.normal_fut_handlers.push((pat, expr));
                },
            }
        }

        Ok(select)
    }
}

// Enum over all the cases in which the `select!` waiting has completed and the result
// can be processed.
//
// `enum __PrivResult<_1, _2, ...> { _1(_1), _2(_2), ..., Complete }`
fn declare_result_enum(
    result_ident: Ident,
    variants: usize,
    complete: bool,
    span: Span
) -> (Vec<Ident>, syn::ItemEnum) {
    // "_0", "_1", "_2"
    let variant_names: Vec<Ident> =
        (0..variants)
            .map(|num| Ident::new(&format!("_{}", num), span))
            .collect();

    let type_parameters = &variant_names;
    let variants = &variant_names;

    let complete_variant = if complete {
        Some(quote!(Complete))
    } else {
        None
    };

    let enum_item = parse_quote! {
        enum #result_ident<#(#type_parameters,)*> {
            #(
                #variants(#type_parameters),
            )*
            #complete_variant
        }
    };

    (variant_names, enum_item)
}

/// The `select!` macro.
#[proc_macro_hack]
pub fn select(input: TokenStream) -> TokenStream {
    let parsed = syn::parse_macro_input!(input as Select);

    let futures_crate: syn::Path = parsed.futures_crate_path.unwrap_or_else(|| parse_quote!(::futures_util));
    let rand_crate: syn::Path = parse_quote!(#futures_crate::rand_reexport);

    // should be def_site, but that's unstable
    let span = Span::call_site();

    let enum_ident = Ident::new("__PrivResult", span);

    let (variant_names, enum_item) = declare_result_enum(
        enum_ident.clone(),
        parsed.normal_fut_exprs.len(),
        parsed.complete.is_some(),
        span,
    );

    // bind non-`Ident` future exprs w/ `let`
    let mut future_let_bindings = Vec::with_capacity(parsed.normal_fut_exprs.len());
    let bound_future_names: Vec<_> = parsed.normal_fut_exprs.into_iter()
        .zip(variant_names.iter())
        .map(|(expr, variant_name)| {
            match expr {
                // Don't bind futures that are already a path.
                // This prevents creating redundant stack space
                // for them.
                syn::Expr::Path(path) => path,
                _ => {
                    future_let_bindings.push(quote! {
                        let mut #variant_name = #expr;
                    });
                    parse_quote! { #variant_name }
                }
            }
        })
        .collect();

    // For each future, make an `&mut dyn FnMut(&LocalWaker) -> Option<Poll<__PrivResult<...>>`
    // to use for polling that individual future. These will then be put in an array.
    let poll_functions = bound_future_names.iter().zip(variant_names.iter())
        .map(|(bound_future_name, variant_name)| {
            quote! {
                let mut #variant_name = |__lw: &_| {
                    if #futures_crate::future::FusedFuture::is_terminated(&#bound_future_name) {
                        None
                    } else {
                        Some(#futures_crate::future::FutureExt::poll_unpin(
                            &mut #bound_future_name,
                            __lw,
                        ).map(#enum_ident::#variant_name))
                    }
                };
                let #variant_name: &mut dyn FnMut(
                    &#futures_crate::task::LocalWaker
                ) -> Option<#futures_crate::task::Poll<_>> = &mut #variant_name;
            }
        });

    let none_polled = if parsed.complete.is_some() {
        quote! {
            #futures_crate::task::Poll::Ready(#enum_ident::Complete)
        }
    } else {
        quote! {
            panic!("all futures in select! were completed,\
                    but no `complete =>` handler was provided")
        }
    };

    let branches = parsed.normal_fut_handlers.into_iter()
        .zip(variant_names.iter())
        .map(|((pat, expr), variant_name)| {
            quote! {
                #enum_ident::#variant_name(#pat) => { #expr },
            }
        });
    let branches = quote! { #( #branches )* };

    let complete_branch = parsed.complete.map(|complete_expr| {
        quote! {
            #enum_ident::Complete => { #complete_expr },
        }
    });

    let branches = quote! {
        #branches
        #complete_branch
    };

    let await_and_select = if let Some(default_expr) = parsed.default {
        quote! {
            if let #futures_crate::task::Poll::Ready(x) =
                __poll_fn(#futures_crate::task::noop_local_waker_ref())
            {
                match x { #branches }
            } else {
                #default_expr
            };
        }
    } else {
        quote! {
            match r#await!(#futures_crate::future::poll_fn(__poll_fn)) {
                #branches
            }
        }
    };

    TokenStream::from(quote! { {
        #enum_item
        #( #future_let_bindings )*

        let mut __poll_fn = |__lw: &#futures_crate::task::LocalWaker| {
            let mut __any_polled = false;

            #( #poll_functions )*

            let mut __select_arr = [#( #variant_names ),*];
            #rand_crate::Rng::shuffle(
                &mut #rand_crate::thread_rng(),
                &mut __select_arr,
            );
            for poller in &mut __select_arr {
                let poller: &mut &mut dyn FnMut(
                    &#futures_crate::task::LocalWaker
                ) -> Option<#futures_crate::task::Poll<_>> = poller;
                match poller(__lw) {
                    Some(x @ #futures_crate::task::Poll::Ready(_)) =>
                        return x,
                    Some(#futures_crate::task::Poll::Pending) => {
                        __any_polled = true;
                    }
                    None => {}
                }
            }

            if !__any_polled {
                #none_polled
            } else {
                #futures_crate::task::Poll::Pending
            }
        };

        #await_and_select
    } })
}
<|endoftext|>