<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Utilities for providing the completion language feature
use crate::lsp::Position;
use common::{FileKey, Span};
use graphql_syntax::{parse, Document, GraphQLSource};
use interner::StringKey;
use log::info;

use relay_compiler::Programs;

use crate::lsp::{
    CompletionItem, CompletionParams, CompletionResponse, Connection, Message, ServerRequestId,
    ServerResponse, TextDocumentPositionParams, Url,
};
use schema::{Schema, Type, TypeWithFields};

use graphql_syntax::{
    ExecutableDefinition, InlineFragment, LinkedField, List, OperationDefinition, OperationKind,
    Selection,
};

pub type GraphQLSourceCache = std::collections::HashMap<Url, Vec<GraphQLSource>>;

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum CompletionKind {
    FieldName,
    FragmentSpread,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct CompletionRequest {
    /// The type of the completion request we're responding to
    kind: CompletionKind,
    /// A list of type metadata that we can use to resolve the leaf
    /// type the request is being made against
    type_path: Vec<TypePathItem>,
}

impl Default for CompletionRequest {
    fn default() -> Self {
        CompletionRequest {
            kind: CompletionKind::FieldName,
            type_path: vec![],
        }
    }
}

impl CompletionRequest {
    fn add_type(&mut self, type_path_item: TypePathItem) {
        self.type_path.push(type_path_item)
    }

    /// Returns the leaf type, which is the type that the completion request is being made against.
    fn resolve_leaf_type(self, schema: &Schema) -> Type {
        let mut type_path = self.type_path;
        type_path.reverse();
        let mut type_ = resolve_root_type(type_path.pop().expect("path must be non-empty"), schema);
        while let Some(path_item) = type_path.pop() {
            type_ = resolve_relative_type(type_, path_item, schema);
        }
        type_
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum TypePathItem {
    Operation(OperationKind),
    FragmentDefinition { type_name: StringKey },
    InlineFragment { type_name: StringKey },
    LinkedField { name: StringKey },
}

pub fn create_completion_request(document: Document, position_span: Span) -> CompletionRequest {
    info!("Building completion path for {:#?}", document);
    let mut completion_request = CompletionRequest::default();

    for definition in document.definitions {
        match &definition {
            ExecutableDefinition::Operation(operation) => {
                if operation.location.contains(position_span) {
                    // TODO don't unwrap here
                    let (_, kind) = operation.operation.clone().unwrap();
                    completion_request.add_type(TypePathItem::Operation(kind));

                    info!(
                        "Completion request is within operation: {:?}",
                        operation.name
                    );
                    let OperationDefinition { selections, .. } = operation;

                    if selections.span.contains(position_span) {
                        // TODO(brandondail) handle when the completion occurs at/within the start token
                        info!("Completion request is within a selection");
                        populate_completion_request_from_selection(
                            selections,
                            position_span,
                            &mut completion_request,
                        );
                    }
                }
                // Check if the position span is within this operation's span
            }
            ExecutableDefinition::Fragment(fragment) => {
                if fragment.location.contains(position_span) {
                    let type_name = fragment.type_condition.type_.value;
                    completion_request.add_type(TypePathItem::FragmentDefinition { type_name });
                    if fragment.selections.span.contains(position_span) {
                        populate_completion_request_from_selection(
                            &fragment.selections,
                            position_span,
                            &mut completion_request,
                        );
                    }
                }
            }
        }
    }

    completion_request
}

/// Resolves the root type of this completion path.
fn resolve_root_type(root_path_item: TypePathItem, schema: &Schema) -> Type {
    match root_path_item {
        TypePathItem::Operation(kind) => match kind {
            OperationKind::Query => schema.query_type().unwrap(),
            OperationKind::Mutation => schema.mutation_type().unwrap(),
            OperationKind::Subscription => schema.subscription_type().unwrap(),
        },
        TypePathItem::FragmentDefinition { type_name } => schema.get_type(type_name).unwrap(),
        TypePathItem::LinkedField { .. } | TypePathItem::InlineFragment { .. } => {
            // TODO(brandondail) fail silently and log here instead
            panic!("Completion paths must start with an operation or fragment")
        }
    }
}

fn resolve_relative_type(parent_type: Type, path_item: TypePathItem, schema: &Schema) -> Type {
    match path_item {
        TypePathItem::Operation(_) => {
            // TODO(brandondail) fail silently and log here instead
            panic!("Operations must only exist at the root of the completion path");
        }
        TypePathItem::FragmentDefinition { .. } => {
            // TODO(brandondail) fail silently and log here instead
            panic!("Fragments must only exist at the root of the completion path");
        }
        TypePathItem::LinkedField { name } => {
            let field_id = schema.named_field(parent_type, name).unwrap();
            let field = schema.field(field_id);
            info!("resolved type for {:?} : {:?}", field.name, field.type_);
            field.type_.inner()
        }
        TypePathItem::InlineFragment { type_name } => schema.get_type(type_name).unwrap(),
    }
}

fn resolve_completion_items_from_fields<T: TypeWithFields>(
    type_: &T,
    schema: &Schema,
) -> Vec<CompletionItem> {
    type_
        .fields()
        .iter()
        .map(|field_id| {
            let field = schema.field(*field_id);
            let name = field.name.to_string();
            CompletionItem::new_simple(name, String::from(""))
        })
        .collect()
}

/// Finds all the valid fragment names for a given type. Used to complete fragment spreads
fn get_valid_fragments_for_type(type_: Type, programs: &Programs<'_>) -> Vec<StringKey> {
    let mut valid_fragment_names = vec![];
    let fragment_map = programs.source.fragment_map();
    for (fragment_name, fragment) in fragment_map {
        if fragment.type_condition == type_ {
            valid_fragment_names.push(*fragment_name);
        }
    }
    info!("get_valid_fragments_for_type {:#?}", valid_fragment_names);
    valid_fragment_names
}

fn resolve_completion_items_for_fragment_spread(
    type_: Type,
    programs: &Programs<'_>,
) -> Vec<CompletionItem> {
    get_valid_fragments_for_type(type_, programs)
        .iter()
        .map(|fragment_name| {
            CompletionItem::new_simple(fragment_name.to_string(), String::from(""))
        })
        .collect()
}

pub fn completion_items_for_request(
    request: CompletionRequest,
    schema: &Schema,
    programs: Option<&Programs<'_>>,
) -> Option<Vec<CompletionItem>> {
    let kind = request.kind;
    let leaf_type = request.resolve_leaf_type(schema);
    info!("completion_items_for_request: {:?} - {:?}", leaf_type, kind);
    match kind {
        CompletionKind::FragmentSpread => {
            if let Some(programs) = programs {
                let items = resolve_completion_items_for_fragment_spread(leaf_type, programs);
                Some(items)
            } else {
                None
            }
        }
        CompletionKind::FieldName => match leaf_type {
            Type::Interface(interface_id) => {
                let interface = schema.interface(interface_id);
                let items = resolve_completion_items_from_fields(interface, schema);
                Some(items)
            }
            Type::Object(object_id) => {
                let object = schema.object(object_id);
                let items = resolve_completion_items_from_fields(object, schema);
                Some(items)
            }
            Type::Enum(_) | Type::InputObject(_) | Type::Scalar(_) | Type::Union(_) => None,
        },
    }
}

fn populate_completion_request_from_selection(
    selections: &List<Selection>,
    position_span: Span,
    completion_request: &mut CompletionRequest,
) {
    for item in &selections.items {
        if item.span().contains(position_span) {
            match item {
                Selection::LinkedField(node) => {
                    completion_request.kind = CompletionKind::FieldName;
                    let LinkedField {
                        name, selections, ..
                    } = node;
                    completion_request.add_type(TypePathItem::LinkedField { name: name.value });
                    populate_completion_request_from_selection(
                        selections,
                        position_span,
                        completion_request,
                    );
                }
                Selection::FragmentSpread(_) => {
                    completion_request.kind = CompletionKind::FragmentSpread;
                }
                Selection::InlineFragment(node) => {
                    let InlineFragment {
                        selections,
                        type_condition,
                        ..
                    } = node;
                    if let Some(type_condition) = type_condition {
                        let type_name = type_condition.type_.value;
                        completion_request.add_type(TypePathItem::InlineFragment { type_name });
                        populate_completion_request_from_selection(
                            selections,
                            position_span,
                            completion_request,
                        )
                    }
                }
                Selection::ScalarField(_node) => {}
            }
        }
    }
}

/// Maps the LSP `Position` type back to a relative span, so we can find out which syntax node(s)
/// this completion request came from
pub fn position_to_span(position: Position, source: &GraphQLSource) -> Option<Span> {
    let mut index_of_last_line = 0;
    let mut line_index = source.line_index as u64;

    let mut chars = source.text.chars().enumerate().peekable();

    while let Some((index, chr)) = chars.next() {
        let is_newline = match chr {
            // Line terminators: https://www.ecma-international.org/ecma-262/#sec-line-terminators
            '\u{000A}' | '\u{000D}' | '\u{2028}' | '\u{2029}' => match (chr, chars.peek()) {
                // <CLRF>
                ('\u{000D}', Some((_, '\u{000D}'))) => false,
                _ => true,
            },
            _ => false,
        };

        if is_newline {
            line_index += 1;
            index_of_last_line = index as u64;
        }

        if line_index == position.line {
            let start_offset = index_of_last_line + position.character;
            return Some(Span::new(start_offset as u32, 0));
        }
    }
    None
}

pub fn send_completion_response(
    items: Vec<CompletionItem>,
    request_id: ServerRequestId,
    connection: &Connection,
) {
    info!("send_completion_response: {:#?}", items);
    // If there are no items, don't send any response
    if items.is_empty() {
        return;
    }
    let completion_response = CompletionResponse::Array(items);
    let result = serde_json::to_value(&completion_response).unwrap();
    let response = ServerResponse {
        id: request_id,
        error: None,
        result: Some(result),
    };
    connection.sender.send(Message::Response(response)).unwrap();
}

/// Return a `CompletionPath` for this request, only if the completion request occurs
// within a GraphQL document. Otherwise return `None`
pub fn get_completion_request(
    params: CompletionParams,
    graphql_source_cache: &GraphQLSourceCache,
) -> Option<CompletionRequest> {
    let CompletionParams {
        text_document_position,
        ..
    } = params;
    let TextDocumentPositionParams {
        text_document,
        position,
    } = text_document_position;
    let url = text_document.uri;
    let graphql_sources = match graphql_source_cache.get(&url) {
        Some(sources) => sources,
        // If we have no sources for this file, do nothing
        None => return None,
    };

    info!(
        "Got completion request for file with sources: {:#?}",
        *graphql_sources
    );

    info!("position: {:?}", position);

    // We have GraphQL documents, now check if the completion request
    // falls within the range of one of these documents.
    let mut target_graphql_source: Option<&GraphQLSource> = None;
    for graphql_source in &*graphql_sources {
        let range = graphql_source.to_range();
        if position >= range.start && position <= range.end {
            target_graphql_source = Some(graphql_source);
            break;
        }
    }

    let graphql_source = match target_graphql_source {
        Some(source) => source,
        // Exit early if this completion request didn't fall within
        // the range of one of our GraphQL documents
        None => return None,
    };

    match parse(&graphql_source.text, FileKey::new(&url.to_string())) {
        Ok(document) => {
            // Now we need to take the `Position` and map that to an offset relative
            // to this GraphQL document, as the `Span`s in the document are relative.
            info!("Successfully parsed the definitions for a target GraphQL source");
            // Map the position to a zero-length span, relative to this GraphQL source.
            let position_span = match position_to_span(position, &graphql_source) {
                Some(span) => span,
                // Exit early if we can't map the position for some reason
                None => return None,
            };
            // Now we need to walk the Document, tracking our path along the way, until
            // we find the position within the document. Note that the GraphQLSource will
            // already be updated *with the characters that triggered the completion request*
            // since the change event fires before completion.
            info!("position_span: {:?}", position_span);
            let completion_request = create_completion_request(document, position_span);
            info!("Completion path: {:#?}", completion_request);
            Some(completion_request)
        }
        Err(err) => {
            info!("Failed to parse this target!");
            info!("{:?}", err);
            None
        }
    }
}
<commit_msg>add snippet completion support for directives on query<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

//! Utilities for providing the completion language feature
use crate::lsp::Position;
use common::{FileKey, Span};
use graphql_syntax::{parse, Document, GraphQLSource};
use interner::StringKey;
use log::info;

use relay_compiler::Programs;

use crate::lsp::{
    CompletionItem, CompletionParams, CompletionResponse, Connection, Message, ServerRequestId,
    ServerResponse, TextDocumentPositionParams, Url,
};
use schema::{Directive as SchemaDirective, DirectiveLocation, Schema, Type, TypeWithFields};

use graphql_syntax::{
    Directive, ExecutableDefinition, InlineFragment, LinkedField, List, OperationDefinition,
    OperationKind, Selection,
};

pub type GraphQLSourceCache = std::collections::HashMap<Url, Vec<GraphQLSource>>;

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
pub enum CompletionKind {
    FieldName,
    FragmentSpread,
    DirectiveName,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct CompletionRequest {
    /// The type of the completion request we're responding to
    kind: CompletionKind,
    /// A list of type metadata that we can use to resolve the leaf
    /// type the request is being made against
    type_path: Vec<TypePathItem>,
}

impl Default for CompletionRequest {
    fn default() -> Self {
        CompletionRequest {
            kind: CompletionKind::FieldName,
            type_path: vec![],
        }
    }
}

impl CompletionRequest {
    fn add_type(&mut self, type_path_item: TypePathItem) {
        self.type_path.push(type_path_item)
    }

    /// Returns the leaf type, which is the type that the completion request is being made against.
    fn resolve_leaf_type(self, schema: &Schema) -> Type {
        let mut type_path = self.type_path;
        type_path.reverse();
        let mut type_ = resolve_root_type(type_path.pop().expect("path must be non-empty"), schema);
        while let Some(path_item) = type_path.pop() {
            type_ = resolve_relative_type(type_, path_item, schema);
        }
        type_
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum TypePathItem {
    Operation(OperationKind),
    FragmentDefinition { type_name: StringKey },
    InlineFragment { type_name: StringKey },
    LinkedField { name: StringKey },
}

pub fn create_completion_request(document: Document, position_span: Span) -> CompletionRequest {
    info!("Building completion path for {:#?}", document);
    let mut completion_request = CompletionRequest::default();

    for definition in document.definitions {
        match &definition {
            ExecutableDefinition::Operation(operation) => {
                if operation.location.contains(position_span) {
                    // TODO don't unwrap here
                    let (_, kind) = operation.operation.clone().unwrap();
                    completion_request.add_type(TypePathItem::Operation(kind));

                    info!(
                        "Completion request is within operation: {:?}",
                        operation.name
                    );
                    let OperationDefinition {
                        selections,
                        directives,
                        ..
                    } = operation;

                    for Directive { span, name, .. } in directives {
                        if span.contains(position_span) {
                            if name.span.contains(position_span) {
                                completion_request.kind = CompletionKind::DirectiveName;
                                return completion_request;
                            }
                        }
                    }

                    if selections.span.contains(position_span) {
                        // TODO(brandondail) handle when the completion occurs at/within the start token
                        info!("Completion request is within a selection");
                        populate_completion_request_from_selection(
                            selections,
                            position_span,
                            &mut completion_request,
                        );
                    }
                }
                // Check if the position span is within this operation's span
            }
            ExecutableDefinition::Fragment(fragment) => {
                if fragment.location.contains(position_span) {
                    let type_name = fragment.type_condition.type_.value;
                    completion_request.add_type(TypePathItem::FragmentDefinition { type_name });
                    if fragment.selections.span.contains(position_span) {
                        populate_completion_request_from_selection(
                            &fragment.selections,
                            position_span,
                            &mut completion_request,
                        );
                    }
                }
            }
        }
    }

    completion_request
}

/// Resolves the root type of this completion path.
fn resolve_root_type(root_path_item: TypePathItem, schema: &Schema) -> Type {
    match root_path_item {
        TypePathItem::Operation(kind) => match kind {
            OperationKind::Query => schema.query_type().unwrap(),
            OperationKind::Mutation => schema.mutation_type().unwrap(),
            OperationKind::Subscription => schema.subscription_type().unwrap(),
        },
        TypePathItem::FragmentDefinition { type_name } => schema.get_type(type_name).unwrap(),
        TypePathItem::LinkedField { .. } | TypePathItem::InlineFragment { .. } => {
            // TODO(brandondail) fail silently and log here instead
            panic!("Completion paths must start with an operation or fragment")
        }
    }
}

fn resolve_relative_type(parent_type: Type, path_item: TypePathItem, schema: &Schema) -> Type {
    match path_item {
        TypePathItem::Operation(_) => {
            // TODO(brandondail) fail silently and log here instead
            panic!("Operations must only exist at the root of the completion path");
        }
        TypePathItem::FragmentDefinition { .. } => {
            // TODO(brandondail) fail silently and log here instead
            panic!("Fragments must only exist at the root of the completion path");
        }
        TypePathItem::LinkedField { name } => {
            let field_id = schema.named_field(parent_type, name).unwrap();
            let field = schema.field(field_id);
            info!("resolved type for {:?} : {:?}", field.name, field.type_);
            field.type_.inner()
        }
        TypePathItem::InlineFragment { type_name } => schema.get_type(type_name).unwrap(),
    }
}

fn resolve_completion_items_from_fields<T: TypeWithFields>(
    type_: &T,
    schema: &Schema,
) -> Vec<CompletionItem> {
    type_
        .fields()
        .iter()
        .map(|field_id| {
            let field = schema.field(*field_id);
            let name = field.name.to_string();
            CompletionItem::new_simple(name, String::from(""))
        })
        .collect()
}

/// Finds all the valid fragment names for a given type. Used to complete fragment spreads
fn get_valid_fragments_for_type(type_: Type, programs: &Programs<'_>) -> Vec<StringKey> {
    let mut valid_fragment_names = vec![];
    let fragment_map = programs.source.fragment_map();
    for (fragment_name, fragment) in fragment_map {
        if fragment.type_condition == type_ {
            valid_fragment_names.push(*fragment_name);
        }
    }
    info!("get_valid_fragments_for_type {:#?}", valid_fragment_names);
    valid_fragment_names
}

fn resolve_completion_items_for_fragment_spread(
    type_: Type,
    programs: &Programs<'_>,
) -> Vec<CompletionItem> {
    get_valid_fragments_for_type(type_, programs)
        .iter()
        .map(|fragment_name| {
            CompletionItem::new_simple(fragment_name.to_string(), String::from(""))
        })
        .collect()
}

pub fn completion_items_for_request(
    request: CompletionRequest,
    schema: &Schema,
    programs: Option<&Programs<'_>>,
) -> Option<Vec<CompletionItem>> {
    let kind = request.kind;
    let leaf_type = request.resolve_leaf_type(schema);
    info!("completion_items_for_request: {:?} - {:?}", leaf_type, kind);
    match kind {
        CompletionKind::FragmentSpread => {
            if let Some(programs) = programs {
                let items = resolve_completion_items_for_fragment_spread(leaf_type, programs);
                Some(items)
            } else {
                None
            }
        }
        CompletionKind::FieldName => match leaf_type {
            Type::Interface(interface_id) => {
                let interface = schema.interface(interface_id);
                let items = resolve_completion_items_from_fields(interface, schema);
                Some(items)
            }
            Type::Object(object_id) => {
                let object = schema.object(object_id);
                let items = resolve_completion_items_from_fields(object, schema);
                Some(items)
            }
            Type::Enum(_) | Type::InputObject(_) | Type::Scalar(_) | Type::Union(_) => None,
        },
        CompletionKind::DirectiveName => match leaf_type {
            Type::Object(_) => {
                let items = schema
                    .get_directives()
                    // TODO don't hardcode the query type assumption
                    .filter(|directive| directive.locations.contains(&DirectiveLocation::Query))
                    .map(|directive| {
                        completion_item_from_directive(directive)
                    })
                    .collect();
                Some(items)
            }
            _ => None,
        },
    }
}

fn populate_completion_request_from_selection(
    selections: &List<Selection>,
    position_span: Span,
    completion_request: &mut CompletionRequest,
) {
    for item in &selections.items {
        if item.span().contains(position_span) {
            match item {
                Selection::LinkedField(node) => {
                    completion_request.kind = CompletionKind::FieldName;
                    let LinkedField {
                        name, selections, ..
                    } = node;
                    completion_request.add_type(TypePathItem::LinkedField { name: name.value });
                    populate_completion_request_from_selection(
                        selections,
                        position_span,
                        completion_request,
                    );
                }
                Selection::FragmentSpread(_) => {
                    completion_request.kind = CompletionKind::FragmentSpread;
                }
                Selection::InlineFragment(node) => {
                    let InlineFragment {
                        selections,
                        type_condition,
                        ..
                    } = node;
                    if let Some(type_condition) = type_condition {
                        let type_name = type_condition.type_.value;
                        completion_request.add_type(TypePathItem::InlineFragment { type_name });
                        populate_completion_request_from_selection(
                            selections,
                            position_span,
                            completion_request,
                        )
                    }
                }
                Selection::ScalarField(_node) => {}
            }
        }
    }
}

fn completion_item_from_directive(directive: &SchemaDirective) -> CompletionItem {
    let SchemaDirective {
        name, arguments, ..
    } = directive;

    use crate::lsp::InsertTextFormat;

    // Always use the name of the directive as the label
    let label = name.to_string();

    // We can return a snippet with the expected arguments of the directive
    let (insert_text, insert_text_format) = if arguments.is_empty() {
        (label.clone(), InsertTextFormat::PlainText)
    } else {
        let mut insert_text = String::from(&label);
        let mut cursor_location = 1;

        insert_text.push('(');

        // Iterate through the arguments and build up a snippet
        for (i, argument) in arguments.iter().enumerate() {
            if i != 0 {
                insert_text.push_str(", ");
            }
            insert_text.push_str(&format!("{}: ${}", argument.name, cursor_location));
            cursor_location += 1;
        }

        insert_text.push(')');
        (insert_text, InsertTextFormat::Snippet)
    };

    CompletionItem {
        label,
        kind: None,
        detail: None,
        documentation: None,
        deprecated: None,
        preselect: None,
        sort_text: None,
        filter_text: None,
        insert_text: Some(insert_text),
        insert_text_format: Some(insert_text_format),
        text_edit: None,
        additional_text_edits: None,
        command: None,
        data: None,
        tags: None,
    }
}

/// Maps the LSP `Position` type back to a relative span, so we can find out which syntax node(s)
/// this completion request came from
pub fn position_to_span(position: Position, source: &GraphQLSource) -> Option<Span> {
    let mut index_of_last_line = 0;
    let mut line_index = source.line_index as u64;

    let mut chars = source.text.chars().enumerate().peekable();

    while let Some((index, chr)) = chars.next() {
        let is_newline = match chr {
            // Line terminators: https://www.ecma-international.org/ecma-262/#sec-line-terminators
            '\u{000A}' | '\u{000D}' | '\u{2028}' | '\u{2029}' => match (chr, chars.peek()) {
                // <CLRF>
                ('\u{000D}', Some((_, '\u{000D}'))) => false,
                _ => true,
            },
            _ => false,
        };

        if is_newline {
            line_index += 1;
            index_of_last_line = index as u64;
        }

        if line_index == position.line {
            let start_offset = index_of_last_line + position.character;
            return Some(Span::new(start_offset as u32, 0));
        }
    }
    None
}

pub fn send_completion_response(
    items: Vec<CompletionItem>,
    request_id: ServerRequestId,
    connection: &Connection,
) {
    // If there are no items, don't send any response
    if items.is_empty() {
        return;
    }
    let completion_response = CompletionResponse::Array(items);
    let result = serde_json::to_value(&completion_response).unwrap();
    let response = ServerResponse {
        id: request_id,
        error: None,
        result: Some(result),
    };
    connection.sender.send(Message::Response(response)).unwrap();
}

/// Return a `CompletionPath` for this request, only if the completion request occurs
// within a GraphQL document. Otherwise return `None`
pub fn get_completion_request(
    params: CompletionParams,
    graphql_source_cache: &GraphQLSourceCache,
) -> Option<CompletionRequest> {
    let CompletionParams {
        text_document_position,
        ..
    } = params;
    let TextDocumentPositionParams {
        text_document,
        position,
    } = text_document_position;
    let url = text_document.uri;
    let graphql_sources = match graphql_source_cache.get(&url) {
        Some(sources) => sources,
        // If we have no sources for this file, do nothing
        None => return None,
    };

    info!(
        "Got completion request for file with sources: {:#?}",
        *graphql_sources
    );

    info!("position: {:?}", position);

    // We have GraphQL documents, now check if the completion request
    // falls within the range of one of these documents.
    let mut target_graphql_source: Option<&GraphQLSource> = None;
    for graphql_source in &*graphql_sources {
        let range = graphql_source.to_range();
        if position >= range.start && position <= range.end {
            target_graphql_source = Some(graphql_source);
            break;
        }
    }

    let graphql_source = match target_graphql_source {
        Some(source) => source,
        // Exit early if this completion request didn't fall within
        // the range of one of our GraphQL documents
        None => return None,
    };

    match parse(&graphql_source.text, FileKey::new(&url.to_string())) {
        Ok(document) => {
            // Now we need to take the `Position` and map that to an offset relative
            // to this GraphQL document, as the `Span`s in the document are relative.
            info!("Successfully parsed the definitions for a target GraphQL source");
            // Map the position to a zero-length span, relative to this GraphQL source.
            let position_span = match position_to_span(position, &graphql_source) {
                Some(span) => span,
                // Exit early if we can't map the position for some reason
                None => return None,
            };
            // Now we need to walk the Document, tracking our path along the way, until
            // we find the position within the document. Note that the GraphQLSource will
            // already be updated *with the characters that triggered the completion request*
            // since the change event fires before completion.
            info!("position_span: {:?}", position_span);
            let completion_request = create_completion_request(document, position_span);
            info!("Completion path: {:#?}", completion_request);
            Some(completion_request)
        }
        Err(err) => {
            info!("Failed to parse this target!");
            info!("{:?}", err);
            None
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    code_action::on_code_action,
    completion::on_completion,
    goto_definition::on_goto_definition,
    hover::on_hover,
    lsp::{
        set_actual_server_status, set_starting_status, CompletionOptions, Connection,
        GotoDefinition, HoverRequest, InitializeParams, Message, ServerCapabilities,
        ServerResponse, TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions,
    },
    lsp_process_error::{LSPProcessError, LSPProcessResult},
    references::on_references,
    shutdown::{on_exit, on_shutdown},
    status_reporting::LSPStatusReporter,
    text_documents::on_did_save_text_document,
    text_documents::{
        on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
    },
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::{SendError, Sender};
use log::info;
use lsp_server::{ErrorCode, Notification, ResponseError};
use lsp_types::{
    notification::{
        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Exit,
    },
    request::{CodeActionRequest, Completion, References, Shutdown},
    CodeActionProviderCapability,
};
use relay_compiler::config::Config;
use std::sync::{Arc, RwLock};
mod lsp_request_dispatch;
use lsp_request_dispatch::LSPRequestDispatch;
mod lsp_notification_dispatch;
use lsp_notification_dispatch::LSPNotificationDispatch;
mod lsp_state;
pub use lsp_state::LSPExtraDataProvider;
pub(crate) use lsp_state::{LSPState, LSPStateError};

/// Initializes an LSP connection, handling the `initialize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> LSPProcessResult<InitializeParams> {
    let mut server_capabilities = ServerCapabilities::default();
    // Enable text document syncing so we can know when files are opened/changed/saved/closed
    server_capabilities.text_document_sync =
        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));

    server_capabilities.completion_provider = Some(CompletionOptions {
        resolve_provider: Some(true),
        trigger_characters: Some(vec![
            "(".into(),
            "{".into(),
            "\n".into(),
            ",".into(),
            "@".into(),
        ]),
        work_done_progress_options: WorkDoneProgressOptions {
            work_done_progress: None,
        },
    });

    server_capabilities.hover_provider = Some(true);
    server_capabilities.definition_provider = Some(true);
    server_capabilities.references_provider = Some(true);
    server_capabilities.code_action_provider = Some(CodeActionProviderCapability::Simple(true));

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub async fn run<TPerfLogger: PerfLogger + 'static>(
    connection: Connection,
    mut config: Config,
    _params: InitializeParams,
    perf_logger: Arc<TPerfLogger>,
    extra_data_provider: Box<dyn LSPExtraDataProvider + Send + Sync>,
) -> LSPProcessResult<()>
where
    TPerfLogger: PerfLogger + 'static,
{
    info!(
        "Running language server with config root {:?}",
        config.root_dir
    );
    set_starting_status(&connection.sender);

    let lsp_state_errors: Arc<RwLock<Vec<LSPStateError>>> = Default::default();

    config.status_reporter = Box::new(LSPStatusReporter::new(
        config.root_dir.clone(),
        connection.sender.clone(),
        Arc::clone(&lsp_state_errors),
    ));

    let mut lsp_state = LSPState::create_state(
        Arc::new(config),
        lsp_state_errors,
        extra_data_provider,
        Arc::clone(&perf_logger),
    )
    .await?;

    set_actual_server_status(&connection.sender, &lsp_state.get_errors());

    for msg in connection.receiver {
        info!("LSP message received {:?}", msg);
        match msg {
            Message::Request(req) => {
                handle_request(&mut lsp_state, req, &connection.sender, &perf_logger)
                    .map_err(LSPProcessError::from)?;
            }
            Message::Notification(notification) => {
                handle_notification(&mut lsp_state, notification, &perf_logger);
            }
            _ => {
                // Ignore responses for now
            }
        }
    }

    Ok(())
}

fn handle_request<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    request: lsp_server::Request,
    sender: &Sender<Message>,
    perf_logger: &Arc<TPerfLogger>,
) -> Result<(), SendError<Message>> {
    let get_server_response_bound = |req| dispatch_request(req, lsp_state);
    let get_response = with_request_logging(perf_logger, get_server_response_bound);

    sender.send(Message::Response(get_response(request)))
}

fn dispatch_request<TPerfLogger: PerfLogger + 'static>(
    request: lsp_server::Request,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> ServerResponse {
    let get_response = || -> Result<_, ServerResponse> {
        let request = LSPRequestDispatch::new(request, lsp_state)
            .on_request_sync::<HoverRequest>(on_hover)?
            .on_request_sync::<GotoDefinition>(on_goto_definition)?
            .on_request_sync::<References>(on_references)?
            .on_request_sync::<Completion>(on_completion)?
            .on_request_sync::<CodeActionRequest>(on_code_action)?
            .on_request_sync::<Shutdown>(on_shutdown)?
            .request();

        // If we have gotten here, we have not handled the request
        Ok(ServerResponse {
            id: request.id,
            result: None,
            error: Some(ResponseError {
                code: ErrorCode::MethodNotFound as i32,
                data: None,
                message: format!("No handler registered for method '{}'", request.method),
            }),
        })
    };
    get_response().unwrap_or_else(|response| response)
}

fn with_request_logging<'a, TPerfLogger: PerfLogger + 'static>(
    perf_logger: &'a Arc<TPerfLogger>,
    get_response: impl FnOnce(lsp_server::Request) -> ServerResponse + 'a,
) -> impl FnOnce(lsp_server::Request) -> ServerResponse + 'a {
    move |request| {
        let lsp_request_event = perf_logger.create_event("lsp_message");
        lsp_request_event.string("lsp_method", request.method.clone());
        lsp_request_event.string("lsp_type", "request".to_string());
        let lsp_request_processing_time = lsp_request_event.start("lsp_message_processing_time");

        let response = get_response(request);

        if response.result.is_some() {
            lsp_request_event.string("lsp_outcome", "success".to_string());
        } else if let Some(error) = &response.error {
            lsp_request_event.string("lsp_outcome", "error".to_string());
            if let Some(data) = &error.data {
                lsp_request_event.string("lsp_error_data", data.to_string());
            }
        }
        // N.B. we don't handle the case where the ServerResponse has neither a result nor
        // an error, which is an invalid state.

        lsp_request_event.stop(lsp_request_processing_time);
        perf_logger.complete_event(lsp_request_event);

        response
    }
}

fn handle_notification<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    notification: Notification,
    perf_logger: &Arc<TPerfLogger>,
) {
    let lsp_notification_event = perf_logger.create_event("lsp_message");
    lsp_notification_event.string("lsp_method", notification.method.clone());
    lsp_notification_event.string("lsp_type", "notification".to_string());
    let lsp_notification_processing_time =
        lsp_notification_event.start("lsp_message_processing_time");

    let notification_result = dispatch_notification(notification, lsp_state);

    // The only possible error (for now) is not handling the notification.
    // N.B. is_ok is correct here.
    if notification_result.is_ok() {
        lsp_notification_event.string("lsp_outcome", "error".to_string());
    }

    lsp_notification_event.stop(lsp_notification_processing_time);
    perf_logger.complete_event(lsp_notification_event);
}

fn dispatch_notification<TPerfLogger: PerfLogger + 'static>(
    notification: lsp_server::Notification,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> Result<(), ()> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

    // If we have gotten here, we have not handled the notification
    info!(
        "Error: no handler registered for notification '{}'",
        notification.method
    );
    Ok(())
}
<commit_msg>Remove "{" from auto complete trigger characters<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::{
    code_action::on_code_action,
    completion::on_completion,
    goto_definition::on_goto_definition,
    hover::on_hover,
    lsp::{
        set_actual_server_status, set_starting_status, CompletionOptions, Connection,
        GotoDefinition, HoverRequest, InitializeParams, Message, ServerCapabilities,
        ServerResponse, TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions,
    },
    lsp_process_error::{LSPProcessError, LSPProcessResult},
    references::on_references,
    shutdown::{on_exit, on_shutdown},
    status_reporting::LSPStatusReporter,
    text_documents::on_did_save_text_document,
    text_documents::{
        on_did_change_text_document, on_did_close_text_document, on_did_open_text_document,
    },
};
use common::{PerfLogEvent, PerfLogger};
use crossbeam::{SendError, Sender};
use log::info;
use lsp_server::{ErrorCode, Notification, ResponseError};
use lsp_types::{
    notification::{
        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, DidSaveTextDocument, Exit,
    },
    request::{CodeActionRequest, Completion, References, Shutdown},
    CodeActionProviderCapability,
};
use relay_compiler::config::Config;
use std::sync::{Arc, RwLock};
mod lsp_request_dispatch;
use lsp_request_dispatch::LSPRequestDispatch;
mod lsp_notification_dispatch;
use lsp_notification_dispatch::LSPNotificationDispatch;
mod lsp_state;
pub use lsp_state::LSPExtraDataProvider;
pub(crate) use lsp_state::{LSPState, LSPStateError};

/// Initializes an LSP connection, handling the `initialize` message and `initialized` notification
/// handshake.
pub fn initialize(connection: &Connection) -> LSPProcessResult<InitializeParams> {
    let mut server_capabilities = ServerCapabilities::default();
    // Enable text document syncing so we can know when files are opened/changed/saved/closed
    server_capabilities.text_document_sync =
        Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::Full));

    server_capabilities.completion_provider = Some(CompletionOptions {
        resolve_provider: Some(true),
        trigger_characters: Some(vec!["(".into(), "\n".into(), ",".into(), "@".into()]),
        work_done_progress_options: WorkDoneProgressOptions {
            work_done_progress: None,
        },
    });

    server_capabilities.hover_provider = Some(true);
    server_capabilities.definition_provider = Some(true);
    server_capabilities.references_provider = Some(true);
    server_capabilities.code_action_provider = Some(CodeActionProviderCapability::Simple(true));

    let server_capabilities = serde_json::to_value(&server_capabilities)?;
    let params = connection.initialize(server_capabilities)?;
    let params: InitializeParams = serde_json::from_value(params)?;
    Ok(params)
}

/// Run the main server loop
pub async fn run<TPerfLogger: PerfLogger + 'static>(
    connection: Connection,
    mut config: Config,
    _params: InitializeParams,
    perf_logger: Arc<TPerfLogger>,
    extra_data_provider: Box<dyn LSPExtraDataProvider + Send + Sync>,
) -> LSPProcessResult<()>
where
    TPerfLogger: PerfLogger + 'static,
{
    info!(
        "Running language server with config root {:?}",
        config.root_dir
    );
    set_starting_status(&connection.sender);

    let lsp_state_errors: Arc<RwLock<Vec<LSPStateError>>> = Default::default();

    config.status_reporter = Box::new(LSPStatusReporter::new(
        config.root_dir.clone(),
        connection.sender.clone(),
        Arc::clone(&lsp_state_errors),
    ));

    let mut lsp_state = LSPState::create_state(
        Arc::new(config),
        lsp_state_errors,
        extra_data_provider,
        Arc::clone(&perf_logger),
    )
    .await?;

    set_actual_server_status(&connection.sender, &lsp_state.get_errors());

    for msg in connection.receiver {
        info!("LSP message received {:?}", msg);
        match msg {
            Message::Request(req) => {
                handle_request(&mut lsp_state, req, &connection.sender, &perf_logger)
                    .map_err(LSPProcessError::from)?;
            }
            Message::Notification(notification) => {
                handle_notification(&mut lsp_state, notification, &perf_logger);
            }
            _ => {
                // Ignore responses for now
            }
        }
    }

    Ok(())
}

fn handle_request<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    request: lsp_server::Request,
    sender: &Sender<Message>,
    perf_logger: &Arc<TPerfLogger>,
) -> Result<(), SendError<Message>> {
    let get_server_response_bound = |req| dispatch_request(req, lsp_state);
    let get_response = with_request_logging(perf_logger, get_server_response_bound);

    sender.send(Message::Response(get_response(request)))
}

fn dispatch_request<TPerfLogger: PerfLogger + 'static>(
    request: lsp_server::Request,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> ServerResponse {
    let get_response = || -> Result<_, ServerResponse> {
        let request = LSPRequestDispatch::new(request, lsp_state)
            .on_request_sync::<HoverRequest>(on_hover)?
            .on_request_sync::<GotoDefinition>(on_goto_definition)?
            .on_request_sync::<References>(on_references)?
            .on_request_sync::<Completion>(on_completion)?
            .on_request_sync::<CodeActionRequest>(on_code_action)?
            .on_request_sync::<Shutdown>(on_shutdown)?
            .request();

        // If we have gotten here, we have not handled the request
        Ok(ServerResponse {
            id: request.id,
            result: None,
            error: Some(ResponseError {
                code: ErrorCode::MethodNotFound as i32,
                data: None,
                message: format!("No handler registered for method '{}'", request.method),
            }),
        })
    };
    get_response().unwrap_or_else(|response| response)
}

fn with_request_logging<'a, TPerfLogger: PerfLogger + 'static>(
    perf_logger: &'a Arc<TPerfLogger>,
    get_response: impl FnOnce(lsp_server::Request) -> ServerResponse + 'a,
) -> impl FnOnce(lsp_server::Request) -> ServerResponse + 'a {
    move |request| {
        let lsp_request_event = perf_logger.create_event("lsp_message");
        lsp_request_event.string("lsp_method", request.method.clone());
        lsp_request_event.string("lsp_type", "request".to_string());
        let lsp_request_processing_time = lsp_request_event.start("lsp_message_processing_time");

        let response = get_response(request);

        if response.result.is_some() {
            lsp_request_event.string("lsp_outcome", "success".to_string());
        } else if let Some(error) = &response.error {
            lsp_request_event.string("lsp_outcome", "error".to_string());
            if let Some(data) = &error.data {
                lsp_request_event.string("lsp_error_data", data.to_string());
            }
        }
        // N.B. we don't handle the case where the ServerResponse has neither a result nor
        // an error, which is an invalid state.

        lsp_request_event.stop(lsp_request_processing_time);
        perf_logger.complete_event(lsp_request_event);

        response
    }
}

fn handle_notification<TPerfLogger: PerfLogger + 'static>(
    lsp_state: &mut LSPState<TPerfLogger>,
    notification: Notification,
    perf_logger: &Arc<TPerfLogger>,
) {
    let lsp_notification_event = perf_logger.create_event("lsp_message");
    lsp_notification_event.string("lsp_method", notification.method.clone());
    lsp_notification_event.string("lsp_type", "notification".to_string());
    let lsp_notification_processing_time =
        lsp_notification_event.start("lsp_message_processing_time");

    let notification_result = dispatch_notification(notification, lsp_state);

    // The only possible error (for now) is not handling the notification.
    // N.B. is_ok is correct here.
    if notification_result.is_ok() {
        lsp_notification_event.string("lsp_outcome", "error".to_string());
    }

    lsp_notification_event.stop(lsp_notification_processing_time);
    perf_logger.complete_event(lsp_notification_event);
}

fn dispatch_notification<TPerfLogger: PerfLogger + 'static>(
    notification: lsp_server::Notification,
    lsp_state: &mut LSPState<TPerfLogger>,
) -> Result<(), ()> {
    let notification = LSPNotificationDispatch::new(notification, lsp_state)
        .on_notification_sync::<DidOpenTextDocument>(on_did_open_text_document)?
        .on_notification_sync::<DidCloseTextDocument>(on_did_close_text_document)?
        .on_notification_sync::<DidChangeTextDocument>(on_did_change_text_document)?
        .on_notification_sync::<DidSaveTextDocument>(on_did_save_text_document)?
        .on_notification_sync::<Exit>(on_exit)?
        .notification();

    // If we have gotten here, we have not handled the notification
    info!(
        "Error: no handler registered for notification '{}'",
        notification.method
    );
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 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::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{CACHE_ARTIFACT_PATH, CACHE_KEY_PATH, find_command};
use hcore::os::process;
use command::studio::enter::ARTIFACT_PATH_ENVVAR;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &'static str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;

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

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            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 !is_serving_windows_containers(&docker_cmd) &&
        (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("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        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);
    }

    check_mounts(&docker_cmd, volumes.iter())?;
    run_container(docker_cmd, args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images").arg(&image_identifier(docker_cmd)).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) -> Result<()> {
    let image = image_identifier(docker_cmd);
    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_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).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) -> 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());
    }

    if !is_serving_windows_containers(docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    cmd_args.push(image_identifier(docker_cmd).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: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> 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 !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            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());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd, cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// 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(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<commit_msg>do not check mounts entering windows container studio<commit_after>// Copyright (c) 2016-2017 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::env;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};

use common::ui::UI;
use hcore::crypto::default_cache_key_path;
use hcore::env as henv;
use hcore::fs::{CACHE_ARTIFACT_PATH, CACHE_KEY_PATH, find_command};
use hcore::os::process;
use command::studio::enter::ARTIFACT_PATH_ENVVAR;

use error::{Error, Result};
use VERSION;

const DOCKER_CMD: &'static str = "docker";
const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
const DOCKER_WINDOWS_IMAGE: &'static str = "habitat-docker-registry.bintray.io/win-studio";
const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
const DOCKER_OPTS_ENVVAR: &'static str = "HAB_DOCKER_OPTS";
const DOCKER_SOCKET: &'static str = "/var/run/docker.sock";
const HAB_STUDIO_SECRET: &'static str = "HAB_STUDIO_SECRET_";

pub fn start_docker_studio(_ui: &mut UI, mut args: Vec<OsString>) -> Result<()> {
    let docker_cmd = find_docker_cmd()?;

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

    let mnt_prefix = match is_serving_windows_containers(&docker_cmd) {
        true => "c:",
        false => "",
    };
    let mut volumes = vec![
        format!(
            "{}:{}{}",
            env::current_dir().unwrap().to_string_lossy(),
            mnt_prefix,
            "/src"
        ),
        format!(
            "{}:{}/{}",
            default_cache_key_path(None).to_string_lossy(),
            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 !is_serving_windows_containers(&docker_cmd) &&
        (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("HAB_AUTH_TOKEN"),
        String::from("HAB_BLDR_URL"),
        String::from("HAB_BLDR_CHANNEL"),
        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 !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    run_container(docker_cmd, args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    cmd.arg("images").arg(&image_identifier(docker_cmd)).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) -> Result<()> {
    let image = image_identifier(docker_cmd);
    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_identifier(docker_cmd).to_string(),
            ));
        } else if err_output.contains("Cannot connect to the Docker daemon") {
            return Err(Error::DockerDaemonDown);
        } else {
            return Err(Error::DockerNetworkDown(
                image_identifier(docker_cmd).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) -> 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("--privileged".into());
    cmd_args.push(image_identifier(docker_cmd).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: Vec<OsString>,
    volumes: I,
    env_vars: J,
) -> 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 !is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("--privileged".into());
    }
    match args.first().map(|f| f.to_str().unwrap_or_default()) {
        Some("build") => {}
        _ => {
            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());
        }
    }
    for vol in volumes {
        cmd_args.push("--volume".into());
        cmd_args.push(vol.as_ref().into());
    }
    cmd_args.push(image_identifier(&docker_cmd).into());
    cmd_args.extend_from_slice(args.as_slice());
    if is_serving_windows_containers(&docker_cmd) {
        cmd_args.push("-w".into());
        cmd_args.push("-n".into());
        cmd_args.push("-o".into());
        cmd_args.push("c:/".into());
    }
    unset_proxy_env_vars();
    Ok(process::become_command(docker_cmd, cmd_args)?)
}

fn unset_proxy_env_vars() {
    for var in vec!["http_proxy", "https_proxy"] {
        if let Ok(_) = henv::var(var) {
            debug!("Unsetting process environment variable '{}'", var);
            env::remove_var(var);
        }
    }
}

/// 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(docker_cmd: &Path) -> String {
    let version: Vec<&str> = VERSION.split("/").collect();
    let img = match is_serving_windows_containers(docker_cmd) {
        true => DOCKER_WINDOWS_IMAGE,
        false => DOCKER_IMAGE,
    };
    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", img, version[0]))
}

#[cfg(test)]
#[cfg(target_os = "linux")]
mod tests {
    use super::{find_docker_cmd, image_identifier, DOCKER_IMAGE};
    use VERSION;

    #[test]
    fn retrieve_image_identifier() {
        let docker_cmd = find_docker_cmd().unwrap();
        assert_eq!(
            image_identifier(&docker_cmd),
            format!("{}:{}", DOCKER_IMAGE, VERSION)
        );
    }
}
<|endoftext|>


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

const DOCKER_CMD: &str = "docker";
const DOCKER_CMD_ENVVAR: &str = "HAB_DOCKER_BINARY";

const DOCKER_IMAGE: &str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &str = "habitat-docker-registry.bintray.io/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 = find_docker_cmd()?;

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

    let mnt_prefix = if is_serving_windows_containers(&docker_cmd) {
        "c:"
    } else {
        ""
    };

    let local_cache_key_path = default_cache_key_path(None);
    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.to_string_lossy().into_owned()
            )
        ));
    }

    let mut volumes = vec![format!("{}:{}{}",
                                   env::current_dir().unwrap().to_string_lossy(),
                                   mnt_prefix,
                                   "/src"),
                           format!("{}:{}/{}",
                                   local_cache_key_path.to_string_lossy(),
                                   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 !is_serving_windows_containers(&docker_cmd)
       && (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 !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    run_container(docker_cmd, &args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or_else(|_| DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    let image = image_identifier_for_active_target(&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) -> Result<()> {
    let image = image_identifier_for_active_target(&docker_cmd);
    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) -> Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<OsStr>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    let image = image_identifier_for_active_target(&docker_cmd);

    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)
                             -> Result<()>
    where I: IntoIterator<Item = S>,
          J: IntoIterator<Item = T>,
          S: AsRef<OsStr>,
          T: AsRef<str>
{
    let using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);
    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(docker_cmd: &Path) -> String {
    image_identifier(is_serving_windows_containers(docker_cmd),
                     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(using_windows_containers: bool, target: target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target) = if using_windows_containers {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS)
    } else {
        let t = match target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2,
            _ => target::X86_64_LINUX,
        };
        (DOCKER_IMAGE, t)
    };

    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or_else(|_| {
                                      format!("{}-{}:{}", img, studio_target, version[0])
                                  })
}

#[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() {
        let windows_container = true;
        assert_eq!(image_identifier(!windows_container, target::X86_64_DARWIN),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX_KERNEL2),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
    }
}
<commit_msg>Removes call to default_cache_key_path()<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,
                         find_command,
                         CACHE_ARTIFACT_PATH,
                         CACHE_KEY_PATH},
                    os::process,
                    package::target},
            license,
            VERSION};
use atty;
use std::{env,
          ffi::{OsStr,
                OsString},
          path::{Path,
                 PathBuf},
          process::{Command,
                    Stdio}};

const DOCKER_CMD: &str = "docker";
const DOCKER_CMD_ENVVAR: &str = "HAB_DOCKER_BINARY";

const DOCKER_IMAGE: &str = "habitat/default-studio";
const DOCKER_WINDOWS_IMAGE: &str = "habitat-docker-registry.bintray.io/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 = find_docker_cmd()?;

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

    let mnt_prefix = if is_serving_windows_containers(&docker_cmd) {
        "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.to_string_lossy().into_owned()
            )
        ));
    }

    let mut volumes = vec![format!("{}:{}{}",
                                   env::current_dir().unwrap().to_string_lossy(),
                                   mnt_prefix,
                                   "/src"),
                           format!("{}:{}/{}",
                                   local_cache_key_path.to_string_lossy(),
                                   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 !is_serving_windows_containers(&docker_cmd)
       && (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 !is_serving_windows_containers(&docker_cmd) {
        check_mounts(&docker_cmd, volumes.iter())?;
    }
    run_container(docker_cmd, &args, volumes.iter(), env_vars.iter())
}

fn find_docker_cmd() -> Result<PathBuf> {
    let docker_cmd = henv::var(DOCKER_CMD_ENVVAR).unwrap_or_else(|_| DOCKER_CMD.to_string());

    match find_command(&docker_cmd) {
        Some(docker_abs_path) => Ok(docker_abs_path),
        None => Err(Error::ExecCommandNotFound(docker_cmd.into())),
    }
}

fn is_image_present(docker_cmd: &Path) -> bool {
    let mut cmd = Command::new(docker_cmd);
    let image = image_identifier_for_active_target(&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) -> Result<()> {
    let image = image_identifier_for_active_target(&docker_cmd);
    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) -> Result<()>
    where I: IntoIterator<Item = S>,
          S: AsRef<OsStr>
{
    let mut cmd_args: Vec<OsString> = vec!["run".into(), "--rm".into()];
    let image = image_identifier_for_active_target(&docker_cmd);

    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)
                             -> Result<()>
    where I: IntoIterator<Item = S>,
          J: IntoIterator<Item = T>,
          S: AsRef<OsStr>,
          T: AsRef<str>
{
    let using_windows_containers = is_serving_windows_containers(&docker_cmd);
    let image = image_identifier_for_active_target(&docker_cmd);
    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(docker_cmd: &Path) -> String {
    image_identifier(is_serving_windows_containers(docker_cmd),
                     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(using_windows_containers: bool, target: target::PackageTarget) -> String {
    let version: Vec<&str> = VERSION.split('/').collect();
    let (img, studio_target) = if using_windows_containers {
        (DOCKER_WINDOWS_IMAGE, target::X86_64_WINDOWS)
    } else {
        let t = match target {
            target::X86_64_LINUX_KERNEL2 => target::X86_64_LINUX_KERNEL2,
            _ => target::X86_64_LINUX,
        };
        (DOCKER_IMAGE, t)
    };

    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or_else(|_| {
                                      format!("{}-{}:{}", img, studio_target, version[0])
                                  })
}

#[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() {
        let windows_container = true;
        assert_eq!(image_identifier(!windows_container, target::X86_64_DARWIN),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_LINUX_KERNEL2),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux-kernel2", VERSION));
        assert_eq!(image_identifier(!windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_IMAGE, "x86_64-linux", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_WINDOWS),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
        assert_eq!(image_identifier(windows_container, target::X86_64_LINUX),
                   format!("{}-{}:{}", DOCKER_WINDOWS_IMAGE, "x86_64-windows", VERSION));
    }
}
<|endoftext|>