<commit_before>use clap::Parser;
use dicom_core::{dicom_value, smallvec};
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_dump::DumpOptions;
use dicom_encoding::transfer_syntax;
use dicom_object::{mem::InMemDicomObject, open_file, StandardDataDictionary};
use dicom_transfer_syntax_registry::{entries, TransferSyntaxRegistry};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use query::parse_queries;
use smallvec::smallvec;
use snafu::{prelude::*, ErrorCompat};
use std::io::{stderr, Read};
use std::path::PathBuf;
use tracing::{debug, error, info, warn, Level};
use transfer_syntax::TransferSyntaxIndex;

mod query;

/// DICOM C-FIND SCU
#[derive(Debug, Parser)]
struct App {
    /// socket address to FIND SCP (example: "127.0.0.1:1045")
    addr: String,
    /// a DICOM file representing the query object
    file: Option<PathBuf>,
    /// sequence of queries
    #[clap(short('q'))]
    query: Vec<String>,

    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling AE title
    #[clap(long = "calling-ae-title", default_value = "FIND-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[clap(long = "called-ae-title")]
    called_ae_title: Option<String>,
    /// the maximum PDU length
    #[clap(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// use patient root information model
    #[clap(long, conflicts_with = "study")]
    patient: bool,
    /// use study root information model (default)
    #[clap(long, conflicts_with = "patient")]
    study: bool,
}

fn report<E: 'static>(err: &E)
where
    E: std::error::Error,
{
    error!("{}", err);
    if let Some(source) = err.source() {
        eprintln!();
        eprintln!("Caused by:");
        for (i, e) in std::iter::successors(Some(source), |e| e.source()).enumerate() {
            eprintln!("   {}: {}", i, e);
        }
    }
}

fn report_backtrace<E: 'static>(err: &E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    let env_backtrace = std::env::var("RUST_BACKTRACE").unwrap_or_default();
    let env_lib_backtrace = std::env::var("RUST_LIB_BACKTRACE").unwrap_or_default();
    if env_lib_backtrace == "1" || (env_backtrace == "1" && env_lib_backtrace != "0") {
        if let Some(backtrace) = ErrorCompat::backtrace(&err) {
            eprintln!();
            eprintln!("Backtrace:");
            eprintln!("{}", backtrace);
        }
    }
}

fn report_with_backtrace<E: 'static>(err: E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    report(&err);
    report_backtrace(&err);
}

fn main() {
    run().unwrap_or_else(|err| {
        report_with_backtrace(err);
        std::process::exit(-2);
    });
}

#[derive(Debug, Snafu)]
enum Error {
    /// Could not initialize SCU
    InitScu {
        source: dicom_ul::association::client::Error,
    },

    /// Could not construct DICOM command
    CreateCommand { source: dicom_object::Error },

    /// Could not read DICOM command
    ReadCommand { source: dicom_object::Error },

    /// Could not dump DICOM output
    DumpOutput { source: std::io::Error },

    #[snafu(whatever, display("{}", message))]
    Other {
        message: String,
        #[snafu(source(from(Box<dyn std::error::Error + 'static>, Some)))]
        source: Option<Box<dyn std::error::Error + 'static>>,
    },
}

fn build_query(
    file: Option<PathBuf>,
    q: Vec<String>,
    patient: bool,
    study: bool,
    verbose: bool,
) -> Result<InMemDicomObject, Error> {
    match (file, q) {
        (Some(file), q) => {
            if !q.is_empty() {
                whatever!("Conflicted file with query terms");
            }

            if verbose {
                info!("Opening file '{}'...", file.display());
            }

            open_file(file)
                .context(CreateCommandSnafu)
                .map(|file| file.into_inner())
        }
        (None, q) => {
            let mut obj =
                parse_queries(&q).whatever_context("Could not build query object from terms")?;

            // (0008,0052) CS QueryRetrieveLevel
            let level = match (patient, study) {
                (true, false) => "PATIENT",
                (false, true) |
                (false, false) => "STUDY",
                _ => unreachable!(),
            };
            obj.put(DataElement::new(
                tags::QUERY_RETRIEVE_LEVEL,
                VR::CS,
                PrimitiveValue::from(level),
            ));

            Ok(obj)
        }
    }
}

fn run() -> Result<(), Error> {
    let App {
        addr,
        file,
        verbose,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        patient,
        study,
        query,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        report(&e);
    });

    let dcm_query = build_query(file, query, patient, study, verbose)?;

    let abstract_syntax = match (patient, study) {
        // Patient Root Query/Retrieve Information Model - FIND
        (true, false) => "1.2.840.10008.5.1.4.1.2.1.1",
        // Study Root Query/Retrieve Information Model – FIND (default)
        (false, false) | (false, true) => "1.2.840.10008.5.1.4.1.2.2.1",
        // Series
        _ => unreachable!("Unexpected flag combination"),
    };

    if verbose {
        info!("Establishing association with '{}'...", &addr);
    }

    let mut scu_opt = ClientAssociationOptions::new()
        .with_abstract_syntax(abstract_syntax)
        .calling_ae_title(calling_ae_title)
        .max_pdu_length(max_pdu_length);

    if let Some(called_ae_title) = called_ae_title {
        scu_opt = scu_opt.called_ae_title(called_ae_title);
    }

    let mut scu = scu_opt.establish_with(&addr).context(InitScuSnafu)?;

    if verbose {
        info!("Association established");
    }

    let pc_selected = if let Some(pc_selected) = scu.presentation_contexts().first() {
        pc_selected
    } else {
        error!("Could not choose a presentation context");
        let _ = scu.abort();
        std::process::exit(-2);
    };
    let pc_selected_id = pc_selected.id;

    let ts = if let Some(ts) = TransferSyntaxRegistry.get(&pc_selected.transfer_syntax) {
        ts
    } else {
        error!("Poorly negotiated transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    if verbose {
        debug!("Transfer Syntax: {}", ts.name());
    }

    let cmd = find_req_command(abstract_syntax, 1);

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(&mut cmd_data, &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased())
        .whatever_context("Failed to write command")?;

    let implicit_vr_le = entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let mut iod_data = Vec::with_capacity(128);
    dcm_query
        .write_dataset_with_ts(&mut iod_data, &implicit_vr_le)
        .whatever_context("failed to write identifier dataset")?;

    let nbytes = cmd_data.len() + iod_data.len();

    if verbose {
        debug!("Sending query ({} B)...", nbytes);
    }

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Command,
            is_last: true,
            data: cmd_data,
        }],
    };
    scu.send(&pdu).whatever_context("Could not send command")?;

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Data,
            is_last: true,
            data: iod_data,
        }],
    };
    scu.send(&pdu)
        .whatever_context("Could not send C-Find request")?;

    if verbose {
        debug!("Awaiting response...");
    }

    let mut i = 0;
    loop {
        let rsp_pdu = scu
            .receive()
            .whatever_context("Failed to receive response from remote node")?;

        match rsp_pdu {
            Pdu::PData { data } => {
                let data_value = &data[0];

                let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                    &data_value.data[..],
                    &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                )
                .context(ReadCommandSnafu)?;
                if verbose {
                    eprintln!("Match #{} Response command:", i);
                    DumpOptions::new()
                        .dump_object_to(stderr(), &cmd_obj)
                        .context(DumpOutputSnafu)?;
                }
                let status = cmd_obj
                    .element(tags::STATUS)
                    .whatever_context("status code from response is missing")?
                    .to_int::<u16>()
                    .whatever_context("failed to read status code")?;
                if status == 0 {
                    if verbose {
                        debug!("Matching is complete");
                    }
                    if i == 0 {
                        info!("No results matching query");
                    }
                    break;
                } else if status == 0xFF00 || status == 0xFF01 {
                    if verbose {
                        debug!("Operation pending: {:x}", status);
                    }

                    // fetch DICOM data

                    let dcm = {
                        let mut rsp = scu.receive_pdata();
                        let mut response_data = Vec::new();
                        rsp.read_to_end(&mut response_data)
                            .whatever_context("Failed to read response data")?;

                        InMemDicomObject::read_dataset_with_ts(&response_data[..], &implicit_vr_le)
                            .whatever_context("Could not read response data set")?
                    };

                    println!("------------ Match #{} ------------", i);
                    DumpOptions::new()
                        .dump_object(&dcm)
                        .context(DumpOutputSnafu)?;
                    i += 1;
                } else {
                    warn!("Operation failed (status code {})", status);
                    break;
                }
            }

            pdu @ Pdu::Unknown { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::ReleaseRQ
            | pdu @ Pdu::ReleaseRP
            | pdu @ Pdu::AbortRQ { .. } => {
                error!("Unexpected SCP response: {:?}", pdu);
                let _ = scu.abort();
                std::process::exit(-2);
            }
        }
    }
    let _ = scu.release();

    Ok(())
}

fn find_req_command(
    sop_class_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + even_len(sop_class_uid.len())   // SOP Class UID
            + 8 + 2 // command field
            + 8 + 2 // message ID
            + 8 + 2 // priority
            + 8 + 2, // data set type
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        PrimitiveValue::from(sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        // 0020H: C-FIND-RQ message
        dicom_value!(U16, [0x0020]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        // medium
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<commit_msg>[findscu] stop findscu on query not specified<commit_after>use clap::Parser;
use dicom_core::{dicom_value, smallvec};
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_dump::DumpOptions;
use dicom_encoding::transfer_syntax;
use dicom_object::{mem::InMemDicomObject, open_file, StandardDataDictionary};
use dicom_transfer_syntax_registry::{entries, TransferSyntaxRegistry};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use query::parse_queries;
use smallvec::smallvec;
use snafu::{prelude::*, ErrorCompat};
use std::io::{stderr, Read};
use std::path::PathBuf;
use tracing::{debug, error, info, warn, Level};
use transfer_syntax::TransferSyntaxIndex;

mod query;

/// DICOM C-FIND SCU
#[derive(Debug, Parser)]
struct App {
    /// socket address to FIND SCP (example: "127.0.0.1:1045")
    addr: String,
    /// a DICOM file representing the query object
    file: Option<PathBuf>,
    /// sequence of queries
    #[clap(short('q'))]
    query: Vec<String>,

    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling AE title
    #[clap(long = "calling-ae-title", default_value = "FIND-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[clap(long = "called-ae-title")]
    called_ae_title: Option<String>,
    /// the maximum PDU length
    #[clap(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// use patient root information model
    #[clap(long, conflicts_with = "study")]
    patient: bool,
    /// use study root information model (default)
    #[clap(long, conflicts_with = "patient")]
    study: bool,
}

fn report<E: 'static>(err: &E)
where
    E: std::error::Error,
{
    error!("{}", err);
    if let Some(source) = err.source() {
        eprintln!();
        eprintln!("Caused by:");
        for (i, e) in std::iter::successors(Some(source), |e| e.source()).enumerate() {
            eprintln!("   {}: {}", i, e);
        }
    }
}

fn report_backtrace<E: 'static>(err: &E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    let env_backtrace = std::env::var("RUST_BACKTRACE").unwrap_or_default();
    let env_lib_backtrace = std::env::var("RUST_LIB_BACKTRACE").unwrap_or_default();
    if env_lib_backtrace == "1" || (env_backtrace == "1" && env_lib_backtrace != "0") {
        if let Some(backtrace) = ErrorCompat::backtrace(&err) {
            eprintln!();
            eprintln!("Backtrace:");
            eprintln!("{}", backtrace);
        }
    }
}

fn report_with_backtrace<E: 'static>(err: E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    report(&err);
    report_backtrace(&err);
}

fn main() {
    run().unwrap_or_else(|err| {
        report_with_backtrace(err);
        std::process::exit(-2);
    });
}

#[derive(Debug, Snafu)]
enum Error {
    /// Could not initialize SCU
    InitScu {
        source: dicom_ul::association::client::Error,
    },

    /// Could not construct DICOM command
    CreateCommand { source: dicom_object::Error },

    /// Could not read DICOM command
    ReadCommand { source: dicom_object::Error },

    /// Could not dump DICOM output
    DumpOutput { source: std::io::Error },

    #[snafu(whatever, display("{}", message))]
    Other {
        message: String,
        #[snafu(source(from(Box<dyn std::error::Error + 'static>, Some)))]
        source: Option<Box<dyn std::error::Error + 'static>>,
    },
}

fn build_query(
    file: Option<PathBuf>,
    q: Vec<String>,
    patient: bool,
    study: bool,
    verbose: bool,
) -> Result<InMemDicomObject, Error> {
    match (file, q) {
        (Some(file), q) => {
            if !q.is_empty() {
                whatever!("Conflicted file with query terms");
            }

            if verbose {
                info!("Opening file '{}'...", file.display());
            }

            open_file(file)
                .context(CreateCommandSnafu)
                .map(|file| file.into_inner())
        }
        (None, q) => {
            if q.is_empty() {
                whatever!("Query not specified");
            }

            let mut obj =
                parse_queries(&q).whatever_context("Could not build query object from terms")?;

            // (0008,0052) CS QueryRetrieveLevel
            let level = match (patient, study) {
                (true, false) => "PATIENT",
                (false, true) |
                (false, false) => "STUDY",
                _ => unreachable!(),
            };
            obj.put(DataElement::new(
                tags::QUERY_RETRIEVE_LEVEL,
                VR::CS,
                PrimitiveValue::from(level),
            ));

            Ok(obj)
        }
    }
}

fn run() -> Result<(), Error> {
    let App {
        addr,
        file,
        verbose,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        patient,
        study,
        query,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        report(&e);
    });

    let dcm_query = build_query(file, query, patient, study, verbose)?;

    let abstract_syntax = match (patient, study) {
        // Patient Root Query/Retrieve Information Model - FIND
        (true, false) => "1.2.840.10008.5.1.4.1.2.1.1",
        // Study Root Query/Retrieve Information Model – FIND (default)
        (false, false) | (false, true) => "1.2.840.10008.5.1.4.1.2.2.1",
        // Series
        _ => unreachable!("Unexpected flag combination"),
    };

    if verbose {
        info!("Establishing association with '{}'...", &addr);
    }

    let mut scu_opt = ClientAssociationOptions::new()
        .with_abstract_syntax(abstract_syntax)
        .calling_ae_title(calling_ae_title)
        .max_pdu_length(max_pdu_length);

    if let Some(called_ae_title) = called_ae_title {
        scu_opt = scu_opt.called_ae_title(called_ae_title);
    }

    let mut scu = scu_opt.establish_with(&addr).context(InitScuSnafu)?;

    if verbose {
        info!("Association established");
    }

    let pc_selected = if let Some(pc_selected) = scu.presentation_contexts().first() {
        pc_selected
    } else {
        error!("Could not choose a presentation context");
        let _ = scu.abort();
        std::process::exit(-2);
    };
    let pc_selected_id = pc_selected.id;

    let ts = if let Some(ts) = TransferSyntaxRegistry.get(&pc_selected.transfer_syntax) {
        ts
    } else {
        error!("Poorly negotiated transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    if verbose {
        debug!("Transfer Syntax: {}", ts.name());
    }

    let cmd = find_req_command(abstract_syntax, 1);

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(&mut cmd_data, &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased())
        .whatever_context("Failed to write command")?;

    let implicit_vr_le = entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let mut iod_data = Vec::with_capacity(128);
    dcm_query
        .write_dataset_with_ts(&mut iod_data, &implicit_vr_le)
        .whatever_context("failed to write identifier dataset")?;

    let nbytes = cmd_data.len() + iod_data.len();

    if verbose {
        debug!("Sending query ({} B)...", nbytes);
    }

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Command,
            is_last: true,
            data: cmd_data,
        }],
    };
    scu.send(&pdu).whatever_context("Could not send command")?;

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Data,
            is_last: true,
            data: iod_data,
        }],
    };
    scu.send(&pdu)
        .whatever_context("Could not send C-Find request")?;

    if verbose {
        debug!("Awaiting response...");
    }

    let mut i = 0;
    loop {
        let rsp_pdu = scu
            .receive()
            .whatever_context("Failed to receive response from remote node")?;

        match rsp_pdu {
            Pdu::PData { data } => {
                let data_value = &data[0];

                let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                    &data_value.data[..],
                    &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                )
                .context(ReadCommandSnafu)?;
                if verbose {
                    eprintln!("Match #{} Response command:", i);
                    DumpOptions::new()
                        .dump_object_to(stderr(), &cmd_obj)
                        .context(DumpOutputSnafu)?;
                }
                let status = cmd_obj
                    .element(tags::STATUS)
                    .whatever_context("status code from response is missing")?
                    .to_int::<u16>()
                    .whatever_context("failed to read status code")?;
                if status == 0 {
                    if verbose {
                        debug!("Matching is complete");
                    }
                    if i == 0 {
                        info!("No results matching query");
                    }
                    break;
                } else if status == 0xFF00 || status == 0xFF01 {
                    if verbose {
                        debug!("Operation pending: {:x}", status);
                    }

                    // fetch DICOM data

                    let dcm = {
                        let mut rsp = scu.receive_pdata();
                        let mut response_data = Vec::new();
                        rsp.read_to_end(&mut response_data)
                            .whatever_context("Failed to read response data")?;

                        InMemDicomObject::read_dataset_with_ts(&response_data[..], &implicit_vr_le)
                            .whatever_context("Could not read response data set")?
                    };

                    println!("------------ Match #{} ------------", i);
                    DumpOptions::new()
                        .dump_object(&dcm)
                        .context(DumpOutputSnafu)?;
                    i += 1;
                } else {
                    warn!("Operation failed (status code {})", status);
                    break;
                }
            }

            pdu @ Pdu::Unknown { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::ReleaseRQ
            | pdu @ Pdu::ReleaseRP
            | pdu @ Pdu::AbortRQ { .. } => {
                error!("Unexpected SCP response: {:?}", pdu);
                let _ = scu.abort();
                std::process::exit(-2);
            }
        }
    }
    let _ = scu.release();

    Ok(())
}

fn find_req_command(
    sop_class_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + even_len(sop_class_uid.len())   // SOP Class UID
            + 8 + 2 // command field
            + 8 + 2 // message ID
            + 8 + 2 // priority
            + 8 + 2, // data set type
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        PrimitiveValue::from(sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        // 0020H: C-FIND-RQ message
        dicom_value!(U16, [0x0020]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        // medium
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::fs;
use std::path::Path;
use std::str::FromStr;

use super::{ExecutionStrategy, Scope};
use common::package_graph::PackageGraph;
use common::ui::{Status, UIWriter, UI};
use error::{Error, Result};
use hcore::error as herror;
use hcore::fs as hfs;
use hcore::package::{all_packages, Identifiable, PackageIdent, PackageInstall};

use hcore::package::list::temp_package_directory;

///
/// Delete a package and all dependencies which are not used by other packages.
/// We do an ordered traverse of the dependencies, updating the graph as we delete a
/// package. This lets us use simple logic where we continually check if the package
/// we're trying to delete currently has no packages depending on it.
///
/// The full logic is:
/// 1. We find the fully qualified package ident and all its dependencies
/// 2. We find all packages on the filesystem and convert them into a graph
/// 3. We do a BFS on the graph to get the dependencies in order
/// 4. We check if the specified package has any reverse deps
///     4a. If there are, we throw an error
///     4b. If not, we delete the package
/// 5. For each dependency we check if there are any packages which depend on it
///     5a. If there are, we skip it
///     5b. If there are not, we delete it from disk and the graph
///
pub fn start(
    ui: &mut UI,
    ident: &PackageIdent,
    fs_root_path: &Path,
    execution_strategy: ExecutionStrategy,
    scope: Scope,
    excludes: Vec<PackageIdent>,
) -> Result<()> {
    // 1.
    let pkg_install = PackageInstall::load(ident, Some(fs_root_path))?;
    let ident = pkg_install.ident();
    ui.begin(format!("Uninstalling {}", &ident))?;

    // 2.
    let packages = load_all_packages(&fs_root_path)?;
    let mut graph = PackageGraph::new();
    graph.build(&packages)?;

    // 3.
    let deps = graph.ordered_deps(&ident);

    // 4.
    match graph.count_rdeps(&ident) {
        None => {
            // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
            // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
            // dependency tree anyway
            ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &ident))?;
        }
        Some(0) => {
            maybe_delete(
                ui,
                &fs_root_path,
                &pkg_install,
                &execution_strategy,
                &excludes,
            )?;
            graph.remove(&ident);
        }
        Some(c) => {
            return Err(Error::CannotRemovePackage(ident.clone(), c));
        }
    }

    // 5.
    let mut count = 0;
    match scope {
        Scope::Package => {
            ui.status(Status::Skipping, "dependencies (--no-deps specified)")?;
        }
        Scope::PackageAndDependencies => {
            for p in &deps {
                match graph.count_rdeps(&p) {
                    None => {
                        // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
                        // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
                        // dependency tree anyway
                        ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &p))?;
                    }
                    Some(0) => {
                        let install = packages.iter().find(|&i| i.ident() == p).unwrap();
                        maybe_delete(ui, &fs_root_path, &install, &execution_strategy, &excludes)?;

                        graph.remove(&p);
                        count = count + 1;
                    }
                    Some(c) => {
                        ui.status(
                            Status::Skipping,
                            format!("{}. It is a dependency of {} packages", &p, c),
                        )?;
                    }
                }
            }
        }
    }

    match execution_strategy {
        ExecutionStrategy::DryRun => {
            ui.end(format!(
                "Would uninstall {} and {} dependencies (Dry run)",
                &ident, count
            ))?;
        }
        ExecutionStrategy::Run => {
            ui.end(format!(
                "Uninstall of {} and {} dependencies complete",
                &ident, count
            ))?;
        }
    };
    Ok(())
}

/// Delete a package from disk, depending upon the ExecutionStrategy supplied
///
/// Returns:
///   Ok(true) - package is deleted
///   Ok(false) - package would be deleted but it's a dry run
///   Err(_) -  IO problem deleting package from filesystem
fn maybe_delete(
    ui: &mut UI,
    fs_root_path: &Path,
    install: &PackageInstall,
    strategy: &ExecutionStrategy,
    excludes: &Vec<PackageIdent>,
) -> Result<bool> {
    let ident = install.ident();
    let pkg_root_path = hfs::pkg_root_path(Some(fs_root_path));
    let hab = PackageIdent::from_str("core/hab")?;

    if ident.satisfies(&hab) {
        ui.status(
            Status::Skipping,
            format!("{}. You can't uninstall core/hab", &ident),
        )?;
        return Ok(false);
    }

    // The excludes list could be looser than the fully qualified idents.  E.g. if core/redis is on the
    // exclude list then we should exclude core/redis/1.1.0/20180608091936.  We use the `Identifiable`
    // trait which supplies this logic for PackageIdents
    let should_exclude = excludes
        .iter()
        .filter(|i| i.satisfies(ident))
        .cloned()
        .count()
        > 0;

    match should_exclude {
        true => {
            ui.status(
                Status::Skipping,
                format!("{}. It is on the exclusion list", &ident),
            )?;
            Ok(false)
        }
        false => match strategy {
            ExecutionStrategy::DryRun => {
                ui.status(Status::DryRunDeleting, &ident)?;
                Ok(false)
            }
            ExecutionStrategy::Run => {
                ui.status(Status::Deleting, &ident)?;
                let pkg_dir = install.installed_path();
                do_clean_delete(&pkg_root_path, &pkg_dir)
            }
        },
    }
}

/// Delete empty parent directories from a given path. don't traverse above
/// the `pkg_root_path`
fn do_clean_delete(pkg_root_path: &Path, real_install_path: &Path) -> Result<bool> {
    // This match will always return Ok(Path) as the install path is always 4 levels
    // below the pkg_root_path
    match real_install_path.parent() {
        Some(real_install_base) => {
            let temp_install_path = temp_package_directory(real_install_path)?;
            fs::rename(&real_install_path, &temp_install_path)?;
            fs::remove_dir_all(&temp_install_path)?;

            for p in real_install_base.ancestors() {
                // Let's be safe and not rm below the package directories
                if p == pkg_root_path {
                    break;
                }
                match p.read_dir() {
                    Ok(contents) => {
                        // This will calculate the amount of items in the directory
                        match contents.collect::<Vec<_>>().len() {
                            0 => fs::remove_dir(&p)?,
                            _ => break,
                        }
                    }
                    Err(e) => return Err(Error::HabitatCore(herror::Error::IO(e))),
                }
            }
            Ok(true)
        }
        None => unreachable!("Install path doesn't have a parent"),
    }
}

fn load_all_packages(fs_root_path: &Path) -> Result<Vec<PackageInstall>> {
    let package_path = hfs::pkg_root_path(Some(&fs_root_path));
    let idents = all_packages(&package_path)?;

    let mut result = Vec::with_capacity(idents.len());
    for i in idents {
        let pkg_install = PackageInstall::load(&i, Some(fs_root_path))?;
        result.push(pkg_install);
    }
    Ok(result)
}
<commit_msg>Review comments - cleanup some logic<commit_after>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::fs;
use std::path::Path;
use std::str::FromStr;

use super::{ExecutionStrategy, Scope};
use common::package_graph::PackageGraph;
use common::ui::{Status, UIWriter, UI};
use error::{Error, Result};
use hcore::error as herror;
use hcore::fs as hfs;
use hcore::package::{all_packages, Identifiable, PackageIdent, PackageInstall};

use hcore::package::list::temp_package_directory;

///
/// Delete a package and all dependencies which are not used by other packages.
/// We do an ordered traverse of the dependencies, updating the graph as we delete a
/// package. This lets us use simple logic where we continually check if the package
/// we're trying to delete currently has no packages depending on it.
///
/// The full logic is:
/// 1. We find the fully qualified package ident and all its dependencies
/// 2. We find all packages on the filesystem and convert them into a graph
/// 3. We do a BFS on the graph to get the dependencies in order
/// 4. We check if the specified package has any reverse deps
///     4a. If there are, we throw an error
///     4b. If not, we delete the package
/// 5. For each dependency we check if there are any packages which depend on it
///     5a. If there are, we skip it
///     5b. If there are not, we delete it from disk and the graph
///
pub fn start(
    ui: &mut UI,
    ident: &PackageIdent,
    fs_root_path: &Path,
    execution_strategy: ExecutionStrategy,
    scope: Scope,
    excludes: Vec<PackageIdent>,
) -> Result<()> {
    // 1.
    let pkg_install = PackageInstall::load(ident, Some(fs_root_path))?;
    let ident = pkg_install.ident();
    ui.begin(format!("Uninstalling {}", &ident))?;

    // 2.
    let packages = load_all_packages(&fs_root_path)?;
    let mut graph = PackageGraph::new();
    graph.build(&packages)?;

    // 3.
    let deps = graph.ordered_deps(&ident);

    // 4.
    match graph.count_rdeps(&ident) {
        None => {
            // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
            // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
            // dependency tree anyway
            ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &ident))?;
        }
        Some(0) => {
            maybe_delete(
                ui,
                &fs_root_path,
                &pkg_install,
                &execution_strategy,
                &excludes,
            )?;
            graph.remove(&ident);
        }
        Some(c) => {
            return Err(Error::CannotRemovePackage(ident.clone(), c));
        }
    }

    // 5.
    let mut count = 0;
    match scope {
        Scope::Package => {
            ui.status(Status::Skipping, "dependencies (--no-deps specified)")?;
        }
        Scope::PackageAndDependencies => {
            for p in &deps {
                match graph.count_rdeps(&p) {
                    None => {
                        // package not in graph - this shouldn't happen but could be a race condition in Step 2 with another hab uninstall. We can
                        // continue as what we wanted (package to be removed) has already happened. We're going to continue and try and delete down through the
                        // dependency tree anyway
                        ui.warn(format!("Tried to find dependant packages of {} but it wasn't in graph.  Maybe another uninstall command was run at the same time?", &p))?;
                    }
                    Some(0) => {
                        let install = packages.iter().find(|&i| i.ident() == p).unwrap();
                        maybe_delete(ui, &fs_root_path, &install, &execution_strategy, &excludes)?;

                        graph.remove(&p);
                        count = count + 1;
                    }
                    Some(c) => {
                        ui.status(
                            Status::Skipping,
                            format!("{}. It is a dependency of {} packages", &p, c),
                        )?;
                    }
                }
            }
        }
    }

    match execution_strategy {
        ExecutionStrategy::DryRun => {
            ui.end(format!(
                "Would uninstall {} and {} dependencies (Dry run)",
                &ident, count
            ))?;
        }
        ExecutionStrategy::Run => {
            ui.end(format!(
                "Uninstall of {} and {} dependencies complete",
                &ident, count
            ))?;
        }
    };
    Ok(())
}

/// Delete a package from disk, depending upon the ExecutionStrategy supplied
///
/// Returns:
///   Ok(true) - package is deleted
///   Ok(false) - package would be deleted but it's a dry run
///   Err(_) -  IO problem deleting package from filesystem
fn maybe_delete(
    ui: &mut UI,
    fs_root_path: &Path,
    install: &PackageInstall,
    strategy: &ExecutionStrategy,
    excludes: &Vec<PackageIdent>,
) -> Result<bool> {
    let ident = install.ident();
    let pkg_root_path = hfs::pkg_root_path(Some(fs_root_path));

    let hab = PackageIdent::from_str("core/hab")?;
    if ident.satisfies(&hab) {
        ui.status(
            Status::Skipping,
            format!("{}. You can't uninstall core/hab", &ident),
        )?;
        return Ok(false);
    }

    // The excludes list could be looser than the fully qualified idents.  E.g. if core/redis is on the
    // exclude list then we should exclude core/redis/1.1.0/20180608091936.  We use the `Identifiable`
    // trait which supplies this logic for PackageIdents
    let should_exclude = excludes.iter().any(|i| i.satisfies(ident));
    if should_exclude {
        ui.status(
            Status::Skipping,
            format!("{}. It is on the exclusion list", &ident),
        )?;
        Ok(false)
    } else {
        match strategy {
            ExecutionStrategy::DryRun => {
                ui.status(Status::DryRunDeleting, &ident)?;
                Ok(false)
            }
            ExecutionStrategy::Run => {
                ui.status(Status::Deleting, &ident)?;
                let pkg_dir = install.installed_path();
                do_clean_delete(&pkg_root_path, &pkg_dir)
            }
        }
    }
}

/// Delete empty parent directories from a given path. don't traverse above
/// the `pkg_root_path`
fn do_clean_delete(pkg_root_path: &Path, real_install_path: &Path) -> Result<bool> {
    // This match will always return Ok(Path) as the install path is always 4 levels
    // below the pkg_root_path
    match real_install_path.parent() {
        Some(real_install_base) => {
            let temp_install_path = temp_package_directory(real_install_path)?;
            fs::rename(&real_install_path, &temp_install_path)?;
            fs::remove_dir_all(&temp_install_path)?;

            for p in real_install_base.ancestors() {
                // Let's be safe and not rm below the package directories
                if p == pkg_root_path {
                    break;
                }
                match p.read_dir() {
                    Ok(contents) => {
                        // This will calculate the amount of items in the directory
                        match contents.collect::<Vec<_>>().len() {
                            0 => fs::remove_dir(&p)?,
                            _ => break,
                        }
                    }
                    Err(e) => return Err(Error::HabitatCore(herror::Error::IO(e))),
                }
            }
            Ok(true)
        }
        None => unreachable!("Install path doesn't have a parent"),
    }
}

fn load_all_packages(fs_root_path: &Path) -> Result<Vec<PackageInstall>> {
    let package_path = hfs::pkg_root_path(Some(&fs_root_path));
    let idents = all_packages(&package_path)?;

    let mut result = Vec::with_capacity(idents.len());
    for i in idents {
        let pkg_install = PackageInstall::load(&i, Some(fs_root_path))?;
        result.push(pkg_install);
    }
    Ok(result)
}
<|endoftext|>


<commit_before>use crate::{command::studio::enter::ARTIFACT_PATH_ENVVAR,
            common::ui::UI,
            error::{Error,
                    Result},
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs::{cache_key_path,
                         CACHE_ARTIFACT_PATH,
                         CACHE_KEY_PATH},
                    os::process,
                    package::target,
                    util::docker},
            license,
            VERSION};
use atty;
use std::{env,
          ffi::{OsStr,
                OsString},
          path::{Path,
                 PathBuf},
          process::{Command,
                    Stdio}};

const DOCKER_IMAGE: &str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &str = "habitat/win-studio";
const DOCKER_IMAGE_ENVVAR: &str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, args: &[OsString]) -> Result<()> {
    let mut args = args.to_vec();
    if args.get(0) == Some(&OsString::from("rm")) {
        return Err(Error::CannotRemoveDockerStudio);
    }

    let docker_cmd = docker::command_path()?;
    let using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(using_windows_containers)?;

    if is_image_present(&docker_cmd, &image) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd, &image)?;
    }

    let mnt_prefix = if using_windows_containers { "c:" } else { "" };

    let local_cache_key_path = match henv::var(CACHE_KEY_PATH_ENV_VAR) {
        Ok(val) => PathBuf::from(val),
        Err(_) => cache_key_path(None::<PathBuf>),
    };
    if !local_cache_key_path.exists() {
        return Err(Error::FileNotFound(format!("{}\nRun `hab setup` to \
                                                create an origin or use \
                                                `hab origin key` to \
                                                configure your keys.",
                                               local_cache_key_path.display())));
    }

    let mut volumes = vec![format!("{}:{}{}",
                                   env::current_dir().unwrap().to_string_lossy(),
                                   mnt_prefix,
                                   "/src"),
                           format!("{}:{}/{}",
                                   local_cache_key_path.display(),
                                   mnt_prefix,
                                   CACHE_KEY_PATH),];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!("{}:{}/{}",
                             cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH));
    }
    if !using_windows_containers
       && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![String::from("DEBUG"),
                            String::from("DO_CHECK"),
                            String::from("HAB_AUTH_TOKEN"),
                            String::from("HAB_BLDR_URL"),
                            String::from("HAB_BLDR_CHANNEL"),
                            String::from("HAB_NOCOLORING"),
                            String::from("HAB_LICENSE"),
                            String::from("HAB_ORIGIN"),
                            String::from("HAB_ORIGIN_KEYS"),
                            String::from("HAB_STUDIO_BACKLINE_PKG"),
                            String::from("HAB_STUDIO_NOSTUDIORC"),
                            String::from("HAB_STUDIO_SUP"),
                            String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
                            String::from("http_proxy"),
                            String::from("https_proxy"),
                            String::from("RUST_LOG"),];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !using_windows_containers {
        check_mounts(&docker_cmd, volumes.iter(), &image)?;
    }
    run_container(docker_cmd,
                  &args,
                  volumes.iter(),
                  env_vars.iter(),
                  image,
                  using_windows_containers)
}

fn is_image_present(docker_cmd: &Path, image: &str) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images").arg(image).arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path, image: &str) -> Result<()> {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
       .arg(image)
       .stdout(Stdio::inherit())
       .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd.spawn()
                    .expect("Docker command failed to spawn")
                    .wait_with_output()
                    .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", image);
    } else {
        debug!("Pulling Docker image '{}' failed with exit code: {:?}",
               image, result.status);

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(image.to_string()));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(image.to_string()));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I, image: &str) -> Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<OsStr>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image.into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd).args(&cmd_args)
                                                 .output()
                                                 .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
       && (stderr.as_str().contains("Mounts denied")
           || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(docker_cmd: PathBuf,
                             args: &[OsString],
                             volumes: I,
                             env_vars: J,
                             image: String,
                             using_windows_containers: bool)
                             -> Result<()>
    where I: IntoIterator<Item = S>,
          J: IntoIterator<Item = T>,
          S: AsRef<OsStr>,
          T: AsRef<str>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];

    if !using_windows_containers {
        cmd_args.push("--privileged".into());
    }

    if atty::is(atty::Stream::Stderr) || atty::is(atty::Stream::Stdout) {
        cmd_args.push("--tty".into());
        cmd_args.push("--interactive".into());
    }

    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts
            .split(' ')
            .map(|v| v.into())
            // Ensure we're not passing something like `--tty` again here.
            .filter(|v| !cmd_args.contains(v))
            .collect::<Vec<_>>();

        if !opts.is_empty() {
            debug!("Adding extra Docker options from {} = {:?}",
                   DOCKER_OPTS_ENVVAR, opts);
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }

    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        } else if var.as_ref() == "HAB_LICENSE" && license::license_exists() {
            debug!("Hab license already accepted. Setting container env var: \
                    HAB_LICENSE=accept-no-persist");
            cmd_args.push("--env".into());
            cmd_args.push("HAB_LICENSE=accept-no-persist".to_string().into());
        }
    }

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }

    cmd_args.push(image.into());
    cmd_args.extend_from_slice(args);

    if using_windows_containers {
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }

    unset_proxy_env_vars();
    process::become_command(docker_cmd, &cmd_args)?;
    Ok(())
}

fn unset_proxy_env_vars() {
    for var in &["http_proxy", "https_proxy"] {
        if henv::var(var).is_ok() {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

fn image_identifier_for_active_target(using_windows_containers: bool) -> Result<String> {
    let windows_base_tag = if using_windows_containers {
        Some(docker::default_base_tag_for_host()?)
    } else {
        None
    };
    Ok(image_identifier(windows_base_tag, target::PackageTarget::active_target()))
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(windows_base_tag: Option<&str>, target: target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target, tag) = if let Some(t) = windows_base_tag {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS, format!("{}-{}", t, version[0]))
    } else {
        let t = match target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2,
            _ => target::X86_64_LINUX,
        };
        (DOCKER_IMAGE, t, version[0].to_string())
    };

    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or_else(|_| format!("{}-{}:{}", img, studio_target, tag))
}

#[cfg(test)]
mod tests {
    use super::{image_identifier,
                DOCKER_IMAGE,
                DOCKER_WINDOWS_IMAGE};
    use crate::VERSION;

    use crate::hcore::package::target;

    #[test]
    fn retrieve_image_identifier() {
        assert_eq!(image_identifier(None, target::X86_64_DARWIN),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(None, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(None, target::X86_64_LINUX_KERNEL2),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION));
        assert_eq!(image_identifier(None, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(Some("ltsc2016"), target::X86_64_WINDOWS),
                   format!("{}-{}:{}-{}",
                           DOCKER_WINDOWS_IMAGE, "x86_64-windows", "ltsc2016", VERSION));
        assert_eq!(image_identifier(Some("ltsc2016"), target::X86_64_LINUX),
                   format!("{}-{}:{}-{}",
                           DOCKER_WINDOWS_IMAGE, "x86_64-windows", "ltsc2016", VERSION));
    }
}
<commit_msg>[hab] Allow flexible whitespace in HAB_DOCKER_OPTS<commit_after>use crate::{command::studio::enter::ARTIFACT_PATH_ENVVAR,
            common::ui::UI,
            error::{Error,
                    Result},
            hcore::{crypto::CACHE_KEY_PATH_ENV_VAR,
                    env as henv,
                    fs::{cache_key_path,
                         CACHE_ARTIFACT_PATH,
                         CACHE_KEY_PATH},
                    os::process,
                    package::target,
                    util::docker},
            license,
            VERSION};
use atty;
use std::{env,
          ffi::{OsStr,
                OsString},
          path::{Path,
                 PathBuf},
          process::{Command,
                    Stdio}};

const DOCKER_IMAGE: &str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &str = "habitat/win-studio";
const DOCKER_IMAGE_ENVVAR: &str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, args: &[OsString]) -> Result<()> {
    let mut args = args.to_vec();
    if args.get(0) == Some(&OsString::from("rm")) {
        return Err(Error::CannotRemoveDockerStudio);
    }

    let docker_cmd = docker::command_path()?;
    let using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(using_windows_containers)?;

    if is_image_present(&docker_cmd, &image) {
        debug!("Found Studio Docker image locally.");
    } else {
        debug!("Failed to find Studio Docker image locally.");
        pull_image(&docker_cmd, &image)?;
    }

    let mnt_prefix = if using_windows_containers { "c:" } else { "" };

    let local_cache_key_path = match henv::var(CACHE_KEY_PATH_ENV_VAR) {
        Ok(val) => PathBuf::from(val),
        Err(_) => cache_key_path(None::<PathBuf>),
    };
    if !local_cache_key_path.exists() {
        return Err(Error::FileNotFound(format!("{}\nRun `hab setup` to \
                                                create an origin or use \
                                                `hab origin key` to \
                                                configure your keys.",
                                               local_cache_key_path.display())));
    }

    let mut volumes = vec![format!("{}:{}{}",
                                   env::current_dir().unwrap().to_string_lossy(),
                                   mnt_prefix,
                                   "/src"),
                           format!("{}:{}/{}",
                                   local_cache_key_path.display(),
                                   mnt_prefix,
                                   CACHE_KEY_PATH),];
    if let Ok(cache_artifact_path) = henv::var(ARTIFACT_PATH_ENVVAR) {
        volumes.push(format!("{}:{}/{}",
                             cache_artifact_path, mnt_prefix, CACHE_ARTIFACT_PATH));
    }
    if !using_windows_containers
       && (Path::new(DOCKER_SOCKET).exists() || cfg!(target_os = "windows"))
    {
        volumes.push(format!("{}:{}", DOCKER_SOCKET, DOCKER_SOCKET));
    }

    let mut env_vars = vec![String::from("DEBUG"),
                            String::from("DO_CHECK"),
                            String::from("HAB_AUTH_TOKEN"),
                            String::from("HAB_BLDR_URL"),
                            String::from("HAB_BLDR_CHANNEL"),
                            String::from("HAB_NOCOLORING"),
                            String::from("HAB_LICENSE"),
                            String::from("HAB_ORIGIN"),
                            String::from("HAB_ORIGIN_KEYS"),
                            String::from("HAB_STUDIO_BACKLINE_PKG"),
                            String::from("HAB_STUDIO_NOSTUDIORC"),
                            String::from("HAB_STUDIO_SUP"),
                            String::from("HAB_UPDATE_STRATEGY_FREQUENCY_MS"),
                            String::from("http_proxy"),
                            String::from("https_proxy"),
                            String::from("RUST_LOG"),];

    for (key, _) in env::vars() {
        if key.starts_with(HAB_STUDIO_SECRET) {
            env_vars.push(key);
        }
    }

    // We need to strip out the -D if it exists to avoid
    // it getting passed to the sup on entering the studio
    let to_cull = OsString::from("-D");
    if let Some(index) = args.iter().position(|x| *x == to_cull) {
        args.remove(index);
    }

    // Windows containers do not use filesystem sharing for
    // local mounts
    if !using_windows_containers {
        check_mounts(&docker_cmd, volumes.iter(), &image)?;
    }
    run_container(docker_cmd,
                  &args,
                  volumes.iter(),
                  env_vars.iter(),
                  image,
                  using_windows_containers)
}

fn is_image_present(docker_cmd: &Path, image: &str) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images").arg(image).arg("-q");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");

    !String::from_utf8_lossy(&result.stdout).as_ref().is_empty()
}

fn is_serving_windows_containers(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("version").arg("--format='{{.Server.Os}}'");
    debug!("Running command: {:?}", cmd);
    let result = cmd.output().expect("Docker command failed to spawn");
    String::from_utf8_lossy(&result.stdout).contains("windows")
}

fn pull_image(docker_cmd: &Path, image: &str) -> Result<()> {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("pull")
       .arg(image)
       .stdout(Stdio::inherit())
       .stderr(Stdio::inherit());
    debug!("Running command: {:?}", cmd);
    let result = cmd.spawn()
                    .expect("Docker command failed to spawn")
                    .wait_with_output()
                    .expect("Failed to wait on child process");

    if result.status.success() {
        debug!("Docker image '{}' is present locally.", image);
    } else {
        debug!("Pulling Docker image '{}' failed with exit code: {:?}",
               image, result.status);

        let err_output = String::from_utf8_lossy(&result.stderr);

        if err_output.contains("image") && err_output.contains("not found") {
            return Err(Error::DockerImageNotFound(image.to_string()));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(image.to_string()));
        }
    }

    Ok(())
}

/// Checks whether or not the volume mounts are working.
///
/// We need to ensure that filesystem sharing has been enabled, otherwise the user will be
/// greeted with a horrible error message that's difficult to make sense of. To mitigate this,
/// we check the studio version. This will cause Docker to go through the mounting steps, so we
/// can watch stderr for failure, but has the advantage of not requiring a TTY.
fn check_mounts<I, S>(docker_cmd: &Path, volumes: I, image: &str) -> Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<OsStr>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image.into());
    cmd_args.push("-V".into());
    let version_output = Command::new(docker_cmd).args(&cmd_args)
                                                 .output()
                                                 .expect("docker failed to start");

    let stderr = String::from_utf8(version_output.stderr).unwrap();
    if !stderr.is_empty()
       && (stderr.as_str().contains("Mounts denied")
           || stderr.as_str().contains("drive is not shared"))
    {
        return Err(Error::DockerFileSharingNotEnabled);
    }
    Ok(())
}

fn run_container<I, J, S, T>(docker_cmd: PathBuf,
                             args: &[OsString],
                             volumes: I,
                             env_vars: J,
                             image: String,
                             using_windows_containers: bool)
                             -> Result<()>
    where I: IntoIterator<Item = S>,
          J: IntoIterator<Item = T>,
          S: AsRef<OsStr>,
          T: AsRef<str>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];

    if !using_windows_containers {
        cmd_args.push("--privileged".into());
    }

    if atty::is(atty::Stream::Stderr) || atty::is(atty::Stream::Stdout) {
        cmd_args.push("--tty".into());
        cmd_args.push("--interactive".into());
    }

    if let Ok(opts) = henv::var(DOCKER_OPTS_ENVVAR) {
        let opts = opts
            .split_whitespace()
            .map(|v| v.into())
            // Ensure we're not passing something like `--tty` again here.
            .filter(|v| !cmd_args.contains(v))
            .collect::<Vec<_>>();

        if !opts.is_empty() {
            debug!("Adding extra Docker options from {} = {:?}",
                   DOCKER_OPTS_ENVVAR, opts);
            cmd_args.extend_from_slice(opts.as_slice());
        }
    }

    for var in env_vars {
        if let Ok(val) = henv::var(var.as_ref()) {
            debug!("Setting container env var: {:?}='{}'", var.as_ref(), val);
            cmd_args.push("--env".into());
            cmd_args.push(format!("{}={}", var.as_ref(), val).into());
        } else if var.as_ref() == "HAB_LICENSE" && license::license_exists() {
            debug!("Hab license already accepted. Setting container env var: \
                    HAB_LICENSE=accept-no-persist");
            cmd_args.push("--env".into());
            cmd_args.push("HAB_LICENSE=accept-no-persist".to_string().into());
        }
    }

    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }

    cmd_args.push(image.into());
    cmd_args.extend_from_slice(args);

    if using_windows_containers {
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }

    unset_proxy_env_vars();
    process::become_command(docker_cmd, &cmd_args)?;
    Ok(())
}

fn unset_proxy_env_vars() {
    for var in &["http_proxy", "https_proxy"] {
        if henv::var(var).is_ok() {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

fn image_identifier_for_active_target(using_windows_containers: bool) -> Result<String> {
    let windows_base_tag = if using_windows_containers {
        Some(docker::default_base_tag_for_host()?)
    } else {
        None
    };
    Ok(image_identifier(windows_base_tag, target::PackageTarget::active_target()))
}

/// Returns the Docker Studio image with tag for the desired version which corresponds to the
/// same version (minus release) as this program.
fn image_identifier(windows_base_tag: Option<&str>, target: target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target, tag) = if let Some(t) = windows_base_tag {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS, format!("{}-{}", t, version[0]))
    } else {
        let t = match target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2,
            _ => target::X86_64_LINUX,
        };
        (DOCKER_IMAGE, t, version[0].to_string())
    };

    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or_else(|_| format!("{}-{}:{}", img, studio_target, tag))
}

#[cfg(test)]
mod tests {
    use super::{image_identifier,
                DOCKER_IMAGE,
                DOCKER_WINDOWS_IMAGE};
    use crate::VERSION;

    use crate::hcore::package::target;

    #[test]
    fn retrieve_image_identifier() {
        assert_eq!(image_identifier(None, target::X86_64_DARWIN),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(None, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(None, target::X86_64_LINUX_KERNEL2),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION));
        assert_eq!(image_identifier(None, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(Some("ltsc2016"), target::X86_64_WINDOWS),
                   format!("{}-{}:{}-{}",
                           DOCKER_WINDOWS_IMAGE, "x86_64-windows", "ltsc2016", VERSION));
        assert_eq!(image_identifier(Some("ltsc2016"), target::X86_64_LINUX),
                   format!("{}-{}:{}-{}",
                           DOCKER_WINDOWS_IMAGE, "x86_64-windows", "ltsc2016", VERSION));
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::env;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::process::Command;

fn main() {
    println!("cargo:rerun-if-changed=.git/HEAD");
    println!("cargo:rerun-if-env-changed=CFG_RELEASE_CHANNEL");

    let out_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());

    File::create(out_dir.join("commit-info.txt"))
        .unwrap()
        .write_all(commit_info().as_bytes())
        .unwrap();
}

// Try to get hash and date of the last commit on a best effort basis. If anything goes wrong
// (git not installed or if this is not a git repository) just return an empty string.
fn commit_info() -> String {
    match (channel(), commit_hash(), commit_date()) {
        (channel, Some(hash), Some(date)) => {
            format!("{} ({} {})", channel, hash.trim_right(), date)
        }
        _ => String::new(),
    }
}

fn channel() -> String {
    if let Ok(channel) = env::var("CFG_RELEASE_CHANNEL") {
        channel
    } else {
        "nightly".to_owned()
    }
}

fn commit_hash() -> Option<String> {
    Command::new("git")
        .args(&["rev-parse", "--short", "HEAD"])
        .output()
        .ok()
        .and_then(|r| String::from_utf8(r.stdout).ok())
}

fn commit_date() -> Option<String> {
    Command::new("git")
        .args(&["log", "-1", "--date=short", "--pretty=format:%cd"])
        .output()
        .ok()
        .and_then(|r| String::from_utf8(r.stdout).ok())
}
<commit_msg>Ignore checking inexistent .git/HEAD in build.rs<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::env;
use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Command;

fn main() {
    // Only check .git/HEAD dirty status if it exists - doing so when
    // building dependent crates may lead to false positives and rebuilds
    if Path::new(".git/HEAD").exists() {
        println!("cargo:rerun-if-changed=.git/HEAD");
    }

    println!("cargo:rerun-if-env-changed=CFG_RELEASE_CHANNEL");

    let out_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());

    File::create(out_dir.join("commit-info.txt"))
        .unwrap()
        .write_all(commit_info().as_bytes())
        .unwrap();
}

// Try to get hash and date of the last commit on a best effort basis. If anything goes wrong
// (git not installed or if this is not a git repository) just return an empty string.
fn commit_info() -> String {
    match (channel(), commit_hash(), commit_date()) {
        (channel, Some(hash), Some(date)) => {
            format!("{} ({} {})", channel, hash.trim_right(), date)
        }
        _ => String::new(),
    }
}

fn channel() -> String {
    if let Ok(channel) = env::var("CFG_RELEASE_CHANNEL") {
        channel
    } else {
        "nightly".to_owned()
    }
}

fn commit_hash() -> Option<String> {
    Command::new("git")
        .args(&["rev-parse", "--short", "HEAD"])
        .output()
        .ok()
        .and_then(|r| String::from_utf8(r.stdout).ok())
}

fn commit_date() -> Option<String> {
    Command::new("git")
        .args(&["log", "-1", "--date=short", "--pretty=format:%cd"])
        .output()
        .ok()
        .and_then(|r| String::from_utf8(r.stdout).ok())
}
<|endoftext|>


<commit_before>#![feature(io)]
#![feature(path)]

use std::fs::File;
use std::io::Write;
use std::process::Command;
use std::path;
use std::env;

fn main()
{
    let dir_var = env::var("OUT_DIR").unwrap();
    let dir = path::Path::new(&dir_var).join("codegen");
    let out_file = dir.to_str().unwrap();
    Command::new("gcc").arg("src/codegen/type_size.c")
                       .arg("-o")
                       .arg(out_file)
                       .status()
                       .unwrap();
    let output = Command::new(out_file).output().unwrap();
    if !output.status.success() {
        panic!("{}", String::from_utf8_lossy(&output.stderr[..]));
    }
    let mut f = match File::create("src/type_size.rs") {
        Ok(f) => f,
        Err(e) => panic!("{:?}", e)
    };
    f.write_all(&output.stdout[..]).unwrap();
}
<commit_msg>Remove all unwraps from build script<commit_after>use std::fs::File;
use std::io::Write;
use std::process::{Command, Output};
use std::path;
use std::env;

fn run_cmd(name: &str, command: &mut Command) -> Output {
    let res = command.output();
    match res {
        Ok(output) => if !output.status.success() {
            panic!("{} failed:\n{}", name, String::from_utf8_lossy(&output.stderr[..]))
        } else {
            output
        },
        Err(err) =>
            panic!("{} failed:\n{}", name, err)
    }
}

fn main()
{
    let dir_var = match env::var("OUT_DIR") {
        Ok(dir) => dir,
        Err(e) => panic!("Could not get env value OUT_DIR: {}", e)
    };

    let out_file = path::Path::new(&dir_var).join("codegen");
    let mut gcc_cmd = Command::new("gcc");
    gcc_cmd.arg("src/codegen/type_size.c").arg("-o").arg(&out_file);

    run_cmd("gcc", &mut gcc_cmd);

    let mut codegen_cmd = Command::new(&out_file);
    let output = run_cmd("codegen", &mut codegen_cmd);

    let mut f = match File::create("src/type_size.rs") {
        Ok(f) => f,
        Err(e) => panic!("Could not open file src/type_size.rs: {}", e)
    };

    match f.write_all(&output.stdout[..]) {
        Err(e) => panic!("Could not write to src/type_size.rs: {}", e),
        Ok(_) => ()
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::path::{Path};
use std::process::{Command};

const SEARCH_OSX: &'static [&'static str] = &[
    "/usr/local/opt/llvm/lib",
    "/Library/Developer/CommandLineTools/usr/lib",
    "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib",
];

const SEARCH_WINDOWS: &'static [&'static str] = &[
    "C:\\Program Files\\LLVM\\bin",
    "C:\\Program Files\\LLVM\\lib",
];

fn error(prefix: &str) -> ! {
    panic!("{:?}, set the LIBCLANG_PATH environment variable to your system's libclang directory", prefix);
}

fn run(command: &str, arguments: &[&str]) -> Option<String> {
    Command::new(command).args(arguments).output().map(|o| {
        String::from_utf8_lossy(&o.stdout).into_owned()
    }).ok()
}

fn find_libclang() -> Option<(String, String)> {
    let search = if let Ok(directory) = env::var("LIBCLANG_PATH") {
        vec![directory]
    } else {
        run("llvm-config", &["--libdir"]).map(|d| vec![d]).unwrap_or_else(|| {
            if cfg!(target_os="osx") {
                SEARCH_OSX
            } else if cfg!(target_os="windows") {
                SEARCH_WINDOWS
            } else {
                if cfg!(any(target_os="freebsd", target_os="linux")) {
                    error("llvm-config was not found");
                } else {
                    error("unsupported operating system");
                }
            }.into_iter().map(|s| s.to_string()).collect()
        })
    };

    let library = format!("{}clang{}", env::consts::DLL_PREFIX, env::consts::DLL_SUFFIX);
    let directory = search.into_iter().find(|d| Path::new(&d).join(&library).exists());

    if directory.is_none() && cfg!(target_os="linux") {
        if let Some(output) = run("llvm-config", &["--libdir"]) {
            output.lines().next().map(|l| (l.into(), library))
        } else {
            None
        }
    } else {
        directory.map(|d| (d, library))
    }
}

const LIBRARIES: &'static [&'static str] = &[
    "LLVMAnalysis",
    "LLVMBitReader",
    "LLVMCore",
    "LLVMLTO",
    "LLVMLinker",
    "LLVMMC",
    "LLVMMCParser",
    "LLVMObjCARCOpts",
    "LLVMObject",
    "LLVMOption",
    "LLVMScalarOpts",
    "LLVMSupport",
    "LLVMTarget",
    "LLVMTransformUtils",
    "LLVMVectorize",
    "LLVMipa",
    "LLVMipo",
    "clang",
    "clangARCMigrate",
    "clangAST",
    "clangASTMatchers",
    "clangAnalysis",
    "clangBasic",
    "clangDriver",
    "clangEdit",
    "clangFormat",
    "clangFrontend",
    "clangIndex",
    "clangLex",
    "clangParse",
    "clangRewrite",
    "clangRewriteFrontend",
    "clangSema",
    "clangSerialization",
    "clangStaticAnalyzerCheckers",
    "clangStaticAnalyzerCore",
    "clangStaticAnalyzerFrontend",
    "clangTooling",
];

fn get_libraries() -> Vec<String> {
    run("llvm-config", &["--libs"]).map(|o| {
        o.split_whitespace().filter_map(|p| {
            Path::new(p).file_stem().map(|l| l.to_string_lossy()[2..].into())
        }).collect()
    }).unwrap_or_else(|| LIBRARIES.iter().map(|l| (*l).into()).collect())
}

fn main() {
    let (directory, file) = match find_libclang() {
        Some((directory, file)) => (directory, file),
        _ => error("unable to find libclang"),
    };

    if cfg!(feature="static") {
        print!("cargo:rustc-flags=");
        if let Ok(directory) = env::var("LIBCLANG_STATIC_PATH") {
            print!("-L {} ", directory);
        }
        for library in get_libraries() {
            print!("-l static={} ", library)
        }
        println!("-L {} -l ncursesw -l z -l stdc++", directory);
    } else {
        println!("cargo:rustc-link-search={}", directory);
        println!("cargo:rustc-link-lib=dylib=:{}", file);
    }
}
<commit_msg>Readd Linux backup search paths<commit_after>use std::env;
use std::path::{Path};
use std::process::{Command};

const SEARCH_LINUX: &'static [&'static str] = &[
    "/usr/lib",
    "/usr/lib/llvm",
    "/usr/lib/llvm-3.8/lib",
    "/usr/lib/llvm-3.7/lib",
    "/usr/lib/llvm-3.6/lib",
    "/usr/lib/llvm-3.5/lib",
    "/usr/lib/llvm-3.4/lib",
    "/usr/lib64/llvm",
    "/usr/lib/x86_64-linux-gnu",
    "/usr/local/lib",
];

const SEARCH_OSX: &'static [&'static str] = &[
    "/usr/local/opt/llvm/lib",
    "/Library/Developer/CommandLineTools/usr/lib",
    "/Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/lib",
];

const SEARCH_WINDOWS: &'static [&'static str] = &[
    "C:\\Program Files\\LLVM\\bin",
    "C:\\Program Files\\LLVM\\lib",
];

fn error(prefix: &str) -> ! {
    panic!("{:?}, set the LIBCLANG_PATH environment variable to your system's libclang directory", prefix);
}

fn run(command: &str, arguments: &[&str]) -> Option<String> {
    Command::new(command).args(arguments).output().map(|o| {
        String::from_utf8_lossy(&o.stdout).into_owned()
    }).ok()
}

fn find_libclang() -> Option<(String, String)> {
    let search = if let Ok(directory) = env::var("LIBCLANG_PATH") {
        vec![directory]
    } else {
        run("llvm-config", &["--libdir"]).map(|d| vec![d]).unwrap_or_else(|| {
            if cfg!(any(target_os="freebsd", target_os="linux")) {
                SEARCH_LINUX
            } else if cfg!(target_os="osx") {
                SEARCH_OSX
            } else if cfg!(target_os="windows") {
                SEARCH_WINDOWS
            } else {
                error("unsupported operating system");
            }.into_iter().map(|s| s.to_string()).collect()
        })
    };

    let library = format!("{}clang{}", env::consts::DLL_PREFIX, env::consts::DLL_SUFFIX);
    let directory = search.into_iter().find(|d| Path::new(&d).join(&library).exists());

    if directory.is_none() && cfg!(target_os="linux") {
        if let Some(output) = run("llvm-config", &["--libdir"]) {
            output.lines().next().map(|l| (l.into(), library))
        } else {
            None
        }
    } else {
        directory.map(|d| (d, library))
    }
}

const LIBRARIES: &'static [&'static str] = &[
    "LLVMAnalysis",
    "LLVMBitReader",
    "LLVMCore",
    "LLVMLTO",
    "LLVMLinker",
    "LLVMMC",
    "LLVMMCParser",
    "LLVMObjCARCOpts",
    "LLVMObject",
    "LLVMOption",
    "LLVMScalarOpts",
    "LLVMSupport",
    "LLVMTarget",
    "LLVMTransformUtils",
    "LLVMVectorize",
    "LLVMipa",
    "LLVMipo",
    "clang",
    "clangARCMigrate",
    "clangAST",
    "clangASTMatchers",
    "clangAnalysis",
    "clangBasic",
    "clangDriver",
    "clangEdit",
    "clangFormat",
    "clangFrontend",
    "clangIndex",
    "clangLex",
    "clangParse",
    "clangRewrite",
    "clangRewriteFrontend",
    "clangSema",
    "clangSerialization",
    "clangStaticAnalyzerCheckers",
    "clangStaticAnalyzerCore",
    "clangStaticAnalyzerFrontend",
    "clangTooling",
];

fn get_libraries() -> Vec<String> {
    run("llvm-config", &["--libs"]).map(|o| {
        o.split_whitespace().filter_map(|p| {
            Path::new(p).file_stem().map(|l| l.to_string_lossy()[2..].into())
        }).collect()
    }).unwrap_or_else(|| LIBRARIES.iter().map(|l| (*l).into()).collect())
}

fn main() {
    let (directory, file) = match find_libclang() {
        Some((directory, file)) => (directory, file),
        _ => error("unable to find libclang"),
    };

    if cfg!(feature="static") {
        print!("cargo:rustc-flags=");
        if let Ok(directory) = env::var("LIBCLANG_STATIC_PATH") {
            print!("-L {} ", directory);
        }
        for library in get_libraries() {
            print!("-l static={} ", library)
        }
        println!("-L {} -l ncursesw -l z -l stdc++", directory);
    } else {
        println!("cargo:rustc-link-search={}", directory);
        println!("cargo:rustc-link-lib=dylib=:{}", file);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2016 Nathan Fiedler
 *
 * 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.
 */

use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::process::Command;

static HEADER: &'static str = "#include <wand/MagickWand.h>\n";
static LIBPATH: &'static str = "/Library/Developer/CommandLineTools/usr/lib";

fn main() {
    //
    // If the MagickWand bindings are missing, generate them using
    // rust-bindgen.
    //
    let bindings_path = Path::new("src/bindings.rs");
    if !bindings_path.exists() {
        let bindgen_path = Path::new("rust-bindgen");
        if !bindgen_path.exists() {
            Command::new("git")
                    .arg("clone")
                    .arg("https://github.com/crabtw/rust-bindgen.git")
                    .status().unwrap();
            // Checkout a version of rust-bindgen that is known to work;
            // more recent versions produce code that does not compile (the
            // commit after 8a51860 changes the way enums are generated).
            Command::new("git")
                    .arg("checkout")
                    .arg("8a51860")
                    .current_dir("rust-bindgen")
                    .status().unwrap();
            Command::new("cargo")
                    .arg("build")
                    .current_dir("rust-bindgen")
                    .status().unwrap();
        }
        // Create the header file that rust-bindgen needs as input.
        let mut gen_h = match File::create("gen.h") {
            Err(why) => panic!("could not create gen.h file: {}", Error::description(&why)),
            Ok(file) => file
        };
        match gen_h.write_all(HEADER.as_bytes()) {
            Err(why) => panic!("could not write to gen.h: {}", Error::description(&why)),
            Ok(_)    => ()
        };
        // Get the compiler and linker flags for the MagickWand library.
        let mw_cflags_output = Command::new("MagickWand-config")
                .arg("--cflags")
                .output().unwrap();
        let mw_cflags = std::str::from_utf8(&mw_cflags_output.stdout).unwrap().trim();
        let mw_cflags_arr: Vec<&str> = mw_cflags.split_whitespace().collect();
        let mw_ldflags_output = Command::new("MagickWand-config")
                .arg("--ldflags")
                .output().unwrap();
        let mw_ldflags = std::str::from_utf8(&mw_ldflags_output.stdout).unwrap().trim();
        let mw_ldflags_arr: Vec<&str> = mw_ldflags.split_whitespace().collect();
        // Combine all of that in the invocation of rust-bindgen.
        let mut cmd = &mut Command::new("./rust-bindgen/target/debug/bindgen");
        if cfg!(target_os = "macos") {
            // Mac requires that the xcode tools are installed so that
            // rustc can find the clang.dylib file. See also issue
            // https://github.com/crabtw/rust-bindgen/issues/89
            let lib_path = Path::new(LIBPATH);
            if !lib_path.exists() {
                panic!("missing {}, run xcode-select --install", LIBPATH);
            }
            cmd.env("DYLD_LIBRARY_PATH", LIBPATH);
        }
        cmd.args(&mw_cflags_arr[..])
           .arg("-builtins")
           .arg("-o")
           .arg("src/bindings.rs")
           .args(&mw_ldflags_arr[..])
           .arg("gen.h")
           .status().unwrap();
        // how to get the output of the command...
        // let output = Commad::new(...).output().unwrap();
        // let out = std::str::from_utf8(&output.stdout).unwrap();
        // println!("cargo:output={}", out);
        // let err = std::str::from_utf8(&output.stderr).unwrap();
        // println!("cargo:error={}", err);
        match std::fs::remove_file("gen.h") {
            Err(why) => panic!("could not remove gen.h: {}", Error::description(&why)),
            Ok(_)    => ()
        }
    }
    // For the sake of easily building and testing on Mac, include the path
    // to MagickWand. Chances are MagickWand is in /usr/local/lib, or
    // somewhere else that rustc can find it.
    println!("cargo:rustc-link-search=native=/usr/local/lib");
}
<commit_msg>Report missing MagickWand-config in build.rs<commit_after>/*
 * Copyright 2016 Nathan Fiedler
 *
 * 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.
 */

use std::error::Error;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::process::Command;

static HEADER: &'static str = "#include <wand/MagickWand.h>\n";
static LIBPATH: &'static str = "/Library/Developer/CommandLineTools/usr/lib";

fn main() {
    //
    // If the MagickWand bindings are missing, generate them using
    // rust-bindgen.
    //
    let bindings_path = Path::new("src/bindings.rs");
    if !bindings_path.exists() {
        let bindgen_path = Path::new("rust-bindgen");
        if !bindgen_path.exists() {
            Command::new("git")
                    .arg("clone")
                    .arg("https://github.com/crabtw/rust-bindgen.git")
                    .status().unwrap();
            // Checkout a version of rust-bindgen that is known to work;
            // more recent versions produce code that does not compile (the
            // commit after 8a51860 changes the way enums are generated).
            Command::new("git")
                    .arg("checkout")
                    .arg("8a51860")
                    .current_dir("rust-bindgen")
                    .status().unwrap();
            Command::new("cargo")
                    .arg("build")
                    .current_dir("rust-bindgen")
                    .status().unwrap();
        }
        // Ensure MagickWand-config is in the PATH and report clearly if not.
        if !Command::new("which").arg("MagickWand-config").status().unwrap().success() {
            panic!("MagickWand-config not in the PATH, please install ImageMagick");
        }
        // Create the header file that rust-bindgen needs as input.
        let mut gen_h = match File::create("gen.h") {
            Err(why) => panic!("could not create gen.h file: {}", Error::description(&why)),
            Ok(file) => file
        };
        match gen_h.write_all(HEADER.as_bytes()) {
            Err(why) => panic!("could not write to gen.h: {}", Error::description(&why)),
            Ok(_)    => ()
        };
        // Get the compiler and linker flags for the MagickWand library.
        let mw_cflags_output = Command::new("MagickWand-config")
                .arg("--cflags")
                .output().unwrap();
        let mw_cflags = std::str::from_utf8(&mw_cflags_output.stdout).unwrap().trim();
        let mw_cflags_arr: Vec<&str> = mw_cflags.split_whitespace().collect();
        let mw_ldflags_output = Command::new("MagickWand-config")
                .arg("--ldflags")
                .output().unwrap();
        let mw_ldflags = std::str::from_utf8(&mw_ldflags_output.stdout).unwrap().trim();
        let mw_ldflags_arr: Vec<&str> = mw_ldflags.split_whitespace().collect();
        // Combine all of that in the invocation of rust-bindgen.
        let mut cmd = &mut Command::new("./rust-bindgen/target/debug/bindgen");
        if cfg!(target_os = "macos") {
            // Mac requires that the xcode tools are installed so that
            // rustc can find the clang.dylib file. See also issue
            // https://github.com/crabtw/rust-bindgen/issues/89
            let lib_path = Path::new(LIBPATH);
            if !lib_path.exists() {
                panic!("missing {}, run xcode-select --install", LIBPATH);
            }
            cmd.env("DYLD_LIBRARY_PATH", LIBPATH);
        }
        cmd.args(&mw_cflags_arr[..])
           .arg("-builtins")
           .arg("-o")
           .arg("src/bindings.rs")
           .args(&mw_ldflags_arr[..])
           .arg("gen.h")
           .status().unwrap();
        // how to get the output of the command...
        // let output = Commad::new(...).output().unwrap();
        // let out = std::str::from_utf8(&output.stdout).unwrap();
        // println!("cargo:output={}", out);
        // let err = std::str::from_utf8(&output.stderr).unwrap();
        // println!("cargo:error={}", err);
        match std::fs::remove_file("gen.h") {
            Err(why) => panic!("could not remove gen.h: {}", Error::description(&why)),
            Ok(_)    => ()
        }
    }
    // For the sake of easily building and testing on Mac, include the path
    // to MagickWand. Chances are MagickWand is in /usr/local/lib, or
    // somewhere else that rustc can find it.
    println!("cargo:rustc-link-search=native=/usr/local/lib");
}
<|endoftext|>