<commit_before>/// Emit diagnostics using the `annotate-snippets` library
///
/// This is the equivalent of `./emitter.rs` but making use of the
/// [`annotate-snippets`][annotate_snippets] library instead of building the output ourselves.
///
/// [annotate_snippets]: https://docs.rs/crate/annotate-snippets/

use syntax_pos::{SourceFile, MultiSpan, Loc};
use crate::{
    Level, CodeSuggestion, DiagnosticBuilder, Emitter,
    SourceMapperDyn, SubDiagnostic, DiagnosticId
};
use crate::emitter::FileWithAnnotatedLines;
use rustc_data_structures::sync::Lrc;
use crate::snippet::Line;
use annotate_snippets::snippet::*;
use annotate_snippets::display_list::DisplayList;
use annotate_snippets::formatter::DisplayListFormatter;

/// Generates diagnostics using annotate-rs
pub struct AnnotateRsEmitterWriter {
    source_map: Option<Lrc<SourceMapperDyn>>,
    /// If true, hides the longer explanation text
    short_message: bool,
    /// If true, will normalize line numbers with LL to prevent noise in UI test diffs.
    ui_testing: bool,
}

impl Emitter for AnnotateRsEmitterWriter {
    /// The entry point for the diagnostics generation
    fn emit_diagnostic(&mut self, db: &DiagnosticBuilder<'_>) {
        let primary_span = db.span.clone();
        let children = db.children.clone();
        // FIXME(#59346): Collect suggestions (see emitter.rs)
        let suggestions: &[_] = &[];

        // FIXME(#59346): Add `fix_multispans_in_std_macros` function from emitter.rs

        self.emit_messages_default(&db.level,
                                   db.message(),
                                   &db.code,
                                   &primary_span,
                                   &children,
                                   &suggestions);
    }

    fn should_show_explain(&self) -> bool {
        !self.short_message
    }
}

/// Collects all the data needed to generate the data structures needed for the
/// `annotate-snippets` library.
struct DiagnosticConverter<'a> {
    source_map: Option<Lrc<SourceMapperDyn>>,
    level: Level,
    message: String,
    code: Option<DiagnosticId>,
    msp: MultiSpan,
    #[allow(dead_code)]
    children: &'a [SubDiagnostic],
    #[allow(dead_code)]
    suggestions: &'a [CodeSuggestion]
}

impl<'a>  DiagnosticConverter<'a> {
    /// Turns rustc Diagnostic information into a `annotate_snippets::snippet::Snippet`.
    fn to_annotation_snippet(&self) -> Option<Snippet> {
        if let Some(source_map) = &self.source_map {
            // Make sure our primary file comes first
            let primary_lo = if let Some(ref primary_span) =
                self.msp.primary_span().as_ref() {
                source_map.lookup_char_pos(primary_span.lo())
            } else {
                // FIXME(#59346): Not sure when this is the case and what
                // should be done if it happens
                return None
            };
            let annotated_files = FileWithAnnotatedLines::collect_annotations(
                &self.msp,
                &self.source_map
            );
            let slices = self.slices_for_files(annotated_files, primary_lo);

            Some(Snippet {
                title: Some(Annotation {
                    label: Some(self.message.to_string()),
                    id: self.code.clone().map(|c| {
                        match c {
                            DiagnosticId::Error(val) | DiagnosticId::Lint(val) => val
                        }
                    }),
                    annotation_type: Self::annotation_type_for_level(self.level),
                }),
                footer: vec![],
                slices: slices,
            })
        } else {
            // FIXME(#59346): Is it ok to return None if there's no source_map?
            None
        }
    }

    fn slices_for_files(
        &self,
        annotated_files: Vec<FileWithAnnotatedLines>,
        primary_lo: Loc
    ) -> Vec<Slice> {
        // FIXME(#59346): Provide a test case where `annotated_files` is > 1
        annotated_files.iter().flat_map(|annotated_file| {
            annotated_file.lines.iter().map(|line| {
                let line_source = Self::source_string(annotated_file.file.clone(), &line);
                Slice {
                    source: line_source,
                    line_start: line.line_index,
                    origin: Some(primary_lo.file.name.to_string()),
                    // FIXME(#59346): Not really sure when `fold` should be true or false
                    fold: false,
                    annotations: line.annotations.iter().map(|a| {
                        self.annotation_to_source_annotation(a.clone())
                    }).collect(),
                }
            }).collect::<Vec<Slice>>()
        }).collect::<Vec<Slice>>()
    }

    /// Turns a `crate::snippet::Annotation` into a `SourceAnnotation`
    fn annotation_to_source_annotation(
        &self,
        annotation: crate::snippet::Annotation
    ) -> SourceAnnotation {
        SourceAnnotation {
            range: (annotation.start_col, annotation.end_col),
            label: annotation.label.unwrap_or("".to_string()),
            annotation_type: Self::annotation_type_for_level(self.level)
        }
    }

    /// Provides the source string for the given `line` of `file`
    fn source_string(file: Lrc<SourceFile>,
                     line: &Line) -> String {
        let source_string = match file.get_line(line.line_index - 1) {
            Some(s) => s.clone(),
            None => return String::new(),
        };
        source_string.to_string()
    }

    /// Maps `Diagnostic::Level` to `snippet::AnnotationType`
    fn annotation_type_for_level(level: Level) -> AnnotationType {
        match level {
            Level::Bug | Level::Fatal | Level::PhaseFatal | Level::Error => AnnotationType::Error,
            Level::Warning => AnnotationType::Warning,
            Level::Note => AnnotationType::Note,
            Level::Help => AnnotationType::Help,
            // FIXME(#59346): Not sure how to map these two levels
            Level::Cancelled | Level::FailureNote => AnnotationType::Error
        }
    }
}

impl AnnotateRsEmitterWriter {
    pub fn new(
        source_map: Option<Lrc<SourceMapperDyn>>,
        short_message: bool
    ) -> Self {
        Self {
            source_map,
            short_message,
            ui_testing: false,
        }
    }

    /// Allows to modify `Self` to enable or disable the `ui_testing` flag.
    ///
    /// If this is set to true, line numbers will be normalized as `LL` in the output.
    // FIXME(#59346): This method is used via the public interface, but setting the `ui_testing`
    // flag currently does not anonymize line numbers. We would have to add the `maybe_anonymized`
    // method from `emitter.rs` and implement rust-lang/annotate-snippets-rs#2 in order to
    // anonymize line numbers.
    pub fn ui_testing(mut self, ui_testing: bool) -> Self {
        self.ui_testing = ui_testing;
        self
    }

    fn emit_messages_default(&mut self,
                             level: &Level,
                             message: String,
                             code: &Option<DiagnosticId>,
                             msp: &MultiSpan,
                             children: &[SubDiagnostic],
                             suggestions: &[CodeSuggestion]
    ) {
        let converter = DiagnosticConverter {
            source_map: self.source_map.clone(),
            level: level.clone(),
            message: message.clone(),
            code: code.clone(),
            msp: msp.clone(),
            children,
            suggestions
        };
        if let Some(snippet) = converter.to_annotation_snippet() {
            let dl = DisplayList::from(snippet);
            let dlf = DisplayListFormatter::new(true);
            // FIXME(#59346): Figure out if we can _always_ print to stderr or not.
            // `emitter.rs` has the `Destination` enum that lists various possible output
            // destinations.
            eprint!("{}", dlf.format(&dl));
        };
    }
}
<commit_msg>Simplify source_string and block-format methods<commit_after>/// Emit diagnostics using the `annotate-snippets` library
///
/// This is the equivalent of `./emitter.rs` but making use of the
/// [`annotate-snippets`][annotate_snippets] library instead of building the output ourselves.
///
/// [annotate_snippets]: https://docs.rs/crate/annotate-snippets/

use syntax_pos::{SourceFile, MultiSpan, Loc};
use crate::{
    Level, CodeSuggestion, DiagnosticBuilder, Emitter,
    SourceMapperDyn, SubDiagnostic, DiagnosticId
};
use crate::emitter::FileWithAnnotatedLines;
use rustc_data_structures::sync::Lrc;
use crate::snippet::Line;
use annotate_snippets::snippet::*;
use annotate_snippets::display_list::DisplayList;
use annotate_snippets::formatter::DisplayListFormatter;

/// Generates diagnostics using annotate-rs
pub struct AnnotateRsEmitterWriter {
    source_map: Option<Lrc<SourceMapperDyn>>,
    /// If true, hides the longer explanation text
    short_message: bool,
    /// If true, will normalize line numbers with LL to prevent noise in UI test diffs.
    ui_testing: bool,
}

impl Emitter for AnnotateRsEmitterWriter {
    /// The entry point for the diagnostics generation
    fn emit_diagnostic(&mut self, db: &DiagnosticBuilder<'_>) {
        let primary_span = db.span.clone();
        let children = db.children.clone();
        // FIXME(#59346): Collect suggestions (see emitter.rs)
        let suggestions: &[_] = &[];

        // FIXME(#59346): Add `fix_multispans_in_std_macros` function from emitter.rs

        self.emit_messages_default(&db.level,
                                   db.message(),
                                   &db.code,
                                   &primary_span,
                                   &children,
                                   &suggestions);
    }

    fn should_show_explain(&self) -> bool {
        !self.short_message
    }
}

/// Collects all the data needed to generate the data structures needed for the
/// `annotate-snippets` library.
struct DiagnosticConverter<'a> {
    source_map: Option<Lrc<SourceMapperDyn>>,
    level: Level,
    message: String,
    code: Option<DiagnosticId>,
    msp: MultiSpan,
    #[allow(dead_code)]
    children: &'a [SubDiagnostic],
    #[allow(dead_code)]
    suggestions: &'a [CodeSuggestion]
}

impl<'a>  DiagnosticConverter<'a> {
    /// Turns rustc Diagnostic information into a `annotate_snippets::snippet::Snippet`.
    fn to_annotation_snippet(&self) -> Option<Snippet> {
        if let Some(source_map) = &self.source_map {
            // Make sure our primary file comes first
            let primary_lo = if let Some(ref primary_span) =
                self.msp.primary_span().as_ref() {
                source_map.lookup_char_pos(primary_span.lo())
            } else {
                // FIXME(#59346): Not sure when this is the case and what
                // should be done if it happens
                return None
            };
            let annotated_files = FileWithAnnotatedLines::collect_annotations(
                &self.msp,
                &self.source_map
            );
            let slices = self.slices_for_files(annotated_files, primary_lo);

            Some(Snippet {
                title: Some(Annotation {
                    label: Some(self.message.to_string()),
                    id: self.code.clone().map(|c| {
                        match c {
                            DiagnosticId::Error(val) | DiagnosticId::Lint(val) => val
                        }
                    }),
                    annotation_type: Self::annotation_type_for_level(self.level),
                }),
                footer: vec![],
                slices: slices,
            })
        } else {
            // FIXME(#59346): Is it ok to return None if there's no source_map?
            None
        }
    }

    fn slices_for_files(
        &self,
        annotated_files: Vec<FileWithAnnotatedLines>,
        primary_lo: Loc
    ) -> Vec<Slice> {
        // FIXME(#59346): Provide a test case where `annotated_files` is > 1
        annotated_files.iter().flat_map(|annotated_file| {
            annotated_file.lines.iter().map(|line| {
                let line_source = Self::source_string(annotated_file.file.clone(), &line);
                Slice {
                    source: line_source,
                    line_start: line.line_index,
                    origin: Some(primary_lo.file.name.to_string()),
                    // FIXME(#59346): Not really sure when `fold` should be true or false
                    fold: false,
                    annotations: line.annotations.iter().map(|a| {
                        self.annotation_to_source_annotation(a.clone())
                    }).collect(),
                }
            }).collect::<Vec<Slice>>()
        }).collect::<Vec<Slice>>()
    }

    /// Turns a `crate::snippet::Annotation` into a `SourceAnnotation`
    fn annotation_to_source_annotation(
        &self,
        annotation: crate::snippet::Annotation
    ) -> SourceAnnotation {
        SourceAnnotation {
            range: (annotation.start_col, annotation.end_col),
            label: annotation.label.unwrap_or("".to_string()),
            annotation_type: Self::annotation_type_for_level(self.level)
        }
    }

    /// Provides the source string for the given `line` of `file`
    fn source_string(
        file: Lrc<SourceFile>,
        line: &Line
    ) -> String {
        file.get_line(line.line_index - 1).map(|a| a.to_string()).unwrap_or(String::new())
    }

    /// Maps `Diagnostic::Level` to `snippet::AnnotationType`
    fn annotation_type_for_level(level: Level) -> AnnotationType {
        match level {
            Level::Bug | Level::Fatal | Level::PhaseFatal | Level::Error => AnnotationType::Error,
            Level::Warning => AnnotationType::Warning,
            Level::Note => AnnotationType::Note,
            Level::Help => AnnotationType::Help,
            // FIXME(#59346): Not sure how to map these two levels
            Level::Cancelled | Level::FailureNote => AnnotationType::Error
        }
    }
}

impl AnnotateRsEmitterWriter {
    pub fn new(
        source_map: Option<Lrc<SourceMapperDyn>>,
        short_message: bool
    ) -> Self {
        Self {
            source_map,
            short_message,
            ui_testing: false,
        }
    }

    /// Allows to modify `Self` to enable or disable the `ui_testing` flag.
    ///
    /// If this is set to true, line numbers will be normalized as `LL` in the output.
    // FIXME(#59346): This method is used via the public interface, but setting the `ui_testing`
    // flag currently does not anonymize line numbers. We would have to add the `maybe_anonymized`
    // method from `emitter.rs` and implement rust-lang/annotate-snippets-rs#2 in order to
    // anonymize line numbers.
    pub fn ui_testing(mut self, ui_testing: bool) -> Self {
        self.ui_testing = ui_testing;
        self
    }

    fn emit_messages_default(
        &mut self,
        level: &Level,
        message: String,
        code: &Option<DiagnosticId>,
        msp: &MultiSpan,
        children: &[SubDiagnostic],
        suggestions: &[CodeSuggestion]
    ) {
        let converter = DiagnosticConverter {
            source_map: self.source_map.clone(),
            level: level.clone(),
            message: message.clone(),
            code: code.clone(),
            msp: msp.clone(),
            children,
            suggestions
        };
        if let Some(snippet) = converter.to_annotation_snippet() {
            let dl = DisplayList::from(snippet);
            let dlf = DisplayListFormatter::new(true);
            // FIXME(#59346): Figure out if we can _always_ print to stderr or not.
            // `emitter.rs` has the `Destination` enum that lists various possible output
            // destinations.
            eprint!("{}", dlf.format(&dl));
        };
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::Write;

use bufstream::BufStream;

use common::consts;
use protocol::cmd::Cmd;
use protocol::cmd::Get;
use protocol::cmd::Resp;
use protocol::cmd::Set;
use protocol::cmd::SetInstr;

use super::conversions::as_number;
use super::conversions::as_string;
use super::errors::TcpTransportError;
use super::stats::TransportStats;
use super::typedefs::TcpTransportResult;

pub struct TcpTransport<T: Read + Write> {
    stream: BufStream<T>,

    pub line_buffer: Vec<u8>,
    pub line_cursor: usize,
    pub line_break_pos: usize,

    key_maxlen: u64,

    stats: TransportStats,
}

impl<T: Read + Write> TcpTransport<T> {
    pub fn new(stream: T) -> TcpTransport<T> {
        TcpTransport {
            key_maxlen: 250, // memcached standard

            // Used to read the first line of a command, which includes a
            // keyword, a key, flags and a bytecount. We don't expect it to be
            // much longer than the key itself. If it is we panic...
            line_buffer: vec![0; 250 + 100],
            line_cursor: 0,
            line_break_pos: 0,

            stats: TransportStats::new(),
            stream: BufStream::new(stream),
        }
    }

    pub fn with_key_maxlen(&mut self,
                           key_maxlen: u64)
                           -> &mut TcpTransport<T> {
        self.key_maxlen = key_maxlen;
        self
    }

    pub fn get_stats_clone(&self) -> TransportStats {
        self.stats.clone()
    }

    pub fn get_stream(&self) -> &T {
        self.stream.get_ref()
    }

    // Basic bytes manipulation and reading from the stream

    pub fn read_byte(&mut self) -> TcpTransportResult<u8> {
        let mut bytes = [0; 1];

        match self.stream.read(&mut bytes) {
            Ok(1) => {
                // Update stats
                self.stats.bytes_read += 1;

                Ok(bytes[0])
            }
            _ => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn read_bytes(&mut self, len: u64) -> TcpTransportResult<Vec<u8>> {
        // TODO: test large values >4kb
        let mut bytes = vec![0; len as usize];

        match self.stream.read(&mut bytes[..]) {
            Ok(n) => {
                // Update stats
                self.stats.bytes_read += n as u64;

                if (n as u64) < len {
                    bytes.truncate(n);
                }

                Ok(bytes)
            }
            _ => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn read_word_in_line(&mut self) -> TcpTransportResult<(Vec<u8>, bool)> {
        let mut word = vec![];
        let mut byte = [0; 1];
        let mut end_of_line = false;

        loop {
            // Read a byte
            let rv = self.stream.read(&mut byte);

            // If there was an error or if there was nothing to read we bail
            if rv.is_err() || rv.unwrap() == 0 {
                return Err(TcpTransportError::StreamReadError);
            }

            // Update stats
            self.stats.bytes_read += 1;

            if byte[0] == consts::BYTE_SPACE {
                // We found a space

                if word.is_empty() {
                    // If it's one or more leading space we ignore it
                    continue;
                }

                // All good, we've found the end of the word
                break;

            } else if byte[0] == consts::BYTE_CARRIAGE_RETURN {
                // We found \r, we think it's the end of the line

                // Try to read \n
                let rv = self.stream.read(&mut byte);

                // If there was an error or if there was nothing to read we bail
                if rv.is_err() || rv.unwrap() == 0 {
                    return Err(TcpTransportError::StreamReadError);
                }

                // Update stats
                self.stats.bytes_read += 1;

                // If it's not a correct end of line we storm out in protest
                if byte[0] != consts::BYTE_LINE_FEED {
                    return Err(TcpTransportError::LineReadError);
                }

                // Else it's all good, we've read the whole line including the
                // terminator
                end_of_line = true;
                break;

            } else {
                // It's not a special char, append to our word
                word.push(byte[0]);
            }
        }

        Ok((word, end_of_line))
    }

    pub fn read_line_as_words(&mut self) -> TcpTransportResult<Vec<Vec<u8>>> {
        let mut words = vec![];

        loop {
            let (word, end_of_line) = try!(self.read_word_in_line());

            // Don't bother if it's an empty word (trailing space before \r\n)
            if !word.is_empty() {
                words.push(word);
            }

            if end_of_line {
                break;
            }
        }

        Ok(words)
    }

    // Writing to the stream

    pub fn flush_writes(&mut self) -> TcpTransportResult<()> {
        match self.stream.flush() {
            Ok(_) => Ok(()),
            Err(_) => Err(TcpTransportError::StreamWriteError),
        }
    }

    pub fn write_bytes(&mut self,
                       bytes: &Vec<u8>)
                       -> TcpTransportResult<usize> {
        match self.stream.write(bytes) {
            Ok(cnt_written) => {
                // Update stats
                self.stats.bytes_written += cnt_written as u64;

                Ok(cnt_written)
            }
            Err(_) => Err(TcpTransportError::StreamWriteError),
        }
    }

    pub fn write_string(&mut self, string: &str) -> TcpTransportResult<usize> {
        let bytes = string.to_string().into_bytes();
        Ok(try!(self.write_bytes(&bytes)))
    }

    // Parse individual commands

    pub fn parse_cmd_get(&mut self) -> TcpTransportResult<Cmd> {
        // parse the key
        let (key, end_of_line) = try!(self.read_word_in_line());
        let key_str = try!(as_string(key));

        if !end_of_line {
            return Err(TcpTransportError::CommandParseError);
        }

        Ok(Cmd::Get(Get { key: key_str }))
    }

    pub fn parse_cmd_set(&mut self) -> TcpTransportResult<Cmd> {
        // parse the key
        let key_str = {
            let (key, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_string(key))
        };

        // parse the flags
        let flags_num = {
            let (flags, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_number::<u16>(flags))
        };

        // parse the exptime
        let exptime_num = {
            let (exptime, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_number::<u32>(exptime))
        };

        // parse the bytelen
        let bytelen_num = {
            let (bytelen, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_number::<u64>(bytelen))
        };

        // parse noreply
        let noreply_flag = {
            let (noreply, end_of_line) = try!(self.read_word_in_line());

            if !end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            let noreply_str = try!(as_string(noreply));
            match noreply_str == "noreply" {
                true => true,
                false => false,
            }
        };

        // We know the byte length, so now read the value
        let value = try!(self.read_bytes(bytelen_num));

        // Read the line termination marker and verify it
        let newline = try!(self.read_bytes(2));
        if !newline.starts_with(&[consts::BYTE_CARRIAGE_RETURN,
                                  consts::BYTE_LINE_FEED]) {
            return Err(TcpTransportError::CommandParseError);
        }

        // We got all the values we expected and there is nothing left
        return Ok(Cmd::Set(Set {
            instr: SetInstr::Set,
            key: key_str,
            flags: flags_num,
            exptime: exptime_num,
            data: value,
            noreply: noreply_flag,
        }));
    }

    // High level functions

    pub fn read_cmd(&mut self) -> TcpTransportResult<Cmd> {
        let keyword_str = {
            let (word, end_of_line) = try!(self.read_word_in_line());
            try!(as_string(word))
        };

        if keyword_str == "get" {
            // TODO check for !eol
            return self.parse_cmd_get();
        } else if keyword_str == "set" {
            // TODO check for !eol
            return self.parse_cmd_set();
        } else if keyword_str == "stats" {
            // TODO check for eol
            return Ok(Cmd::Stats);
        }

        Err(TcpTransportError::InvalidCmd)
    }

    pub fn write_resp(&mut self, resp: &Resp) -> TcpTransportResult<()> {
        match *resp {
            Resp::Empty => (),
            Resp::Error => {
                try!(self.write_string("ERROR\r\n"));
            }
            Resp::Stats(ref stats) => {
                for stat in stats {
                    try!(self.write_string("STAT "));
                    try!(self.write_string(&stat.key));
                    try!(self.write_string(" "));
                    try!(self.write_string(&stat.value));
                    try!(self.write_string("\r\n"));
                }
                try!(self.write_string("END\r\n"));
            }
            Resp::Stored => {
                try!(self.write_string("STORED\r\n"));
            }
            Resp::Value(ref value) => {
                try!(self.write_string("VALUE ")); // keyword
                try!(self.write_string(&value.key)); // key
                try!(self.write_string(" ")); // space
                try!(self.write_string(&value.flags.to_string())); // flags
                try!(self.write_string(" ")); // space
                try!(self.write_string(&value.data.len().to_string())); // bytelen
                try!(self.write_string(&"\r\n".to_string())); // newline
                try!(self.write_bytes(&value.data)); // data block
                try!(self.write_string(&"\r\n".to_string())); // newline
                try!(self.write_string(&"END\r\n".to_string())); // END + newline
            }
            _ => {
                return Err(TcpTransportError::StreamWriteError);
            }
        }

        // Make sure all bytes were actually sent
        self.flush_writes()
    }
}
<commit_msg>check the size of the value<commit_after>use std::io::Read;
use std::io::Write;

use bufstream::BufStream;

use common::consts;
use protocol::cmd::Cmd;
use protocol::cmd::Get;
use protocol::cmd::Resp;
use protocol::cmd::Set;
use protocol::cmd::SetInstr;

use super::conversions::as_number;
use super::conversions::as_string;
use super::errors::TcpTransportError;
use super::stats::TransportStats;
use super::typedefs::TcpTransportResult;

pub struct TcpTransport<T: Read + Write> {
    stream: BufStream<T>,

    pub line_buffer: Vec<u8>,
    pub line_cursor: usize,
    pub line_break_pos: usize,

    key_maxlen: u64,

    stats: TransportStats,
}

impl<T: Read + Write> TcpTransport<T> {
    pub fn new(stream: T) -> TcpTransport<T> {
        TcpTransport {
            key_maxlen: 250, // memcached standard

            // Used to read the first line of a command, which includes a
            // keyword, a key, flags and a bytecount. We don't expect it to be
            // much longer than the key itself. If it is we panic...
            line_buffer: vec![0; 250 + 100],
            line_cursor: 0,
            line_break_pos: 0,

            stats: TransportStats::new(),
            stream: BufStream::new(stream),
        }
    }

    pub fn with_key_maxlen(&mut self,
                           key_maxlen: u64)
                           -> &mut TcpTransport<T> {
        self.key_maxlen = key_maxlen;
        self
    }

    pub fn get_stats_clone(&self) -> TransportStats {
        self.stats.clone()
    }

    pub fn get_stream(&self) -> &T {
        self.stream.get_ref()
    }

    // Basic bytes manipulation and reading from the stream

    pub fn read_byte(&mut self) -> TcpTransportResult<u8> {
        let mut bytes = [0; 1];

        match self.stream.read(&mut bytes) {
            Ok(1) => {
                // Update stats
                self.stats.bytes_read += 1;

                Ok(bytes[0])
            }
            _ => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn read_bytes(&mut self, len: u64) -> TcpTransportResult<Vec<u8>> {
        // TODO: test large values >4kb
        let mut bytes = vec![0; len as usize];

        match self.stream.read(&mut bytes[..]) {
            Ok(n) => {
                // Update stats
                self.stats.bytes_read += n as u64;

                if (n as u64) < len {
                    bytes.truncate(n);
                }

                Ok(bytes)
            }
            _ => Err(TcpTransportError::StreamReadError),
        }
    }

    pub fn read_word_in_line(&mut self) -> TcpTransportResult<(Vec<u8>, bool)> {
        let mut word = vec![];
        let mut byte = [0; 1];
        let mut end_of_line = false;

        loop {
            // Read a byte
            let rv = self.stream.read(&mut byte);

            // If there was an error or if there was nothing to read we bail
            if rv.is_err() || rv.unwrap() == 0 {
                return Err(TcpTransportError::StreamReadError);
            }

            // Update stats
            self.stats.bytes_read += 1;

            if byte[0] == consts::BYTE_SPACE {
                // We found a space

                if word.is_empty() {
                    // If it's one or more leading space we ignore it
                    continue;
                }

                // All good, we've found the end of the word
                break;

            } else if byte[0] == consts::BYTE_CARRIAGE_RETURN {
                // We found \r, we think it's the end of the line

                // Try to read \n
                let rv = self.stream.read(&mut byte);

                // If there was an error or if there was nothing to read we bail
                if rv.is_err() || rv.unwrap() == 0 {
                    return Err(TcpTransportError::StreamReadError);
                }

                // Update stats
                self.stats.bytes_read += 1;

                // If it's not a correct end of line we storm out in protest
                if byte[0] != consts::BYTE_LINE_FEED {
                    return Err(TcpTransportError::LineReadError);
                }

                // Else it's all good, we've read the whole line including the
                // terminator
                end_of_line = true;
                break;

            } else {
                // It's not a special char, append to our word
                word.push(byte[0]);
            }
        }

        Ok((word, end_of_line))
    }

    pub fn read_line_as_words(&mut self) -> TcpTransportResult<Vec<Vec<u8>>> {
        let mut words = vec![];

        loop {
            let (word, end_of_line) = try!(self.read_word_in_line());

            // Don't bother if it's an empty word (trailing space before \r\n)
            if !word.is_empty() {
                words.push(word);
            }

            if end_of_line {
                break;
            }
        }

        Ok(words)
    }

    // Writing to the stream

    pub fn flush_writes(&mut self) -> TcpTransportResult<()> {
        match self.stream.flush() {
            Ok(_) => Ok(()),
            Err(_) => Err(TcpTransportError::StreamWriteError),
        }
    }

    pub fn write_bytes(&mut self,
                       bytes: &Vec<u8>)
                       -> TcpTransportResult<usize> {
        match self.stream.write(bytes) {
            Ok(cnt_written) => {
                // Update stats
                self.stats.bytes_written += cnt_written as u64;

                Ok(cnt_written)
            }
            Err(_) => Err(TcpTransportError::StreamWriteError),
        }
    }

    pub fn write_string(&mut self, string: &str) -> TcpTransportResult<usize> {
        let bytes = string.to_string().into_bytes();
        Ok(try!(self.write_bytes(&bytes)))
    }

    // Parse individual commands

    pub fn parse_cmd_get(&mut self) -> TcpTransportResult<Cmd> {
        // parse the key
        let (key, end_of_line) = try!(self.read_word_in_line());
        let key_str = try!(as_string(key));

        if !end_of_line {
            return Err(TcpTransportError::CommandParseError);
        }

        Ok(Cmd::Get(Get { key: key_str }))
    }

    pub fn parse_cmd_set(&mut self) -> TcpTransportResult<Cmd> {
        // parse the key
        let key_str = {
            let (key, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_string(key))
        };

        // parse the flags
        let flags_num = {
            let (flags, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_number::<u16>(flags))
        };

        // parse the exptime
        let exptime_num = {
            let (exptime, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_number::<u32>(exptime))
        };

        // parse the bytelen
        let bytelen_num = {
            let (bytelen, end_of_line) = try!(self.read_word_in_line());

            if end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            try!(as_number::<u64>(bytelen))
        };

        // parse noreply
        let noreply_flag = {
            let (noreply, end_of_line) = try!(self.read_word_in_line());

            if !end_of_line {
                return Err(TcpTransportError::CommandParseError);
            }

            let noreply_str = try!(as_string(noreply));
            match noreply_str == "noreply" {
                true => true,
                false => false,
            }
        };

        // We now know the byte length, so read the value
        let value = try!(self.read_bytes(bytelen_num));

        // The value is the wrong size
        if value.len() as u64 != bytelen_num {
            return Err(TcpTransportError::CommandParseError);
        }

        // Verify that we found the line terminator
        let terminator = try!(self.read_bytes(2));
        if !terminator.ends_with(&[consts::BYTE_CARRIAGE_RETURN,
                                   consts::BYTE_LINE_FEED]) {
            return Err(TcpTransportError::CommandParseError);
        }

        // We got all the values we expected and there is nothing left
        return Ok(Cmd::Set(Set {
            instr: SetInstr::Set,
            key: key_str,
            flags: flags_num,
            exptime: exptime_num,
            data: value,
            noreply: noreply_flag,
        }));
    }

    // High level functions

    pub fn read_cmd(&mut self) -> TcpTransportResult<Cmd> {
        let keyword_str = {
            let (word, end_of_line) = try!(self.read_word_in_line());
            try!(as_string(word))
        };

        if keyword_str == "get" {
            // TODO check for !eol
            return self.parse_cmd_get();
        } else if keyword_str == "set" {
            // TODO check for !eol
            return self.parse_cmd_set();
        } else if keyword_str == "stats" {
            // TODO check for eol
            return Ok(Cmd::Stats);
        }

        Err(TcpTransportError::InvalidCmd)
    }

    pub fn write_resp(&mut self, resp: &Resp) -> TcpTransportResult<()> {
        match *resp {
            Resp::Empty => (),
            Resp::Error => {
                try!(self.write_string("ERROR\r\n"));
            }
            Resp::Stats(ref stats) => {
                for stat in stats {
                    try!(self.write_string("STAT "));
                    try!(self.write_string(&stat.key));
                    try!(self.write_string(" "));
                    try!(self.write_string(&stat.value));
                    try!(self.write_string("\r\n"));
                }
                try!(self.write_string("END\r\n"));
            }
            Resp::Stored => {
                try!(self.write_string("STORED\r\n"));
            }
            Resp::Value(ref value) => {
                try!(self.write_string("VALUE ")); // keyword
                try!(self.write_string(&value.key)); // key
                try!(self.write_string(" ")); // space
                try!(self.write_string(&value.flags.to_string())); // flags
                try!(self.write_string(" ")); // space
                try!(self.write_string(&value.data.len().to_string())); // bytelen
                try!(self.write_string(&"\r\n".to_string())); // newline
                try!(self.write_bytes(&value.data)); // data block
                try!(self.write_string(&"\r\n".to_string())); // newline
                try!(self.write_string(&"END\r\n".to_string())); // END + newline
            }
            _ => {
                return Err(TcpTransportError::StreamWriteError);
            }
        }

        // Make sure all bytes were actually sent
        self.flush_writes()
    }
}
<|endoftext|>


<commit_before>use crate::syntax::atom::Atom::{self, *};
use crate::syntax::namespace::Namespace;
use crate::syntax::report::Errors;
use crate::syntax::{
    error, ident, Api, Enum, ExternFn, ExternType, Lang, Receiver, Ref, Slice, Struct, Ty1, Type,
    Types,
};
use proc_macro2::{Delimiter, Group, Ident, TokenStream};
use quote::{quote, ToTokens};
use std::fmt::Display;

pub(crate) struct Check<'a> {
    namespace: &'a Namespace,
    apis: &'a [Api],
    types: &'a Types<'a>,
    errors: &'a mut Errors,
}

pub(crate) fn typecheck(cx: &mut Errors, namespace: &Namespace, apis: &[Api], types: &Types) {
    do_typecheck(&mut Check {
        namespace,
        apis,
        types,
        errors: cx,
    });
}

fn do_typecheck(cx: &mut Check) {
    ident::check_all(cx, cx.namespace, cx.apis);

    for ty in cx.types {
        match ty {
            Type::Ident(ident) => check_type_ident(cx, ident),
            Type::RustBox(ptr) => check_type_box(cx, ptr),
            Type::RustVec(ty) => check_type_rust_vec(cx, ty),
            Type::UniquePtr(ptr) => check_type_unique_ptr(cx, ptr),
            Type::CxxVector(ptr) => check_type_cxx_vector(cx, ptr),
            Type::Ref(ty) => check_type_ref(cx, ty),
            Type::Slice(ty) => check_type_slice(cx, ty),
            _ => {}
        }
    }

    for api in cx.apis {
        match api {
            Api::Struct(strct) => check_api_struct(cx, strct),
            Api::Enum(enm) => check_api_enum(cx, enm),
            Api::CxxType(ty) | Api::RustType(ty) => check_api_type(cx, ty),
            Api::CxxFunction(efn) | Api::RustFunction(efn) => check_api_fn(cx, efn),
            _ => {}
        }
    }
}

impl Check<'_> {
    pub(crate) fn error(&mut self, sp: impl ToTokens, msg: impl Display) {
        self.errors.error(sp, msg);
    }
}

fn check_type_ident(cx: &mut Check, ident: &Ident) {
    if Atom::from(ident).is_none()
        && !cx.types.structs.contains_key(ident)
        && !cx.types.enums.contains_key(ident)
        && !cx.types.cxx.contains(ident)
        && !cx.types.rust.contains(ident)
    {
        cx.error(ident, "unsupported type");
    }
}

fn check_type_box(cx: &mut Check, ptr: &Ty1) {
    if let Type::Ident(ident) = &ptr.inner {
        if cx.types.cxx.contains(ident) {
            cx.error(ptr, error::BOX_CXX_TYPE.msg);
        }

        if Atom::from(ident).is_none() {
            return;
        }
    }

    cx.error(ptr, "unsupported target type of Box");
}

fn check_type_rust_vec(cx: &mut Check, ty: &Ty1) {
    if let Type::Ident(ident) = &ty.inner {
        if cx.types.cxx.contains(ident) {
            cx.error(ty, "Rust Vec containing C++ type is not supported yet");
            return;
        }

        match Atom::from(ident) {
            None | Some(U8) | Some(U16) | Some(U32) | Some(U64) | Some(Usize) | Some(I8)
            | Some(I16) | Some(I32) | Some(I64) | Some(Isize) | Some(F32) | Some(F64)
            | Some(RustString) => return,
            Some(Bool) => { /* todo */ }
            Some(CxxString) => {}
        }
    }

    cx.error(ty, "unsupported element type of Vec");
}

fn check_type_unique_ptr(cx: &mut Check, ptr: &Ty1) {
    if let Type::Ident(ident) = &ptr.inner {
        if cx.types.rust.contains(ident) {
            cx.error(ptr, "unique_ptr of a Rust type is not supported yet");
        }

        match Atom::from(ident) {
            None | Some(CxxString) => return,
            _ => {}
        }
    } else if let Type::CxxVector(_) = &ptr.inner {
        return;
    }

    cx.error(ptr, "unsupported unique_ptr target type");
}

fn check_type_cxx_vector(cx: &mut Check, ptr: &Ty1) {
    if let Type::Ident(ident) = &ptr.inner {
        if cx.types.rust.contains(ident) {
            cx.error(
                ptr,
                "C++ vector containing a Rust type is not supported yet",
            );
        }

        match Atom::from(ident) {
            None | Some(U8) | Some(U16) | Some(U32) | Some(U64) | Some(Usize) | Some(I8)
            | Some(I16) | Some(I32) | Some(I64) | Some(Isize) | Some(F32) | Some(F64)
            | Some(CxxString) => return,
            Some(Bool) | Some(RustString) => {}
        }
    }

    cx.error(ptr, "unsupported vector target type");
}

fn check_type_ref(cx: &mut Check, ty: &Ref) {
    if ty.lifetime.is_some() {
        cx.error(ty, "references with explicit lifetimes are not supported");
    }

    match ty.inner {
        Type::Fn(_) | Type::Void(_) => {}
        Type::Ref(_) => {
            cx.error(ty, "C++ does not allow references to references");
            return;
        }
        _ => return,
    }

    cx.error(ty, "unsupported reference type");
}

fn check_type_slice(cx: &mut Check, ty: &Slice) {
    cx.error(ty, "only &[u8] is supported so far, not other slice types");
}

fn check_api_struct(cx: &mut Check, strct: &Struct) {
    check_reserved_name(cx, &strct.ident);

    if strct.fields.is_empty() {
        let span = span_for_struct_error(strct);
        cx.error(span, "structs without any fields are not supported");
    }

    for field in &strct.fields {
        if is_unsized(cx, &field.ty) {
            let desc = describe(cx, &field.ty);
            let msg = format!("using {} by value is not supported", desc);
            cx.error(field, msg);
        }
        if let Type::Fn(_) = field.ty {
            cx.error(
                field,
                "function pointers in a struct field are not implemented yet",
            );
        }
    }
}

fn check_api_enum(cx: &mut Check, enm: &Enum) {
    check_reserved_name(cx, &enm.ident);

    if enm.variants.is_empty() {
        let span = span_for_enum_error(enm);
        cx.error(span, "enums without any variants are not supported");
    }
}

fn check_api_type(cx: &mut Check, ty: &ExternType) {
    check_reserved_name(cx, &ty.ident);
}

fn check_api_fn(cx: &mut Check, efn: &ExternFn) {
    if let Some(receiver) = &efn.receiver {
        let ref span = span_for_receiver_error(receiver);

        if receiver.ty == "Self" {
            let mutability = match receiver.mutability {
                Some(_) => "mut ",
                None => "",
            };
            let msg = format!(
                "unnamed receiver type is only allowed if the surrounding \
                 extern block contains exactly one extern type; \
                 use `self: &{mutability}TheType`",
                mutability = mutability,
            );
            cx.error(span, msg);
        } else if !cx.types.structs.contains_key(&receiver.ty)
            && !cx.types.cxx.contains(&receiver.ty)
            && !cx.types.rust.contains(&receiver.ty)
        {
            cx.error(span, "unrecognized receiver type");
        }

        if receiver.lifetime.is_some() {
            cx.error(span, "references with explicit lifetimes are not supported");
        }
    }

    for arg in &efn.args {
        if is_unsized(cx, &arg.ty) {
            let desc = describe(cx, &arg.ty);
            let msg = format!("passing {} by value is not supported", desc);
            cx.error(arg, msg);
        }
        if let Type::Fn(_) = arg.ty {
            if efn.lang == Lang::Rust {
                cx.error(
                    arg,
                    "passing a function pointer from C++ to Rust is not implemented yet",
                );
            }
        }
    }

    if let Some(ty) = &efn.ret {
        if is_unsized(cx, ty) {
            let desc = describe(cx, ty);
            let msg = format!("returning {} by value is not supported", desc);
            cx.error(ty, msg);
        }
        if let Type::Fn(_) = ty {
            cx.error(ty, "returning a function pointer is not implemented yet");
        }
    }

    if efn.lang == Lang::Cxx {
        check_mut_return_restriction(cx, efn);
    }

    check_multiple_arg_lifetimes(cx, efn);
}

fn check_mut_return_restriction(cx: &mut Check, efn: &ExternFn) {
    match &efn.ret {
        Some(Type::Ref(ty)) if ty.mutability.is_some() => {}
        _ => return,
    }

    for arg in &efn.args {
        if let Type::Ref(ty) = &arg.ty {
            if ty.mutability.is_some() {
                return;
            }
        }
    }

    cx.error(
        efn,
        "&mut return type is not allowed unless there is a &mut argument",
    );
}

fn check_multiple_arg_lifetimes(cx: &mut Check, efn: &ExternFn) {
    match &efn.ret {
        Some(Type::Ref(_)) => {}
        _ => return,
    }

    let mut reference_args = 0;
    for arg in &efn.args {
        if let Type::Ref(_) = &arg.ty {
            reference_args += 1;
        }
    }

    if efn.receiver.is_some() {
        reference_args += 1;
    }

    if reference_args != 1 {
        cx.error(
            efn,
            "functions that return a reference must take exactly one input reference",
        );
    }
}

fn check_reserved_name(cx: &mut Check, ident: &Ident) {
    if ident == "Box"
        || ident == "UniquePtr"
        || ident == "Vec"
        || ident == "CxxVector"
        || Atom::from(ident).is_some()
    {
        cx.error(ident, "reserved name");
    }
}

fn is_unsized(cx: &mut Check, ty: &Type) -> bool {
    let ident = match ty {
        Type::Ident(ident) => ident,
        Type::CxxVector(_) | Type::Slice(_) | Type::Void(_) => return true,
        _ => return false,
    };
    ident == CxxString || cx.types.cxx.contains(ident) || cx.types.rust.contains(ident)
}

fn span_for_struct_error(strct: &Struct) -> TokenStream {
    let struct_token = strct.struct_token;
    let mut brace_token = Group::new(Delimiter::Brace, TokenStream::new());
    brace_token.set_span(strct.brace_token.span);
    quote!(#struct_token #brace_token)
}

fn span_for_enum_error(enm: &Enum) -> TokenStream {
    let enum_token = enm.enum_token;
    let mut brace_token = Group::new(Delimiter::Brace, TokenStream::new());
    brace_token.set_span(enm.brace_token.span);
    quote!(#enum_token #brace_token)
}

fn span_for_receiver_error(receiver: &Receiver) -> TokenStream {
    let ampersand = receiver.ampersand;
    let lifetime = &receiver.lifetime;
    let mutability = receiver.mutability;
    if receiver.shorthand {
        let var = receiver.var;
        quote!(#ampersand #lifetime #mutability #var)
    } else {
        let ty = &receiver.ty;
        quote!(#ampersand #lifetime #mutability #ty)
    }
}

fn describe(cx: &mut Check, ty: &Type) -> String {
    match ty {
        Type::Ident(ident) => {
            if cx.types.structs.contains_key(ident) {
                "struct".to_owned()
            } else if cx.types.enums.contains_key(ident) {
                "enum".to_owned()
            } else if cx.types.cxx.contains(ident) {
                "C++ type".to_owned()
            } else if cx.types.rust.contains(ident) {
                "opaque Rust type".to_owned()
            } else if Atom::from(ident) == Some(CxxString) {
                "C++ string".to_owned()
            } else {
                ident.to_string()
            }
        }
        Type::RustBox(_) => "Box".to_owned(),
        Type::RustVec(_) => "Vec".to_owned(),
        Type::UniquePtr(_) => "unique_ptr".to_owned(),
        Type::Ref(_) => "reference".to_owned(),
        Type::Str(_) => "&str".to_owned(),
        Type::CxxVector(_) => "C++ vector".to_owned(),
        Type::Slice(_) => "slice".to_owned(),
        Type::SliceRefU8(_) => "&[u8]".to_owned(),
        Type::Fn(_) => "function pointer".to_owned(),
        Type::Void(_) => "()".to_owned(),
    }
}
<commit_msg>Box, Vec of an extern enum is supposed to work<commit_after>use crate::syntax::atom::Atom::{self, *};
use crate::syntax::namespace::Namespace;
use crate::syntax::report::Errors;
use crate::syntax::{
    error, ident, Api, Enum, ExternFn, ExternType, Lang, Receiver, Ref, Slice, Struct, Ty1, Type,
    Types,
};
use proc_macro2::{Delimiter, Group, Ident, TokenStream};
use quote::{quote, ToTokens};
use std::fmt::Display;

pub(crate) struct Check<'a> {
    namespace: &'a Namespace,
    apis: &'a [Api],
    types: &'a Types<'a>,
    errors: &'a mut Errors,
}

pub(crate) fn typecheck(cx: &mut Errors, namespace: &Namespace, apis: &[Api], types: &Types) {
    do_typecheck(&mut Check {
        namespace,
        apis,
        types,
        errors: cx,
    });
}

fn do_typecheck(cx: &mut Check) {
    ident::check_all(cx, cx.namespace, cx.apis);

    for ty in cx.types {
        match ty {
            Type::Ident(ident) => check_type_ident(cx, ident),
            Type::RustBox(ptr) => check_type_box(cx, ptr),
            Type::RustVec(ty) => check_type_rust_vec(cx, ty),
            Type::UniquePtr(ptr) => check_type_unique_ptr(cx, ptr),
            Type::CxxVector(ptr) => check_type_cxx_vector(cx, ptr),
            Type::Ref(ty) => check_type_ref(cx, ty),
            Type::Slice(ty) => check_type_slice(cx, ty),
            _ => {}
        }
    }

    for api in cx.apis {
        match api {
            Api::Struct(strct) => check_api_struct(cx, strct),
            Api::Enum(enm) => check_api_enum(cx, enm),
            Api::CxxType(ty) | Api::RustType(ty) => check_api_type(cx, ty),
            Api::CxxFunction(efn) | Api::RustFunction(efn) => check_api_fn(cx, efn),
            _ => {}
        }
    }
}

impl Check<'_> {
    pub(crate) fn error(&mut self, sp: impl ToTokens, msg: impl Display) {
        self.errors.error(sp, msg);
    }
}

fn check_type_ident(cx: &mut Check, ident: &Ident) {
    if Atom::from(ident).is_none()
        && !cx.types.structs.contains_key(ident)
        && !cx.types.enums.contains_key(ident)
        && !cx.types.cxx.contains(ident)
        && !cx.types.rust.contains(ident)
    {
        cx.error(ident, "unsupported type");
    }
}

fn check_type_box(cx: &mut Check, ptr: &Ty1) {
    if let Type::Ident(ident) = &ptr.inner {
        if cx.types.cxx.contains(ident) && !cx.types.enums.contains_key(ident) {
            cx.error(ptr, error::BOX_CXX_TYPE.msg);
        }

        if Atom::from(ident).is_none() {
            return;
        }
    }

    cx.error(ptr, "unsupported target type of Box");
}

fn check_type_rust_vec(cx: &mut Check, ty: &Ty1) {
    if let Type::Ident(ident) = &ty.inner {
        if cx.types.cxx.contains(ident) && !cx.types.enums.contains_key(ident) {
            cx.error(ty, "Rust Vec containing C++ type is not supported yet");
            return;
        }

        match Atom::from(ident) {
            None | Some(U8) | Some(U16) | Some(U32) | Some(U64) | Some(Usize) | Some(I8)
            | Some(I16) | Some(I32) | Some(I64) | Some(Isize) | Some(F32) | Some(F64)
            | Some(RustString) => return,
            Some(Bool) => { /* todo */ }
            Some(CxxString) => {}
        }
    }

    cx.error(ty, "unsupported element type of Vec");
}

fn check_type_unique_ptr(cx: &mut Check, ptr: &Ty1) {
    if let Type::Ident(ident) = &ptr.inner {
        if cx.types.rust.contains(ident) {
            cx.error(ptr, "unique_ptr of a Rust type is not supported yet");
        }

        match Atom::from(ident) {
            None | Some(CxxString) => return,
            _ => {}
        }
    } else if let Type::CxxVector(_) = &ptr.inner {
        return;
    }

    cx.error(ptr, "unsupported unique_ptr target type");
}

fn check_type_cxx_vector(cx: &mut Check, ptr: &Ty1) {
    if let Type::Ident(ident) = &ptr.inner {
        if cx.types.rust.contains(ident) {
            cx.error(
                ptr,
                "C++ vector containing a Rust type is not supported yet",
            );
        }

        match Atom::from(ident) {
            None | Some(U8) | Some(U16) | Some(U32) | Some(U64) | Some(Usize) | Some(I8)
            | Some(I16) | Some(I32) | Some(I64) | Some(Isize) | Some(F32) | Some(F64)
            | Some(CxxString) => return,
            Some(Bool) | Some(RustString) => {}
        }
    }

    cx.error(ptr, "unsupported vector target type");
}

fn check_type_ref(cx: &mut Check, ty: &Ref) {
    if ty.lifetime.is_some() {
        cx.error(ty, "references with explicit lifetimes are not supported");
    }

    match ty.inner {
        Type::Fn(_) | Type::Void(_) => {}
        Type::Ref(_) => {
            cx.error(ty, "C++ does not allow references to references");
            return;
        }
        _ => return,
    }

    cx.error(ty, "unsupported reference type");
}

fn check_type_slice(cx: &mut Check, ty: &Slice) {
    cx.error(ty, "only &[u8] is supported so far, not other slice types");
}

fn check_api_struct(cx: &mut Check, strct: &Struct) {
    check_reserved_name(cx, &strct.ident);

    if strct.fields.is_empty() {
        let span = span_for_struct_error(strct);
        cx.error(span, "structs without any fields are not supported");
    }

    for field in &strct.fields {
        if is_unsized(cx, &field.ty) {
            let desc = describe(cx, &field.ty);
            let msg = format!("using {} by value is not supported", desc);
            cx.error(field, msg);
        }
        if let Type::Fn(_) = field.ty {
            cx.error(
                field,
                "function pointers in a struct field are not implemented yet",
            );
        }
    }
}

fn check_api_enum(cx: &mut Check, enm: &Enum) {
    check_reserved_name(cx, &enm.ident);

    if enm.variants.is_empty() {
        let span = span_for_enum_error(enm);
        cx.error(span, "enums without any variants are not supported");
    }
}

fn check_api_type(cx: &mut Check, ty: &ExternType) {
    check_reserved_name(cx, &ty.ident);
}

fn check_api_fn(cx: &mut Check, efn: &ExternFn) {
    if let Some(receiver) = &efn.receiver {
        let ref span = span_for_receiver_error(receiver);

        if receiver.ty == "Self" {
            let mutability = match receiver.mutability {
                Some(_) => "mut ",
                None => "",
            };
            let msg = format!(
                "unnamed receiver type is only allowed if the surrounding \
                 extern block contains exactly one extern type; \
                 use `self: &{mutability}TheType`",
                mutability = mutability,
            );
            cx.error(span, msg);
        } else if !cx.types.structs.contains_key(&receiver.ty)
            && !cx.types.cxx.contains(&receiver.ty)
            && !cx.types.rust.contains(&receiver.ty)
        {
            cx.error(span, "unrecognized receiver type");
        }

        if receiver.lifetime.is_some() {
            cx.error(span, "references with explicit lifetimes are not supported");
        }
    }

    for arg in &efn.args {
        if is_unsized(cx, &arg.ty) {
            let desc = describe(cx, &arg.ty);
            let msg = format!("passing {} by value is not supported", desc);
            cx.error(arg, msg);
        }
        if let Type::Fn(_) = arg.ty {
            if efn.lang == Lang::Rust {
                cx.error(
                    arg,
                    "passing a function pointer from C++ to Rust is not implemented yet",
                );
            }
        }
    }

    if let Some(ty) = &efn.ret {
        if is_unsized(cx, ty) {
            let desc = describe(cx, ty);
            let msg = format!("returning {} by value is not supported", desc);
            cx.error(ty, msg);
        }
        if let Type::Fn(_) = ty {
            cx.error(ty, "returning a function pointer is not implemented yet");
        }
    }

    if efn.lang == Lang::Cxx {
        check_mut_return_restriction(cx, efn);
    }

    check_multiple_arg_lifetimes(cx, efn);
}

fn check_mut_return_restriction(cx: &mut Check, efn: &ExternFn) {
    match &efn.ret {
        Some(Type::Ref(ty)) if ty.mutability.is_some() => {}
        _ => return,
    }

    for arg in &efn.args {
        if let Type::Ref(ty) = &arg.ty {
            if ty.mutability.is_some() {
                return;
            }
        }
    }

    cx.error(
        efn,
        "&mut return type is not allowed unless there is a &mut argument",
    );
}

fn check_multiple_arg_lifetimes(cx: &mut Check, efn: &ExternFn) {
    match &efn.ret {
        Some(Type::Ref(_)) => {}
        _ => return,
    }

    let mut reference_args = 0;
    for arg in &efn.args {
        if let Type::Ref(_) = &arg.ty {
            reference_args += 1;
        }
    }

    if efn.receiver.is_some() {
        reference_args += 1;
    }

    if reference_args != 1 {
        cx.error(
            efn,
            "functions that return a reference must take exactly one input reference",
        );
    }
}

fn check_reserved_name(cx: &mut Check, ident: &Ident) {
    if ident == "Box"
        || ident == "UniquePtr"
        || ident == "Vec"
        || ident == "CxxVector"
        || Atom::from(ident).is_some()
    {
        cx.error(ident, "reserved name");
    }
}

fn is_unsized(cx: &mut Check, ty: &Type) -> bool {
    let ident = match ty {
        Type::Ident(ident) => ident,
        Type::CxxVector(_) | Type::Slice(_) | Type::Void(_) => return true,
        _ => return false,
    };
    ident == CxxString
        || cx.types.cxx.contains(ident) && !cx.types.enums.contains_key(ident)
        || cx.types.rust.contains(ident)
}

fn span_for_struct_error(strct: &Struct) -> TokenStream {
    let struct_token = strct.struct_token;
    let mut brace_token = Group::new(Delimiter::Brace, TokenStream::new());
    brace_token.set_span(strct.brace_token.span);
    quote!(#struct_token #brace_token)
}

fn span_for_enum_error(enm: &Enum) -> TokenStream {
    let enum_token = enm.enum_token;
    let mut brace_token = Group::new(Delimiter::Brace, TokenStream::new());
    brace_token.set_span(enm.brace_token.span);
    quote!(#enum_token #brace_token)
}

fn span_for_receiver_error(receiver: &Receiver) -> TokenStream {
    let ampersand = receiver.ampersand;
    let lifetime = &receiver.lifetime;
    let mutability = receiver.mutability;
    if receiver.shorthand {
        let var = receiver.var;
        quote!(#ampersand #lifetime #mutability #var)
    } else {
        let ty = &receiver.ty;
        quote!(#ampersand #lifetime #mutability #ty)
    }
}

fn describe(cx: &mut Check, ty: &Type) -> String {
    match ty {
        Type::Ident(ident) => {
            if cx.types.structs.contains_key(ident) {
                "struct".to_owned()
            } else if cx.types.enums.contains_key(ident) {
                "enum".to_owned()
            } else if cx.types.cxx.contains(ident) {
                "C++ type".to_owned()
            } else if cx.types.rust.contains(ident) {
                "opaque Rust type".to_owned()
            } else if Atom::from(ident) == Some(CxxString) {
                "C++ string".to_owned()
            } else {
                ident.to_string()
            }
        }
        Type::RustBox(_) => "Box".to_owned(),
        Type::RustVec(_) => "Vec".to_owned(),
        Type::UniquePtr(_) => "unique_ptr".to_owned(),
        Type::Ref(_) => "reference".to_owned(),
        Type::Str(_) => "&str".to_owned(),
        Type::CxxVector(_) => "C++ vector".to_owned(),
        Type::Slice(_) => "slice".to_owned(),
        Type::SliceRefU8(_) => "&[u8]".to_owned(),
        Type::Fn(_) => "function pointer".to_owned(),
        Type::Void(_) => "()".to_owned(),
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interface listing implementation for all non-Windows platforms.

use {MacAddr, NetworkInterface};

use ipnetwork::{ip_mask_to_prefix, IpNetwork};
use pnet_sys;

use libc;

use std::ffi::{CStr, CString};
use std::mem;
use std::net::IpAddr;
use std::os::raw::c_char;
use std::str::from_utf8_unchecked;

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    fn merge(old: &mut NetworkInterface, new: &NetworkInterface) {
        old.mac = match new.mac {
            None => old.mac,
            _ => new.mac,
        };
        old.ips.extend_from_slice(&new.ips[..]);
        old.flags = old.flags | new.flags;
    }

    let mut ifaces: Vec<NetworkInterface> = Vec::new();
    unsafe {
        #[allow(deprecated)]
        let mut addrs: *mut libc::ifaddrs = mem::uninitialized();
        if libc::getifaddrs(&mut addrs) != 0 {
            return ifaces;
        }
        let mut addr = addrs;
        while !addr.is_null() {
            let c_str = (*addr).ifa_name as *const c_char;
            let bytes = CStr::from_ptr(c_str).to_bytes();
            let name = from_utf8_unchecked(bytes).to_owned();
            let (mac, ip) = sockaddr_to_network_addr((*addr).ifa_addr as *const libc::sockaddr);
            let (_, netmask) =
                sockaddr_to_network_addr((*addr).ifa_netmask as *const libc::sockaddr);
            let prefix = netmask
                .and_then(|netmask| ip_mask_to_prefix(netmask).ok())
                .unwrap_or(0);
            let network = ip.and_then(|ip| IpNetwork::new(ip, prefix).ok());
            let ni = NetworkInterface {
                name: name.clone(),
                description: "".to_string(),
                index: 0,
                mac: mac,
                ips: network.into_iter().collect(),
                flags: (*addr).ifa_flags,
            };
            let mut found: bool = false;
            for iface in &mut ifaces {
                if name == iface.name {
                    merge(iface, &ni);
                    found = true;
                }
            }
            if !found {
                ifaces.push(ni);
            }

            addr = (*addr).ifa_next;
        }
        libc::freeifaddrs(addrs);

        for iface in &mut ifaces {
            let name = CString::new(iface.name.as_bytes());
            iface.index = libc::if_nametoindex(name.unwrap().as_ptr());
        }

        ifaces
    }
}

#[cfg(any(target_os = "linux", target_os = "android"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr(
                (*sll).sll_addr[0],
                (*sll).sll_addr[1],
                (*sll).sll_addr[2],
                (*sll).sll_addr[3],
                (*sll).sll_addr[4],
                (*sll).sll_addr[5],
            );

            (Some(mac), None)
        } else {
            let addr = pnet_sys::sockaddr_to_addr(
                mem::transmute(sa),
                mem::size_of::<libc::sockaddr_storage>(),
            );

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

#[cfg(any(target_os = "openbsd", target_os = "freebsd", target_os = "netbsd", target_os = "macos", target_os = "ios"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as usize;
            let mac = MacAddr(
                (*sdl).sdl_data[nlen] as u8,
                (*sdl).sdl_data[nlen + 1] as u8,
                (*sdl).sdl_data[nlen + 2] as u8,
                (*sdl).sdl_data[nlen + 3] as u8,
                (*sdl).sdl_data[nlen + 4] as u8,
                (*sdl).sdl_data[nlen + 5] as u8,
            );

            (Some(mac), None)
        } else {
            let addr = pnet_sys::sockaddr_to_addr(
                mem::transmute(sa),
                mem::size_of::<libc::sockaddr_storage>(),
            );

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}
<commit_msg>Fix invalid pointer use in unix interfaces<commit_after>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interface listing implementation for all non-Windows platforms.

use {MacAddr, NetworkInterface};

use ipnetwork::{ip_mask_to_prefix, IpNetwork};
use pnet_sys;

use libc;

use std::ffi::{CStr, CString};
use std::mem;
use std::net::IpAddr;
use std::os::raw::c_char;
use std::str::from_utf8_unchecked;

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    fn merge(old: &mut NetworkInterface, new: &NetworkInterface) {
        old.mac = match new.mac {
            None => old.mac,
            _ => new.mac,
        };
        old.ips.extend_from_slice(&new.ips[..]);
        old.flags = old.flags | new.flags;
    }

    let mut ifaces: Vec<NetworkInterface> = Vec::new();
    unsafe {
        #[allow(deprecated)]
        let mut addrs: *mut libc::ifaddrs = mem::uninitialized();
        if libc::getifaddrs(&mut addrs) != 0 {
            return ifaces;
        }
        let mut addr = addrs;
        while !addr.is_null() {
            let c_str = (*addr).ifa_name as *const c_char;
            let bytes = CStr::from_ptr(c_str).to_bytes();
            let name = from_utf8_unchecked(bytes).to_owned();
            let (mac, ip) = sockaddr_to_network_addr((*addr).ifa_addr as *const libc::sockaddr);
            let (_, netmask) =
                sockaddr_to_network_addr((*addr).ifa_netmask as *const libc::sockaddr);
            let prefix = netmask
                .and_then(|netmask| ip_mask_to_prefix(netmask).ok())
                .unwrap_or(0);
            let network = ip.and_then(|ip| IpNetwork::new(ip, prefix).ok());
            let ni = NetworkInterface {
                name: name.clone(),
                description: "".to_string(),
                index: 0,
                mac: mac,
                ips: network.into_iter().collect(),
                flags: (*addr).ifa_flags,
            };
            let mut found: bool = false;
            for iface in &mut ifaces {
                if name == iface.name {
                    merge(iface, &ni);
                    found = true;
                }
            }
            if !found {
                ifaces.push(ni);
            }

            addr = (*addr).ifa_next;
        }
        libc::freeifaddrs(addrs);

        for iface in &mut ifaces {
            let name = CString::new(iface.name.as_bytes()).unwrap();
            iface.index = libc::if_nametoindex(name.as_ptr());
        }

        ifaces
    }
}

#[cfg(any(target_os = "linux", target_os = "android"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == libc::AF_PACKET {
            let sll: *const libc::sockaddr_ll = mem::transmute(sa);
            let mac = MacAddr(
                (*sll).sll_addr[0],
                (*sll).sll_addr[1],
                (*sll).sll_addr[2],
                (*sll).sll_addr[3],
                (*sll).sll_addr[4],
                (*sll).sll_addr[5],
            );

            (Some(mac), None)
        } else {
            let addr = pnet_sys::sockaddr_to_addr(
                mem::transmute(sa),
                mem::size_of::<libc::sockaddr_storage>(),
            );

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}

#[cfg(any(target_os = "openbsd", target_os = "freebsd", target_os = "netbsd", target_os = "macos", target_os = "ios"))]
fn sockaddr_to_network_addr(sa: *const libc::sockaddr) -> (Option<MacAddr>, Option<IpAddr>) {
    use bindings::bpf;
    use std::net::SocketAddr;

    unsafe {
        if sa.is_null() {
            (None, None)
        } else if (*sa).sa_family as libc::c_int == bpf::AF_LINK {
            let sdl: *const bpf::sockaddr_dl = mem::transmute(sa);
            let nlen = (*sdl).sdl_nlen as usize;
            let mac = MacAddr(
                (*sdl).sdl_data[nlen] as u8,
                (*sdl).sdl_data[nlen + 1] as u8,
                (*sdl).sdl_data[nlen + 2] as u8,
                (*sdl).sdl_data[nlen + 3] as u8,
                (*sdl).sdl_data[nlen + 4] as u8,
                (*sdl).sdl_data[nlen + 5] as u8,
            );

            (Some(mac), None)
        } else {
            let addr = pnet_sys::sockaddr_to_addr(
                mem::transmute(sa),
                mem::size_of::<libc::sockaddr_storage>(),
            );

            match addr {
                Ok(SocketAddr::V4(sa)) => (None, Some(IpAddr::V4(*sa.ip()))),
                Ok(SocketAddr::V6(sa)) => (None, Some(IpAddr::V6(*sa.ip()))),
                Err(_) => (None, None),
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate serde_json;
extern crate serde;
extern crate regex;
extern crate docopt;
extern crate rustc_serialize;
extern crate env_logger;
extern crate ansi_term;
#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate log;
#[macro_use]
extern crate sonicd;
extern crate rpassword;
extern crate pbr;

mod util {
    #[cfg(feature = "serde_macros")]
    include!("util.rs.in");

    #[cfg(not(feature = "serde_macros"))]
    include!(concat!(env!("OUT_DIR"), "/util.rs"));
}

use std::path::PathBuf;
use std::process;
use std::io::{Write, stderr, stdout};
use std::cell::RefCell;

use docopt::Docopt;
use pbr::ProgressBar;
use sonicd::{Authenticate, SonicMessage};
use rpassword::read_password;

#[cfg_attr(rustfmt, rustfmt_skip)]
const USAGE: &'static str = "
.
                           d8b
                           Y8P

.d8888b   .d88b.  88888b.  888  .d8888b
88K      d88\"\"88b 888 \"88b 888 d88P\"
\"Y8888b. 888  888 888  888 888 888
     X88 Y88..88P 888  888 888 Y88b.
 88888P'  \"Y88P\"  888  888 888  \"Y8888P

Usage:
  sonic <source> [-d <foo=bar>...] [options] -e <query>
  sonic <source> [-d <foo=bar>...] [options] -f <file>
  sonic login [options]
  sonic -h | --help
  sonic --version

Options:
  -e, --execute         Run command literal
  -f, --file            Run command from file
  -c <config>           Use a different configuration file (default: $HOME/.sonicrc)
  -d <foo=bar>          Replace variable in query in the form of `${foo}` with value `var`
  -r, --rows-only       Skip printing column names
  -S, --silent          Skip printing query progress bar
  -v, --verbose         Enable debug logging
  -h, --help            Print this message
  --version             Print version
";

static VERSION: &'static str = env!("CARGO_PKG_VERSION");
static COMMIT: Option<&'static str> = option_env!("SONIC_COMMIT");

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_file: String,
    arg_query: String,
    flag_silent: bool,
    flag_rows_only: bool,
    flag_d: Vec<String>,
    flag_file: bool,
    flag_c: String,
    flag_verbose: bool,
    arg_source: String,
    cmd_login: bool,
    flag_execute: bool,
    flag_version: bool,
    flag_help: bool,
}

error_chain! {
    types {
        Error, ErrorKind, ChainErr, Result;
    }

    links {
        sonicd::Error, sonicd::ErrorKind, SonicdError;
    }

    foreign_links {
        ::std::io::Error, IoError;
        ::serde_json::Error, Json;
        ::std::sync::mpsc::RecvError, RecvError;
    }

    errors {
        InvalidFormat(co: String, msg: String) {
            description("invalid format error")
                display("invalid format {}: {}", co, msg)
        }
    }
}

fn hide<T: Write>(pb: &mut ProgressBar<T>) {
    pb.show_bar = false;
    pb.show_speed = false;
    pb.show_percent = false;
    pb.show_counter = false;
    pb.show_time_left = false;
    pb.show_tick = false;
    pb.show_message = false;
}

fn show<T: Write>(pb: &mut ProgressBar<T>) {
    pb.show_bar = true;
    pb.show_speed = false;
    pb.show_percent = true;
    pb.show_counter = true;
    pb.show_time_left = true;
    pb.show_tick = true;
    pb.show_message = true;
}

fn exec(
    host: &str,
    port: &u16,
    query: SonicMessage,
    rows_only: bool,
    silent: bool
) -> Result<()> {

    let out = RefCell::new(stdout());
    let mut buf: Vec<SonicMessage> = Vec::new();
    let mut pb = ProgressBar::on(stderr(), 0);
    pb.format("╢░░_╟");
    pb.tick_format("▀▐▄▌");
    show(&mut pb);

    let (tx, rx) = ::std::sync::mpsc::channel();

    try!(sonicd::stream((host, *port), query, tx));

    let draw = |b: &mut Vec<SonicMessage>| {
        let len = b.len();
        for msg in b.drain(..len) {
            let cols = match msg {
                SonicMessage::OutputChunk(data) => {
                    data.iter().fold(String::new(),
                                     |acc, val| format!("{}{:?}\t", acc, val))
                }
                SonicMessage::TypeMetadata(data) => {
                    data.iter().fold(String::new(),
                                     |acc, col| format!("{}{:?}\t", acc, col.0))
                }
                _ => panic!("not possible!"),
            };
            let row = format!("{}\n", cols.trim_right());
            let mut bout = out.borrow_mut();
            bout.write_all(row.as_bytes()).unwrap();
            bout.flush().unwrap();
        }
    };

    let res: Result<()>;

    loop {
        match try!(rx.recv()) {
            Ok(SonicMessage::QueryStarted(trace_id)) => {
                debug!("started query with trace_id: {}", trace_id);
            }
            Ok(msg @ SonicMessage::TypeMetadata(_)) => {
                if !rows_only {
                    buf.push(msg);
                }
            }
            Ok(msg @ SonicMessage::OutputChunk(_)) => {
                buf.push(msg);
                if !silent && !pb.is_finish {
                    // tick format breaks suffix length
                    // so we need to disable it before finishing bar
                    hide(&mut pb);
                    pb.tick();
                    pb.finish();
                }
                draw(&mut buf);
            }
            Ok(SonicMessage::QueryProgress { status, progress, total, .. }) => {
                if !silent && !pb.is_finish {

                    pb.message(&format!("{:?} ", status));
                    debug!("{:?}: {:?}/{:?}", status, progress, total);

                    if let Some(total) = total {
                        pb.total = total as u64;
                    }

                    if progress >= 1.0 {
                        pb.add(progress as u64);
                    } else {
                        pb.tick();
                    };
                }
            }
            Ok(SonicMessage::Done(None, trace_id)) => {
                debug!("stream '{}' completed successfully", trace_id);
                res = Ok(());
                break;
            }
            Ok(SonicMessage::Done(Some(cause), trace_id)) => {
                debug!("stream '{}' failed with error {:?}", &cause, trace_id);
                let error: Error = cause.into();
                res = Err(error).chain_err(|| {
                    format!("error when running query; trace_id: {}", trace_id)
                });
                break;
            }
            Err(e) => {
                let error: Error = e.into();
                res = Err(error).chain_err(|| "unexpected channel recv error");
                break;
            }
            Ok(a) => debug!("ignoring msg {:?}", a),
        }
    }

    draw(&mut buf);

    res
}

pub fn login(host: &str, tcp_port: &u16) -> Result<()> {

    let user = std::env::var("USER").unwrap_or("unknown".to_owned());

    try!(stdout().write(b"Enter key: "));
    try!(stdout().flush());

    let key = try!(read_password());

    let (tx, rx) = ::std::sync::mpsc::channel();

    let cmd = SonicMessage::AuthenticateMsg(Authenticate::new(user.to_owned(),
                                                              key,
                                                              None));

    try!(sonicd::stream((host, *tcp_port), cmd, tx));

    let token: String;

    loop {
        match try!(rx.recv()) {
            Ok(SonicMessage::OutputChunk(data)) => {
                token = try!(util::parse_token(data));
                break;
            }
            Ok(SonicMessage::Done(None, trace_id)) => {
                return Err(format!("protocol error: server returned no data \
                                    for '{}'",
                                   trace_id)
                    .into());
            }
            Ok(SonicMessage::Done(Some(cause), trace_id)) => {
                let error: Error = cause.into();
                return Err(error).chain_err(|| {
                    format!("error when running login command; trace_id: {}",
                            trace_id)
                });
            }
            Ok(_) => {}
            Err(e) => return Err(e.into()),
        };
    }

    debug!("generated token for user {:?}: {:?}", &user, &token);

    let path = util::get_config_path();
    let config = try!(util::read_config(&path));

    let new_config = util::ClientConfig {
        auth: Some(token),
        ..config
    };

    try!(util::write_config(&new_config, &path));

    println!("OK");
    Ok(())
}

fn _main(args: Args) -> Result<()> {

    let Args { arg_file,
               arg_query,
               flag_silent,
               flag_rows_only,
               flag_d,
               flag_file,
               flag_c,
               arg_source,
               cmd_login,
               flag_execute,
               flag_version,
               .. } = args;

    let util::ClientConfig { sonicd, tcp_port, sources, auth } = if flag_c !=
                                                                    "" {
        debug!("sourcing passed config in path '{:?}'", &flag_c);
        try!(util::read_config(&PathBuf::from(flag_c)))
    } else {
        debug!("sourcing default config in path '$HOME/.sonicrc'");
        try!(util::get_default_config())
    };

    if flag_file {

        let query_str =
            try!(util::read_file_contents(&PathBuf::from(&arg_file)));
        let split = try!(util::split_key_value(&flag_d));
        let injected = try!(util::inject_vars(&query_str, &split));
        let query = try!(util::build(arg_source, sources, auth, injected));

        exec(&sonicd, &tcp_port, query, flag_rows_only, flag_silent)

    } else if flag_execute {

        let query = try!(util::build(arg_source, sources, auth, arg_query));

        exec(&sonicd, &tcp_port, query, flag_rows_only, flag_silent)

    } else if cmd_login {

        login(&sonicd, &tcp_port)

    } else if flag_version {

        println!("Sonic CLI version {} ({})",
                 VERSION,
                 COMMIT.unwrap_or_else(|| "dev"));

        Ok(())

    } else {

        panic!("unexpected args");

    }
}

fn main() {

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let verbose: bool = args.flag_verbose;

    if verbose {
        let mut builder = env_logger::LogBuilder::new();
        builder.parse("debug");
        builder.init().unwrap();
    } else {
        env_logger::init().unwrap();
    }

    debug!("parsed args {:?}", args);

    match _main(args) {
        Ok(_) => {}
        Err(e) => {
            util::report_error(&e, verbose);
            process::exit(1)
        }
    }
}
<commit_msg>improved CLI help messages<commit_after>extern crate serde_json;
extern crate serde;
extern crate regex;
extern crate docopt;
extern crate rustc_serialize;
extern crate env_logger;
extern crate ansi_term;
#[macro_use]
extern crate error_chain;
#[macro_use]
extern crate log;
#[macro_use]
extern crate sonicd;
extern crate rpassword;
extern crate pbr;

mod util {
    #[cfg(feature = "serde_macros")]
    include!("util.rs.in");

    #[cfg(not(feature = "serde_macros"))]
    include!(concat!(env!("OUT_DIR"), "/util.rs"));
}

use std::path::PathBuf;
use std::process;
use std::io::{Write, stderr, stdout};
use std::cell::RefCell;

use docopt::Docopt;
use pbr::ProgressBar;
use sonicd::{Authenticate, SonicMessage};
use rpassword::read_password;

#[cfg_attr(rustfmt, rustfmt_skip)]
const USAGE: &'static str = "
.
                           d8b
                           Y8P

.d8888b   .d88b.  88888b.  888  .d8888b
88K      d88\"\"88b 888 \"88b 888 d88P\"
\"Y8888b. 888  888 888  888 888 888
     X88 Y88..88P 888  888 888 Y88b.
 88888P'  \"Y88P\"  888  888 888  \"Y8888P

Usage:
  sonic <source> [options] -e <query>
  sonic <source> [options] -f <file>
  sonic login [options]
  sonic -h | --help
  sonic --version

Options:
  -e, --execute         Run query literal
  -f, --file            Run file contents as query
  -c <config>           Override default configuration file ($HOME/.sonicrc)
  -d <foo=var> ...      Replace variable in query in the form of `${foo}` with value `var`
  -r, --rows-only       Skip printing column names
  -S, --silent          Disable progress bar
  -v, --verbose         Enable debug level logging
  -h, --help            Print this message
  --version             Print version
";

static VERSION: &'static str = env!("CARGO_PKG_VERSION");
static COMMIT: Option<&'static str> = option_env!("SONIC_COMMIT");

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_file: String,
    arg_query: String,
    flag_silent: bool,
    flag_rows_only: bool,
    flag_d: Vec<String>,
    flag_file: bool,
    flag_c: String,
    flag_verbose: bool,
    arg_source: String,
    cmd_login: bool,
    flag_execute: bool,
    flag_version: bool,
    flag_help: bool,
}

error_chain! {
    types {
        Error, ErrorKind, ChainErr, Result;
    }

    links {
        sonicd::Error, sonicd::ErrorKind, SonicdError;
    }

    foreign_links {
        ::std::io::Error, IoError;
        ::serde_json::Error, Json;
        ::std::sync::mpsc::RecvError, RecvError;
    }

    errors {
        InvalidFormat(co: String, msg: String) {
            description("invalid format error")
                display("invalid format {}: {}", co, msg)
        }
    }
}

fn hide<T: Write>(pb: &mut ProgressBar<T>) {
    pb.show_bar = false;
    pb.show_speed = false;
    pb.show_percent = false;
    pb.show_counter = false;
    pb.show_time_left = false;
    pb.show_tick = false;
    pb.show_message = false;
}

fn show<T: Write>(pb: &mut ProgressBar<T>) {
    pb.show_bar = true;
    pb.show_speed = false;
    pb.show_percent = true;
    pb.show_counter = true;
    pb.show_time_left = true;
    pb.show_tick = true;
    pb.show_message = true;
}

fn exec(
    host: &str,
    port: &u16,
    query: SonicMessage,
    rows_only: bool,
    silent: bool
) -> Result<()> {

    let out = RefCell::new(stdout());
    let mut buf: Vec<SonicMessage> = Vec::new();
    let mut pb = ProgressBar::on(stderr(), 0);
    pb.format("╢░░_╟");
    pb.tick_format("▀▐▄▌");
    show(&mut pb);

    let (tx, rx) = ::std::sync::mpsc::channel();

    try!(sonicd::stream((host, *port), query, tx));

    let draw = |b: &mut Vec<SonicMessage>| {
        let len = b.len();
        for msg in b.drain(..len) {
            let cols = match msg {
                SonicMessage::OutputChunk(data) => {
                    data.iter().fold(String::new(),
                                     |acc, val| format!("{}{:?}\t", acc, val))
                }
                SonicMessage::TypeMetadata(data) => {
                    data.iter().fold(String::new(),
                                     |acc, col| format!("{}{:?}\t", acc, col.0))
                }
                _ => panic!("not possible!"),
            };
            let row = format!("{}\n", cols.trim_right());
            let mut bout = out.borrow_mut();
            bout.write_all(row.as_bytes()).unwrap();
            bout.flush().unwrap();
        }
    };

    let res: Result<()>;

    loop {
        match try!(rx.recv()) {
            Ok(SonicMessage::QueryStarted(trace_id)) => {
                debug!("started query with trace_id: {}", trace_id);
            }
            Ok(msg @ SonicMessage::TypeMetadata(_)) => {
                if !rows_only {
                    buf.push(msg);
                }
            }
            Ok(msg @ SonicMessage::OutputChunk(_)) => {
                buf.push(msg);
                if !silent && !pb.is_finish {
                    // tick format breaks suffix length
                    // so we need to disable it before finishing bar
                    hide(&mut pb);
                    pb.tick();
                    pb.finish();
                }
                draw(&mut buf);
            }
            Ok(SonicMessage::QueryProgress { status, progress, total, .. }) => {
                if !silent && !pb.is_finish {

                    pb.message(&format!("{:?} ", status));
                    debug!("{:?}: {:?}/{:?}", status, progress, total);

                    if let Some(total) = total {
                        pb.total = total as u64;
                    }

                    if progress >= 1.0 {
                        pb.add(progress as u64);
                    } else {
                        pb.tick();
                    };
                }
            }
            Ok(SonicMessage::Done(None, trace_id)) => {
                debug!("stream '{}' completed successfully", trace_id);
                res = Ok(());
                break;
            }
            Ok(SonicMessage::Done(Some(cause), trace_id)) => {
                debug!("stream '{}' failed with error {:?}", &cause, trace_id);
                let error: Error = cause.into();
                res = Err(error).chain_err(|| {
                    format!("error when running query; trace_id: {}", trace_id)
                });
                break;
            }
            Err(e) => {
                let error: Error = e.into();
                res = Err(error).chain_err(|| "unexpected channel recv error");
                break;
            }
            Ok(a) => debug!("ignoring msg {:?}", a),
        }
    }

    draw(&mut buf);

    res
}

pub fn login(host: &str, tcp_port: &u16) -> Result<()> {

    let user = std::env::var("USER").unwrap_or("unknown".to_owned());

    try!(stdout().write(b"Enter key: "));
    try!(stdout().flush());

    let key = try!(read_password());

    let (tx, rx) = ::std::sync::mpsc::channel();

    let cmd = SonicMessage::AuthenticateMsg(Authenticate::new(user.to_owned(),
                                                              key,
                                                              None));

    try!(sonicd::stream((host, *tcp_port), cmd, tx));

    let token: String;

    loop {
        match try!(rx.recv()) {
            Ok(SonicMessage::OutputChunk(data)) => {
                token = try!(util::parse_token(data));
                break;
            }
            Ok(SonicMessage::Done(None, trace_id)) => {
                return Err(format!("protocol error: server returned no data \
                                    for '{}'",
                                   trace_id)
                    .into());
            }
            Ok(SonicMessage::Done(Some(cause), trace_id)) => {
                let error: Error = cause.into();
                return Err(error).chain_err(|| {
                    format!("error when running login command; trace_id: {}",
                            trace_id)
                });
            }
            Ok(_) => {}
            Err(e) => return Err(e.into()),
        };
    }

    debug!("generated token for user {:?}: {:?}", &user, &token);

    let path = util::get_config_path();
    let config = try!(util::read_config(&path));

    let new_config = util::ClientConfig {
        auth: Some(token),
        ..config
    };

    try!(util::write_config(&new_config, &path));

    println!("OK");
    Ok(())
}

fn _main(args: Args) -> Result<()> {

    let Args { arg_file,
               arg_query,
               flag_silent,
               flag_rows_only,
               flag_d,
               flag_file,
               flag_c,
               arg_source,
               cmd_login,
               flag_execute,
               flag_version,
               .. } = args;

    let util::ClientConfig { sonicd, tcp_port, sources, auth } = if flag_c !=
                                                                    "" {
        debug!("sourcing passed config in path '{:?}'", &flag_c);
        try!(util::read_config(&PathBuf::from(flag_c)))
    } else {
        debug!("sourcing default config in path '$HOME/.sonicrc'");
        try!(util::get_default_config())
    };

    if flag_file {

        let query_str =
            try!(util::read_file_contents(&PathBuf::from(&arg_file)));
        let split = try!(util::split_key_value(&flag_d));
        let injected = try!(util::inject_vars(&query_str, &split));
        let query = try!(util::build(arg_source, sources, auth, injected));

        exec(&sonicd, &tcp_port, query, flag_rows_only, flag_silent)

    } else if flag_execute {

        let query = try!(util::build(arg_source, sources, auth, arg_query));

        exec(&sonicd, &tcp_port, query, flag_rows_only, flag_silent)

    } else if cmd_login {

        login(&sonicd, &tcp_port)

    } else if flag_version {

        println!("Sonic CLI version {} ({})",
                 VERSION,
                 COMMIT.unwrap_or_else(|| "dev"));

        Ok(())

    } else {

        panic!("unexpected args");

    }
}

fn main() {

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    let verbose: bool = args.flag_verbose;

    if verbose {
        let mut builder = env_logger::LogBuilder::new();
        builder.parse("debug");
        builder.init().unwrap();
    } else {
        env_logger::init().unwrap();
    }

    debug!("parsed args {:?}", args);

    match _main(args) {
        Ok(_) => {}
        Err(e) => {
            util::report_error(&e, verbose);
            process::exit(1)
        }
    }
}
<|endoftext|>