<commit_before>
#![allow( non_upper_case_globals )]

use languageserver_types::*;
use serde;
use serde_json;
use std::{
    io,
    mem,
    str
};
use std::collections::{
    HashMap
};
use tokio_core::io::{
    Codec,
    EasyBuf
};

macro_rules! server_messages {
    (
        methods {
            $( $method_name : ident => $method_type : ident $( ( $method_param : ty ) )*; )*
        },
        notifications {
            $( $notif_name : ident => $notif_type : ident $( ( $notif_param : ty ) )*; )*
        }
    ) => {
        #[derive( Debug )]
        pub enum RequestMethod {
            $( $method_type $( ( $method_param ) )*, )*
        }

        #[derive( Debug )]
        pub enum NotificationMethod {
            $( $notif_type $( ( $notif_param ) )*, )*
        }

        pub fn parse_server_message( msg_json_str : &str ) -> io::Result< ServerMessage > {
            let msg_json : serde_json::Value = match serde_json::from_str( msg_json_str ) {
                Ok( msg_json ) => msg_json,
                Err( _ ) => return Err( new_invalid_data_error( "Invalid JSON data." ) )
            };

            if let Some( method_val ) = msg_json.lookup( "method" ) {
                let method = match method_val.as_str( ) {
                    Some( method ) => method,
                    None => return Err( new_invalid_data_error( "Invalid 'method' field, could not parse as String." ) )
                };
                match method {
                    $(
                        $method_name => {
                            let msg_id = match msg_json.lookup( "id" ) {
                                Some( msg_id ) => {
                                    match msg_id.as_i64( ) {
                                        Some( msg_id ) => msg_id,
                                        None => return Err( new_invalid_data_error( "Invalid 'id' field, could not parse as number." ) )
                                    }
                                },
                                None => return Err( new_invalid_data_error( "Missing required 'id' field." ) )
                            };
                            let msg = RequestMessage {
                                id     : msg_id,
                                method : RequestMethod::$method_type $( ( parse_parameters::< $method_param >( &msg_json )? ) )*
                            };
                            Ok( ServerMessage::Request( msg ) )
                        }
                    ),*
                    $(
                        $notif_name => {
                            let notif = NotificationMessage {
                                method : NotificationMethod::$notif_type $( ( parse_parameters::< $notif_param >( &msg_json )? ) )*
                            };
                            Ok( ServerMessage::Notification( notif ) )
                        }
                    ),*
                    _ => Err( new_invalid_data_error( "Invalid request method." ) )
                }
            }
            else {
                return Err( new_invalid_data_error( "Missing required 'method' field." ) );
            }
        }
    };
}

macro_rules! server_responses {
    (
        datas {
            $( $data_type : ident( $data_param : ty ); )*
        },
        errors {
            $( $err_name : ident( $err_val : expr ); )*
        }
    ) => {
        #[derive( Debug )]
        pub enum RequestResponseData {
            $( $data_type ( $data_param ) ),*
        }

        $( pub static $err_name : i64 = $err_val; )*

        impl serde::Serialize for RequestResponseData {

            fn serialize< S >( &self, serializer : &mut S ) -> Result< ( ), S::Error > where S : serde::Serializer {
                match *self {
                    $(
                        RequestResponseData::$data_type( ref param ) => {
                            param.serialize( serializer )
                        }
                    ),*
                }
            }

        }
    }
}

#[derive( Debug )]
pub struct IncomingServerMessage {
    headers : HashMap< String, String >,
    message : ServerMessage
}

#[derive( Debug )]
pub struct OutgoingServerMessage {
    headers : HashMap< String, String >,
    message : ClientMessage
}

#[derive( Debug )]
pub enum ServerMessage {
    Request( RequestMessage ),
    Notification( NotificationMessage )
}

#[derive( Debug )]
pub enum ClientMessage {
    Response( RequestResponseMessage )
}

#[derive( Debug )]
pub struct RequestMessage {
    pub id     : i64,
    pub method : RequestMethod
}

#[derive( Debug )]
pub struct NotificationMessage {
    pub method : NotificationMethod
}

#[derive( Debug, Serialize )]
pub struct RequestResponseError {
    pub code    : i64,
    pub message : String
}

#[derive( Debug )]
pub struct RequestResponseMessage {
    id     : i64,
    result : Option< RequestResponseData >,
    error  : Option< RequestResponseError >
}

server_messages! {
    methods {
        REQUEST__Initialize                        => Initialize( InitializeParams );
        REQUEST__Shutdown                          => Shutdown;
        REQUEST__Completion                        => Completion( TextDocumentPositionParams );
        REQUEST__ResolveCompletionItem             => CompletionResolve( CompletionItem );
        REQUEST__Hover                             => Hover( TextDocumentPositionParams );
        REQUEST__SignatureHelp                     => SignatureHelp( TextDocumentPositionParams );
        REQUEST__GotoDefinition                    => GotoDefinition( TextDocumentPositionParams );
        REQUEST__References                        => FindReferences( ReferenceParams );
        REQUEST__DocumentHighlight                 => DocumentHighlights( TextDocumentPositionParams );
        REQUEST__DocumentSymbols                   => DocumentSymbols( DocumentSymbolParams );
        REQUEST__WorkspaceSymbols                  => WorkspaceSymbols( WorkspaceSymbolParams );
        REQUEST__CodeAction                        => CodeAction( CodeActionParams );
        REQUEST__CodeLens                          => CodeLens( CodeLensParams );
        REQUEST__CodeLensResolve                   => CodeLensResolve( CodeLens );
        REQUEST__DocumentLink                      => DocumentLink( DocumentLinkParams );
        REQUEST__DocumentLinkResolve               => DocumentLinkResolve( DocumentLink );
        REQUEST__Formatting                        => DocumentFormatting( DocumentFormattingParams );
        REQUEST__RangeFormatting                   => DocumentRangeFormatting( DocumentRangeFormattingParams );
        REQUEST__OnTypeFormatting                  => DocumentOnTypeFormatting( DocumentOnTypeFormattingParams );
        REQUEST__Rename                            => Rename( RenameParams );
    },
    notifications {
        NOTIFICATION__Cancel                       => Cancel( CancelParams );
        NOTIFICATION__Exit                         => Exit;
        NOTIFICATION__WorkspaceChangeConfiguration => DidChangeConfiguration( DidChangeConfigurationParams );
        NOTIFICATION__DidOpenTextDocument          => DidOpenTextDocument( DidOpenTextDocumentParams );
        NOTIFICATION__DidChangeTextDocument        => DidChangeTextDocument( DidChangeTextDocumentParams );
        NOTIFICATION__DidCloseTextDocument         => DidCloseTextDocument( DidCloseTextDocumentParams );
        NOTIFICATION__DidSaveTextDocument          => DidSaveTextDocument( DidSaveTextDocumentParams );
        NOTIFICATION__DidChangeWatchedFiles        => DidChangeWatchedFiles( DidChangeWatchedFilesParams );
    }
}

server_responses! {
    datas {
        Init( InitializeResult );
        SymbolInfo( Vec< SymbolInformation > );
        CompletionItems( Vec< CompletionItem > );
        WorkspaceEdit( WorkspaceEdit );
        TextEdit( [ TextEdit; 1 ] );
        Locations( Vec< Location > );
        Highlights( Vec< DocumentHighlight > );
        HoverSuccess( Hover );
    },
    errors {
        PARSE_ERROR( -32700 );
        INVALID_REQUEST( -32600 );
        METHOD_NOT_FOUND( -32601 );
        INVALID_PARAMS( -32602 );
        INTERNAL_ERROR( -32603 );
        SERVER_ERROR_START( -32099 );
        SERVER_ERROR_END( -32000 );
        SERVER_NOT_INITIALIZED( -32002 );
        UNKNOWN_ERROR_CODE( -32001 );
    }
}

fn new_invalid_data_error( error_msg : &'static str ) -> io::Error {
    io::Error::new( io::ErrorKind::InvalidData, error_msg )
}

fn parse_parameters< T >( msg_json : &serde_json::Value ) -> io::Result< T > where T : serde::Deserialize {
    let params = match msg_json.lookup( "params" ) {
        Some( params ) => params.to_owned( ),
        None => return Err( new_invalid_data_error( "Missing required 'params' field." ) )
    };

    match serde_json::from_value( params ) {
        Ok( value ) => Ok( value ),
        Err( _ ) => Err( new_invalid_data_error( "Unable to parse parameter field." ) )
    }
}

fn easybuf_to_utf8( buf : &EasyBuf ) -> io::Result< &str > {
    match str::from_utf8( buf.as_ref( ) ) {
        Ok( s ) => Ok( s ),
        Err( _ ) => Err( new_invalid_data_error( "Unable to parse bytes as UTF-8." ) )
    }
}

enum RLSCodecState {
    Headers,
    Body
}

pub struct RLSCodec {
    codec_state    : RLSCodecState,
    headers        : HashMap< String, String >,
    content_length : usize
}

impl RLSCodec {

    pub fn new( ) -> Self {
        RLSCodec {
            codec_state    : RLSCodecState::Headers,
            headers        : HashMap::new( ),
            content_length : 0
        }
    }

    fn read_header( &mut self, buf : &mut EasyBuf ) -> io::Result< bool > {
        if let Some( pos ) = buf.as_ref( ).windows( 2 ).position( | bytes | bytes == b"\r\n" ) {
            let header_bytes = buf.drain_to( pos );
            buf.drain_to( 2 );

            if pos == 0 {
                self.content_length = self.parse_content_length( )?;
                self.codec_state    = RLSCodecState::Body;

                trace!( "Finished parsing headers. Waiting for '{}' bytes for body.", self.content_length );
            }
            else {
                let header_str = easybuf_to_utf8( &header_bytes )?;
                trace!( "Received header line: {}", header_str );

                let mut header_parts = header_str.split( ": " );

                if header_parts.clone( ).count( ) != 2 {
                    return Err( new_invalid_data_error( "Unable to parse invalid header." ) );
                }

                let header_key = header_parts.next( ).unwrap( );
                let header_val = header_parts.next( ).unwrap( );

                self.headers.insert( header_key.to_string( ), header_val.to_string( ) );
            }
            Ok( true )
        }
        else {
            Ok( false )
        }
    }

    fn parse_content_length( &self ) -> io::Result< usize > {
        if let Some( content_length_str ) = self.headers.get( "Content-Length" ) {
            if let Ok( content_length ) = content_length_str.parse::< usize >( ) {
                Ok( content_length )
            }
            else {
                Err( new_invalid_data_error( "Unable to parse 'Content-Length' header into usize." ) )
            }
        }
        else {
            Err( new_invalid_data_error( "Missing required 'Content-Length' header." ) )
        }
    }

    fn reset_codec( &mut self ) -> HashMap< String, String > {
        let mut new_headers = HashMap::new( );
        mem::swap( &mut self.headers, &mut new_headers );

        self.codec_state = RLSCodecState::Headers;

        new_headers
    }

    fn parse_body_contents( &mut self, buf : &mut EasyBuf ) -> io::Result< Option< IncomingServerMessage > > {
        if buf.len( ) < self.content_length {
            return Ok( None );
        }

        let body_bytes = buf.drain_to( self.content_length );
        let body_str   = easybuf_to_utf8( &body_bytes )?;
        trace!( "Received body contents: {}", body_str );

        let message = parse_server_message( body_str )?;
        trace!( "Parsed server message: {:?}", message );

        let headers = self.reset_codec( );

        Ok( Some( IncomingServerMessage{
            headers : headers,
            message : message
        } ) )
    }

}

impl Codec for RLSCodec {

    type In  = IncomingServerMessage;
    type Out = OutgoingServerMessage;

    fn decode( &mut self, buf : &mut EasyBuf ) -> io::Result< Option< Self::In > > {
        loop {
            match self.codec_state {
                RLSCodecState::Headers => {
                    if !self.read_header( buf )? {
                        return Ok( None );
                    }
                },
                RLSCodecState::Body => {
                    return self.parse_body_contents( buf );
                }
            }
        }
    }

    fn encode( &mut self, msg : OutgoingServerMessage, buffer : &mut Vec< u8 > ) -> io::Result< ( ) > {
        let OutgoingServerMessage{ mut headers, message } = msg;

        let msg_json = message.into_json( );
        headers.insert( "Content-Length".to_string( ), msg_json.as_bytes( ).len( ).to_string( ) );
        headers.insert( "Content-Type".to_string( ), "application/vscode-jsonrpc; charset=utf8".to_string( ) );

        for ( ref key, ref val ) in headers {
            trace!( "Sending header: {} = {}", key, val );

            buffer.extend_from_slice( key.as_bytes( ) );
            buffer.extend_from_slice( b": " );
            buffer.extend_from_slice( val.as_bytes( ) );
            buffer.extend_from_slice( b"\r\n" );
        }
        buffer.extend_from_slice( b"\r\n" );

        trace!( "Sending body: {}", msg_json );
        buffer.extend_from_slice( msg_json.as_bytes( ) );

        Ok( ( ) )
    }
}

impl ClientMessage {

    fn into_json( self ) -> String {
        match self {
            ClientMessage::Response( value ) => {

                #[derive( Serialize )]
                struct ResponseMessage {
                    jsonrpc : &'static str,
                    id      : i64,
                    result  : Option< RequestResponseData >,
                    error   : Option< RequestResponseError >
                }

                serde_json::to_string( &ResponseMessage {
                    jsonrpc : "2.0",
                    id      : value.id,
                    result  : value.result,
                    error   : value.error
                } ).unwrap( )
            }
        }
    }

}<commit_msg>Add some documentation and missing pub members.<commit_after>//! Implements encoding/decoding support and types for the server side of the Language Server protocol.
//!
//! Provides an implementation of tokio_core::io::Codec to decode IncomingServerMessages and encode
//! OutgoingServerMessages

#![allow( non_upper_case_globals )]

use languageserver_types::*;
use serde;
use serde_json;
use std::{
    io,
    mem,
    str
};
use std::collections::{
    HashMap
};
use tokio_core::io::{
    Codec,
    EasyBuf
};

macro_rules! server_messages {
    (
        methods {
            $( $method_name : ident => $method_type : ident $( ( $method_param : ty ) )*; )*
        },
        notifications {
            $( $notif_name : ident => $notif_type : ident $( ( $notif_param : ty ) )*; )*
        }
    ) => {
        #[derive( Debug )]
        pub enum RequestMethod {
            $( $method_type $( ( $method_param ) )*, )*
        }

        #[derive( Debug )]
        pub enum NotificationMethod {
            $( $notif_type $( ( $notif_param ) )*, )*
        }

        pub fn parse_server_message( msg_json_str : &str ) -> io::Result< ServerMessage > {
            let msg_json : serde_json::Value = match serde_json::from_str( msg_json_str ) {
                Ok( msg_json ) => msg_json,
                Err( _ ) => return Err( new_invalid_data_error( "Invalid JSON data." ) )
            };

            if let Some( method_val ) = msg_json.lookup( "method" ) {
                let method = match method_val.as_str( ) {
                    Some( method ) => method,
                    None => return Err( new_invalid_data_error( "Invalid 'method' field, could not parse as String." ) )
                };
                match method {
                    $(
                        $method_name => {
                            let msg_id = match msg_json.lookup( "id" ) {
                                Some( msg_id ) => {
                                    match msg_id.as_i64( ) {
                                        Some( msg_id ) => msg_id,
                                        None => return Err( new_invalid_data_error( "Invalid 'id' field, could not parse as number." ) )
                                    }
                                },
                                None => return Err( new_invalid_data_error( "Missing required 'id' field." ) )
                            };
                            let msg = RequestMessage {
                                id     : msg_id,
                                method : RequestMethod::$method_type $( ( parse_parameters::< $method_param >( &msg_json )? ) )*
                            };
                            Ok( ServerMessage::Request( msg ) )
                        }
                    ),*
                    $(
                        $notif_name => {
                            let notif = NotificationMessage {
                                method : NotificationMethod::$notif_type $( ( parse_parameters::< $notif_param >( &msg_json )? ) )*
                            };
                            Ok( ServerMessage::Notification( notif ) )
                        }
                    ),*
                    _ => Err( new_invalid_data_error( "Invalid request method." ) )
                }
            }
            else {
                return Err( new_invalid_data_error( "Missing required 'method' field." ) );
            }
        }
    };
}

macro_rules! server_responses {
    (
        datas {
            $( $data_type : ident( $data_param : ty ); )*
        },
        errors {
            $( $err_name : ident( $err_val : expr ); )*
        }
    ) => {
        #[derive( Debug )]
        pub enum RequestResponseData {
            $( $data_type ( $data_param ) ),*
        }

        $( pub static $err_name : i64 = $err_val; )*

        impl serde::Serialize for RequestResponseData {

            fn serialize< S >( &self, serializer : &mut S ) -> Result< ( ), S::Error > where S : serde::Serializer {
                match *self {
                    $(
                        RequestResponseData::$data_type( ref param ) => {
                            param.serialize( serializer )
                        }
                    ),*
                }
            }

        }
    }
}

/// A completely parsed message and its headers sent from the client to the server.
#[derive( Debug )]
pub struct IncomingServerMessage {
    /// The headers that accompanied the message.
    pub headers : HashMap< String, String >,

    /// The parsed message.
    pub message : ServerMessage
}

/// A complete message to send from the server to the client.
#[derive( Debug )]
pub struct OutgoingServerMessage {
    /// The headers that should accompany this message.
    ///
    /// Content-Length and Content-Type will be ignored when writing all headers.
    pub headers : HashMap< String, String >,

    /// The message to send to the client.
    pub message : ClientMessage
}

/// Represents either a request or a notification sent from the client to the server.
#[derive( Debug )]
pub enum ServerMessage {
    /// A RequestMessage sent from the client to the server.
    Request( RequestMessage ),

    /// A NotificationMessage sent from the client to the server.
    Notification( NotificationMessage )
}

/// Represents a response from the server to the client.
#[derive( Debug )]
pub enum ClientMessage {
    /// A RequestResponseMessage that contains the response for a client's previously made request.
    Response( RequestResponseMessage )
}

/// A Request from the client to the server.
///
/// Requests must be responded to in the order they were received, but they do not have to be
/// processed in any specific order.
#[derive( Debug )]
pub struct RequestMessage {
    /// The ID of the request that is used to distinguish multiple parallel request's and their responses.
    pub id     : i64,
    /// The method and it's parameters to process.
    pub method : RequestMethod
}

/// A Notification from the client to the server.
///
/// Notifications must not be responded to and can be handled in any order the server sees fit.
#[derive( Debug )]
pub struct NotificationMessage {
    /// The method and it's parameters to process.
    pub method : NotificationMethod
}

/// An error response from the server to the client in response to a specific request.
#[derive( Debug, Serialize )]
pub struct RequestResponseError {
    /// The error code.
    pub code    : i64,
    /// An accompanying error message to present to the client.
    pub message : String
}

/// A message from the server to the client in response to a specific request.
///
/// May contain either a result or an error for either success or failure of a specific client request.
#[derive( Debug )]
pub struct RequestResponseMessage {
    /// The id of the request that this message responds to.
    pub id     : i64,
    /// The result of a successful completion of the request.
    pub result : Option< RequestResponseData >,
    /// The error of a request that failed.
    pub error  : Option< RequestResponseError >
}

enum RLSCodecState {
    Headers,
    Body
}

/// Implements decoding of IncomingServerMessage and encoding of OutgoingServerMessage
///
/// Implements tokio_core::io::Core
pub struct RLSCodec {
    codec_state    : RLSCodecState,
    headers        : HashMap< String, String >,
    content_length : usize
}

server_messages! {
    methods {
        REQUEST__Initialize                        => Initialize( InitializeParams );
        REQUEST__Shutdown                          => Shutdown;
        REQUEST__Completion                        => Completion( TextDocumentPositionParams );
        REQUEST__ResolveCompletionItem             => CompletionResolve( CompletionItem );
        REQUEST__Hover                             => Hover( TextDocumentPositionParams );
        REQUEST__SignatureHelp                     => SignatureHelp( TextDocumentPositionParams );
        REQUEST__GotoDefinition                    => GotoDefinition( TextDocumentPositionParams );
        REQUEST__References                        => FindReferences( ReferenceParams );
        REQUEST__DocumentHighlight                 => DocumentHighlights( TextDocumentPositionParams );
        REQUEST__DocumentSymbols                   => DocumentSymbols( DocumentSymbolParams );
        REQUEST__WorkspaceSymbols                  => WorkspaceSymbols( WorkspaceSymbolParams );
        REQUEST__CodeAction                        => CodeAction( CodeActionParams );
        REQUEST__CodeLens                          => CodeLens( CodeLensParams );
        REQUEST__CodeLensResolve                   => CodeLensResolve( CodeLens );
        REQUEST__DocumentLink                      => DocumentLink( DocumentLinkParams );
        REQUEST__DocumentLinkResolve               => DocumentLinkResolve( DocumentLink );
        REQUEST__Formatting                        => DocumentFormatting( DocumentFormattingParams );
        REQUEST__RangeFormatting                   => DocumentRangeFormatting( DocumentRangeFormattingParams );
        REQUEST__OnTypeFormatting                  => DocumentOnTypeFormatting( DocumentOnTypeFormattingParams );
        REQUEST__Rename                            => Rename( RenameParams );
    },
    notifications {
        NOTIFICATION__Cancel                       => Cancel( CancelParams );
        NOTIFICATION__Exit                         => Exit;
        NOTIFICATION__WorkspaceChangeConfiguration => DidChangeConfiguration( DidChangeConfigurationParams );
        NOTIFICATION__DidOpenTextDocument          => DidOpenTextDocument( DidOpenTextDocumentParams );
        NOTIFICATION__DidChangeTextDocument        => DidChangeTextDocument( DidChangeTextDocumentParams );
        NOTIFICATION__DidCloseTextDocument         => DidCloseTextDocument( DidCloseTextDocumentParams );
        NOTIFICATION__DidSaveTextDocument          => DidSaveTextDocument( DidSaveTextDocumentParams );
        NOTIFICATION__DidChangeWatchedFiles        => DidChangeWatchedFiles( DidChangeWatchedFilesParams );
    }
}

server_responses! {
    datas {
        Init( InitializeResult );
        SymbolInfo( Vec< SymbolInformation > );
        CompletionItems( Vec< CompletionItem > );
        WorkspaceEdit( WorkspaceEdit );
        TextEdit( [ TextEdit; 1 ] );
        Locations( Vec< Location > );
        Highlights( Vec< DocumentHighlight > );
        HoverSuccess( Hover );
    },
    errors {
        PARSE_ERROR( -32700 );
        INVALID_REQUEST( -32600 );
        METHOD_NOT_FOUND( -32601 );
        INVALID_PARAMS( -32602 );
        INTERNAL_ERROR( -32603 );
        SERVER_ERROR_START( -32099 );
        SERVER_ERROR_END( -32000 );
        SERVER_NOT_INITIALIZED( -32002 );
        UNKNOWN_ERROR_CODE( -32001 );
    }
}

fn new_invalid_data_error( error_msg : &'static str ) -> io::Error {
    io::Error::new( io::ErrorKind::InvalidData, error_msg )
}

fn parse_parameters< T >( msg_json : &serde_json::Value ) -> io::Result< T > where T : serde::Deserialize {
    let params = match msg_json.lookup( "params" ) {
        Some( params ) => params.to_owned( ),
        None => return Err( new_invalid_data_error( "Missing required 'params' field." ) )
    };

    match serde_json::from_value( params ) {
        Ok( value ) => Ok( value ),
        Err( _ ) => Err( new_invalid_data_error( "Unable to parse parameter field." ) )
    }
}

fn easybuf_to_utf8( buf : &EasyBuf ) -> io::Result< &str > {
    match str::from_utf8( buf.as_ref( ) ) {
        Ok( s ) => Ok( s ),
        Err( _ ) => Err( new_invalid_data_error( "Unable to parse bytes as UTF-8." ) )
    }
}

impl RLSCodec {

    pub fn new( ) -> Self {
        RLSCodec {
            codec_state    : RLSCodecState::Headers,
            headers        : HashMap::new( ),
            content_length : 0
        }
    }

    fn read_header( &mut self, buf : &mut EasyBuf ) -> io::Result< bool > {
        if let Some( pos ) = buf.as_ref( ).windows( 2 ).position( | bytes | bytes == b"\r\n" ) {
            let header_bytes = buf.drain_to( pos );
            buf.drain_to( 2 );

            if pos == 0 {
                self.content_length = self.parse_content_length( )?;
                self.codec_state    = RLSCodecState::Body;

                trace!( "Finished parsing headers. Waiting for '{}' bytes for body.", self.content_length );
            }
            else {
                let header_str = easybuf_to_utf8( &header_bytes )?;
                trace!( "Received header line: {}", header_str );

                let mut header_parts = header_str.split( ": " );

                if header_parts.clone( ).count( ) != 2 {
                    return Err( new_invalid_data_error( "Unable to parse invalid header." ) );
                }

                let header_key = header_parts.next( ).unwrap( );
                let header_val = header_parts.next( ).unwrap( );

                self.headers.insert( header_key.to_string( ), header_val.to_string( ) );
            }
            Ok( true )
        }
        else {
            Ok( false )
        }
    }

    fn parse_content_length( &self ) -> io::Result< usize > {
        if let Some( content_length_str ) = self.headers.get( "Content-Length" ) {
            if let Ok( content_length ) = content_length_str.parse::< usize >( ) {
                Ok( content_length )
            }
            else {
                Err( new_invalid_data_error( "Unable to parse 'Content-Length' header into usize." ) )
            }
        }
        else {
            Err( new_invalid_data_error( "Missing required 'Content-Length' header." ) )
        }
    }

    fn reset_codec( &mut self ) -> HashMap< String, String > {
        let mut new_headers = HashMap::new( );
        mem::swap( &mut self.headers, &mut new_headers );

        self.codec_state = RLSCodecState::Headers;

        new_headers
    }

    fn parse_body_contents( &mut self, buf : &mut EasyBuf ) -> io::Result< Option< IncomingServerMessage > > {
        if buf.len( ) < self.content_length {
            return Ok( None );
        }

        let body_bytes = buf.drain_to( self.content_length );
        let body_str   = easybuf_to_utf8( &body_bytes )?;
        trace!( "Received body contents: {}", body_str );

        let message = parse_server_message( body_str )?;
        trace!( "Parsed server message: {:?}", message );

        let headers = self.reset_codec( );

        Ok( Some( IncomingServerMessage{
            headers : headers,
            message : message
        } ) )
    }

}

impl Codec for RLSCodec {

    type In  = IncomingServerMessage;
    type Out = OutgoingServerMessage;

    fn decode( &mut self, buf : &mut EasyBuf ) -> io::Result< Option< Self::In > > {
        loop {
            match self.codec_state {
                RLSCodecState::Headers => {
                    if !self.read_header( buf )? {
                        return Ok( None );
                    }
                },
                RLSCodecState::Body => {
                    return self.parse_body_contents( buf );
                }
            }
        }
    }

    fn encode( &mut self, msg : OutgoingServerMessage, buffer : &mut Vec< u8 > ) -> io::Result< ( ) > {
        let OutgoingServerMessage{ mut headers, message } = msg;

        let msg_json = message.into_json( );
        headers.insert( "Content-Length".to_string( ), msg_json.as_bytes( ).len( ).to_string( ) );
        headers.insert( "Content-Type".to_string( ), "application/vscode-jsonrpc; charset=utf8".to_string( ) );

        for ( ref key, ref val ) in headers {
            trace!( "Sending header: {} = {}", key, val );

            buffer.extend_from_slice( key.as_bytes( ) );
            buffer.extend_from_slice( b": " );
            buffer.extend_from_slice( val.as_bytes( ) );
            buffer.extend_from_slice( b"\r\n" );
        }
        buffer.extend_from_slice( b"\r\n" );

        trace!( "Sending body: {}", msg_json );
        buffer.extend_from_slice( msg_json.as_bytes( ) );

        Ok( ( ) )
    }
}

impl ClientMessage {

    fn into_json( self ) -> String {
        match self {
            ClientMessage::Response( value ) => {

                #[derive( Serialize )]
                struct ResponseMessage {
                    jsonrpc : &'static str,
                    id      : i64,
                    result  : Option< RequestResponseData >,
                    error   : Option< RequestResponseError >
                }

                serde_json::to_string( &ResponseMessage {
                    jsonrpc : "2.0",
                    id      : value.id,
                    result  : value.result,
                    error   : value.error
                } ).unwrap( )
            }
        }
    }

}<|endoftext|>


<commit_before>use std::io::Write;
use std::net::{TcpListener, TcpStream};
use std::thread;

fn handle_client(mut stream: TcpStream) {
    println!("new conn");
    stream.write(b"Hello World\r\n").unwrap();
}

fn main () {
    let listener = TcpListener::bind("0.0.0.0:3451").unwrap();
    println!("listening");
    for stream in listener.incoming() {
        match stream {
            Ok(stream) => {
                thread::spawn(move|| {
                    handle_client(stream)
                });
            }
            Err(e) => println!("connection failed!")
        }
    }
}
<commit_msg>Delete server.rs<commit_after><|endoftext|>


<commit_before>use std::thread;
use std::thread::JoinHandle;
use std::path::Path;

use websocket::{WebSocketStream, Server};
use websocket::Message as WebMessage;
use websocket::Sender as WebSender;
use websocket::client::Sender as WebClientSender;
use websocket::Receiver as WebReceiver;
use websocket::message::Type as WebMessageType;
use websocket::server::Connection;
use websocket::header::WebSocketProtocol;

use git2::Repository;
use git2::ObjectType;

use payload;
use State;

extern crate fsevent;
use self::fsevent::Event as FsEvent;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;

use serde_xml as xml;
use yaml_rust::yaml::Hash;
use yaml_rust::{Yaml, YamlEmitter};

#[derive(Debug)]
enum NotifierMessage<'a> {
    WebMessage(WebMessage<'a>),
    FsEvent(FsEvent)
}

pub fn start() {
    let server = Server::bind("127.0.0.1:2794").unwrap();
    thread::spawn(move || {
        for connection in server {
            thread::spawn(move || connect(connection.unwrap()));
        }
    });
}

fn connect(connection: Connection<WebSocketStream, WebSocketStream>) {
    let request = connection.read_request().unwrap();
    let headers = request.headers.clone();
    request.validate().unwrap();
    let mut response = request.accept();
    println!("Connected");

    if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
        if protocols.contains(&("rust-websocket".to_string())) {
            response.headers.set(WebSocketProtocol(vec!["superconductor".to_string()]));
        }
    }

    let mut client = response.send().unwrap();
    let message = WebMessage::text(payload::generate(None));
    client.send_message(&message).unwrap();

    let (tx, rx) = channel::<NotifierMessage>();
    let (mut sender, mut receiver) = client.split();
    let notifier = thread::spawn(move || start_notifier(rx, sender));
    let txc = tx.clone();
    let monitor = thread::spawn(move || start_monitor(txc));
    for message in receiver.incoming_messages() {
        let message: WebMessage = message.unwrap_or(WebMessage::close());
        match message.opcode {
            WebMessageType::Close => {
                tx.send(NotifierMessage::WebMessage(message));
                break;
            },
            _ => tx.send(NotifierMessage::WebMessage(message))
        };
    }
    monitor.join();
    notifier.join();
}

fn start_monitor(tx: Sender<NotifierMessage>) {
    let (ftx, rx) = channel();
    let observer = thread::spawn(move || {
        println!("Monitoring");
        let fsevent = fsevent::FsEvent::new(ftx);
        fsevent.append_path(".");
        fsevent.observe();
    });
    loop {
        let event = rx.recv().unwrap();
        tx.send(NotifierMessage::FsEvent(event));
    }
    observer.join();
}

fn start_notifier(rx: Receiver<NotifierMessage>, mut sender: WebClientSender<WebSocketStream>) {
    let mut last_state: Option<State> = None;
    loop {
        let value = rx.recv().unwrap();
        match value {
            NotifierMessage::WebMessage(message) => {
                match message.opcode {
                    WebMessageType::Close => {
                        break;
                    },
                    WebMessageType::Ping => {
                        let message = WebMessage::pong(message.payload);
                        sender.send_message(&message).unwrap();
                    },
                    _ => {
                        let payload = String::from_utf8_lossy(message.payload.as_ref());
                        println!("{}", payload);
                        let state: State = xml::from_str(&payload).unwrap();
                        println!("{:?}", state);

                        let repo = Repository::discover(".").unwrap();
                        let head = repo.head().unwrap().peel(ObjectType::Commit).unwrap();
                        let mut index = repo.index().unwrap();

                        let to_remove = index.iter().fold(vec![], |mut acc, entry| {
                            let entry_path = String::from_utf8_lossy(entry.path.as_ref());
                            match state.include.iter().find(|i| i.as_ref() == entry_path) {
                                None => acc.push(entry_path.into_owned()),
                                _ => {}
                            };
                            acc
                        });

                        repo.reset_default(Some(&head), to_remove.iter()).unwrap();
                        for change in state.clone().include {
                            let path = Path::new(&change);
                            index.add_path(path).unwrap();
                        }
                        index.write().unwrap();

                        if let Some(event) = state.save_update.clone() {
                            let author = repo.signature().unwrap();
                            let mut yaml = String::new();
                            {
                                // Constructing the properties YAML
                                let mut tasks = Hash::new();
                                let mut properties = Hash::new();
                                let mut emitter = YamlEmitter::new(&mut yaml);
                                for property in state.property {
                                    properties.insert(Yaml::String(property.name), Yaml::String(property.value));
                                }
                                tasks.insert(Yaml::String(String::from("BLBA-1234")), Yaml::Hash(properties));
                                emitter.dump(&Yaml::Hash(tasks)).unwrap();
                            }
                            let message = state.message + "\n" + &yaml;

                            index.read(false);
                            let tree_oid = index.write_tree().unwrap();
                            let tree = repo.find_tree(tree_oid).unwrap();
                            repo.commit(Some("HEAD"), &author, &author, &message, &tree, &[&head.as_commit().unwrap()]);
                            let message = WebMessage::text(payload::generate(None));
                            sender.send_message(&message).unwrap();
                            last_state = None;
                        } else {
                            last_state = Some(state);
                        }
                        let message = WebMessage::text(payload::generate(last_state.clone()));
                        sender.send_message(&message).unwrap();
                    }
                }
            },
            NotifierMessage::FsEvent(event) => {
                let repo = Repository::discover(".").unwrap();
                let message = WebMessage::text(payload::generate(last_state.clone()));
                sender.send_message(&message).unwrap();
            }
        }
    }
}
<commit_msg>Adds fs aggregator to prevent over accumulating notifier events. --- "BLBA-1234":    Description: ""   Developer: Sprint   Estimate: "5"   Status: Done<commit_after>use std::thread;
use std::thread::JoinHandle;
use std::path::Path;
use std::time::Duration;

use websocket::{WebSocketStream, Server};
use websocket::Message as WebMessage;
use websocket::Sender as WebSender;
use websocket::client::Sender as WebClientSender;
use websocket::Receiver as WebReceiver;
use websocket::message::Type as WebMessageType;
use websocket::server::Connection;
use websocket::header::WebSocketProtocol;

use git2::Repository;
use git2::ObjectType;

use payload;
use State;

extern crate fsevent;
use self::fsevent::Event as FsEvent;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use std::sync::mpsc::Receiver;

use serde_xml as xml;
use yaml_rust::yaml::Hash;
use yaml_rust::{Yaml, YamlEmitter};

#[derive(Debug)]
enum NotifierMessage<'a> {
    WebMessage(WebMessage<'a>),
    FsEvent(FsEvent)
}

pub fn start() {
    let server = Server::bind("127.0.0.1:2794").unwrap();
    thread::spawn(move || {
        for connection in server {
            thread::spawn(move || connect(connection.unwrap()));
        }
    });
}

fn connect(connection: Connection<WebSocketStream, WebSocketStream>) {
    let request = connection.read_request().unwrap();
    let headers = request.headers.clone();
    request.validate().unwrap();
    let mut response = request.accept();
    println!("Connected");

    if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
        if protocols.contains(&("rust-websocket".to_string())) {
            response.headers.set(WebSocketProtocol(vec!["superconductor".to_string()]));
        }
    }

    let mut client = response.send().unwrap();
    let message = WebMessage::text(payload::generate(None));
    client.send_message(&message).unwrap();

    let (tx, rx) = channel::<NotifierMessage>();
    let (mut sender, mut receiver) = client.split();
    let notifier = thread::spawn(move || start_notifier(rx, sender));
    let txc = tx.clone();
    let monitor = thread::spawn(move || start_monitor(txc));
    for message in receiver.incoming_messages() {
        let message: WebMessage = message.unwrap_or(WebMessage::close());
        match message.opcode {
            WebMessageType::Close => {
                tx.send(NotifierMessage::WebMessage(message));
                break;
            },
            _ => tx.send(NotifierMessage::WebMessage(message))
        };
    }
    monitor.join();
    notifier.join();
}

fn start_monitor(tx: Sender<NotifierMessage>) {
    let (ftx, rx) = channel();
    let observer = thread::spawn(move || {
        println!("Monitoring");
        let fsevent = fsevent::FsEvent::new(ftx);
        fsevent.append_path(".");
        fsevent.observe();
    });
    loop {
        let event = rx.recv().unwrap();
        thread::sleep(Duration::from_millis(500));
        let mut aggregator = vec![];
        while let Ok(agg_event) = rx.try_recv() {
            aggregator.push(agg_event);
        }
        tx.send(NotifierMessage::FsEvent(event));
    }
    observer.join();
}

fn start_notifier(rx: Receiver<NotifierMessage>, mut sender: WebClientSender<WebSocketStream>) {
    let mut last_state: Option<State> = None;
    loop {
        let value = rx.recv().unwrap();
        match value {
            NotifierMessage::WebMessage(message) => {
                match message.opcode {
                    WebMessageType::Close => {
                        break;
                    },
                    WebMessageType::Ping => {
                        let message = WebMessage::pong(message.payload);
                        sender.send_message(&message).unwrap();
                    },
                    _ => {
                        let payload = String::from_utf8_lossy(message.payload.as_ref());
                        println!("{}", payload);
                        let state: State = xml::from_str(&payload).unwrap();
                        println!("{:?}", state);

                        let repo = Repository::discover(".").unwrap();
                        let head = repo.head().unwrap().peel(ObjectType::Commit).unwrap();
                        let mut index = repo.index().unwrap();

                        let to_remove = index.iter().fold(vec![], |mut acc, entry| {
                            let entry_path = String::from_utf8_lossy(entry.path.as_ref());
                            match state.include.iter().find(|i| i.as_ref() == entry_path) {
                                None => acc.push(entry_path.into_owned()),
                                _ => {}
                            };
                            acc
                        });

                        repo.reset_default(Some(&head), to_remove.iter()).unwrap();
                        for change in state.clone().include {
                            let path = Path::new(&change);
                            index.add_path(path).unwrap();
                        }
                        index.write().unwrap();

                        if let Some(event) = state.save_update.clone() {
                            let author = repo.signature().unwrap();
                            let mut yaml = String::new();
                            {
                                // Constructing the properties YAML
                                let mut tasks = Hash::new();
                                let mut properties = Hash::new();
                                let mut emitter = YamlEmitter::new(&mut yaml);
                                for property in state.property {
                                    properties.insert(Yaml::String(property.name), Yaml::String(property.value));
                                }
                                tasks.insert(Yaml::String(String::from("BLBA-1234")), Yaml::Hash(properties));
                                emitter.dump(&Yaml::Hash(tasks)).unwrap();
                            }
                            let message = state.message + "\n" + &yaml;

                            index.read(false);
                            let tree_oid = index.write_tree().unwrap();
                            let tree = repo.find_tree(tree_oid).unwrap();
                            repo.commit(Some("HEAD"), &author, &author, &message, &tree, &[&head.as_commit().unwrap()]);
                            let message = WebMessage::text(payload::generate(None));
                            sender.send_message(&message).unwrap();
                            last_state = None;
                        } else {
                            last_state = Some(state);
                        }
                        let message = WebMessage::text(payload::generate(last_state.clone()));
                        sender.send_message(&message).unwrap();
                    }
                }
            },
            NotifierMessage::FsEvent(event) => {
                let repo = Repository::discover(".").unwrap();
                let message = WebMessage::text(payload::generate(last_state.clone()));
                sender.send_message(&message).unwrap();
            }
        }
    }
}
<|endoftext|>


<commit_before>use grain_capnp::{PowerboxCapability, UiView, UiSession};
use web_session_capnp::{WebSession};

use collections::hashmap::HashMap;
use capnp::capability::{ClientHook, FromServer};
use capnp::AnyPointer;
use capnp_rpc::rpc::{RpcConnectionState, SturdyRefRestorer};
use capnp_rpc::capability::{LocalClient};

use sqlite3;

pub struct UiViewImpl;

impl PowerboxCapability::Server for UiViewImpl {
    fn get_powerbox_info(&mut self, context : PowerboxCapability::GetPowerboxInfoContext) {
        context.done()
    }
}

impl UiView::Server for UiViewImpl {
    fn get_view_info(&mut self, context : UiView::GetViewInfoContext) {
        context.done()
    }

    fn new_session(&mut self, mut context : UiView::NewSessionContext) {
        println!("asked for a new session!");
        let (_, results) = context.get();

        let client : WebSession::Client = FromServer::new(None::<LocalClient>, ~WebSessionImpl::new());
        // we need to do this dance to upcast.
        results.set_session(UiSession::Client { client : client.client});

        context.done()
    }
}

pub struct WebSessionImpl {
    db : sqlite3::Database,
}

impl WebSessionImpl {
    pub fn new() -> WebSessionImpl {
        WebSessionImpl {
            db : sqlite3::open("/var/data.db").unwrap(),
        }
    }
}

impl UiSession::Server for WebSessionImpl {

}

struct Path {
    path : ~str,
    query : HashMap<~str, ~str>,
}

impl Path {
    fn new() -> Path {
        Path { path : ~"", query : HashMap::new() }
    }
}

fn parse_path(path : &str) -> Path {
    let mut result = Path::new();

    let v : ~[&str] = path.splitn('?', 2).collect();
    if v.len() == 0 {
        return result;
    }
    result.path = v[0].into_owned();
    if v.len() == 1 {
        return result;
    }
    for attr in v[1].split('&') {
        let a : ~[&str] = attr.splitn('=', 2).collect();
        if a.len() == 2 {
            result.query.insert(a[0].into_owned(),
                                a[1].into_owned());
        }
    }
    return result;
}

impl WebSessionImpl {

    fn checked<T>(&self, result : sqlite3::SqliteResult<T>) -> T {
        match result {
            Err(e) => {fail!("database error: {}, {:?}", e, self.db.get_errmsg()) }
            Ok(v) => { return v; }
        }
    }

    fn is_word(&self, word : &str) -> sqlite3::SqliteResult<bool> {

        if ! word.is_alphanumeric() { return Ok(false); }

        let cursor = try!(self.db.prepare(
            format!("SELECT * FROM Words WHERE Word = \"{}\";", word),
            &None));

        return Ok(try!(cursor.step_row()).is_some());
    }

    fn validate_def(&self, word : &str, definition : &[&str]) -> sqlite3::SqliteResult<bool> {

        if definition.len() != word.len() { return Ok(false); }

        let mut idx = 0;
        for &d in definition.iter() {
            if !(try!(self.is_word(d)) && d.len() > 0 && d[0] == word[idx]) {
                return Ok(false);
            }

            idx += 1;
        }

        return Ok(true);
    }

    fn write_def(&self, word : &str, definition : &[&str]) -> sqlite3::SqliteResult<()> {

        let time : i64 = ::time::get_time().sec;
        let mut query = StrBuf::new();
        query.push_str(format!("BEGIN; DELETE FROM Definitions WHERE Definee =\"{}\"; ", word));
        query.push_str("INSERT INTO Definitions(Definee, Idx, Definer) VALUES");
        let mut idx = 0;
        for &d in definition.iter() {
            if idx != 0 { query.push_str(","); }
            query.push_str(format!("(\"{}\", {}, \"{}\")", word, idx, d));
            idx += 1;
        }
        query.push_str(";");
        query.push_str(format!("INSERT Into Log(Word, Timestamp) VALUES(\"{}\",{});", word, time));
        query.push_str("COMMIT;");

        println!("query: {}", query);

        try!(self.db.exec(query.as_slice()));

        return Ok(());
    }

    fn get_def(&self, word : &str) -> sqlite3::SqliteResult<~str> {

        let cursor = try!(self.db.prepare(
            format!("SELECT * FROM Definitions WHERE Definee = \"{}\";", word),
            &None));

        let mut map = HashMap::<int, ~str>::new();

        loop {
            match try!(cursor.step_row()) {
                None => break,
                Some(row) => {
                    let definer = match row.get(&~"Definer") { &sqlite3::Text(ref t) => t.clone(), _ => fail!(), };
                    let idx = match row.get(&~"Idx") { &sqlite3::Integer(ref i) => i.clone(), _ => fail!(), };

                    map.insert(idx, definer);
                }
            }
        }

        if map.len() != word.len() {
            return Ok(~"
this word has no definition yet
");
        } else {

            let mut result = StrBuf::new();
            result.push_str("
");
            for idx in range::<int>(0, word.len() as int) {
                let definer : &str = map.get(&idx).as_slice();
                result.push_str(format!(" {word} ", word=definer));
            }
            result.push_str("
");
            return Ok(result.into_owned());
        }
    }

    fn count_defs(&self) -> sqlite3::SqliteResult<(int, int, Vec<~str>)> {
        let cursor = try!(self.db.prepare("SELECT COUNT(*) FROM Words;", &None));
        assert!(cursor.step() == sqlite3::SQLITE_ROW);
        let num_words = cursor.get_int(0);

        let cursor = try!(self.db.prepare("SELECT COUNT(*) FROM Definitions WHERE Idx = 0;", &None));
        assert!(cursor.step() == sqlite3::SQLITE_ROW);
        let defined_words = cursor.get_int(0);

        let mut recent_words = Vec::new();
        let cursor = try!(
            self.db.prepare("SELECT Word, Timestamp FROM Log ORDER BY Timestamp DESC LIMIT 5;", &None));
        loop {
            match try!(cursor.step_row()) {
                None => break,
                Some(row) => {
                    let word : ~str = match row.get(&~"Word") {&sqlite3::Text(ref t) => t.clone(), _ => fail!(),};
                    recent_words.push(word);
                }
            }
        }
        Ok((defined_words, num_words, recent_words))

    }

}

static main_css : &'static str =
    "body { font-family: Helvetica, Sans, Arial;
            font-size: medium;
             margin-left: auto;
             margin-right: auto;
             width: 600px;
             text-align: center;
     }
     div {
          padding-bottom: 10pt;
     }
    .word {
        text-align: center;
        font-size: 500%;
     }
     .err {
       font-size: 90%;
       color: #AA0000;
     }
     .title {
       text-align: center;
       font-size:500%;
     }
     ";

static header : &'static str =
  r#"<title> acronymy </title><link rel="stylesheet" type="text/css" href="main.css" >
 <meta http-equiv="Content-Type" content="text/html;charset=utf-8" >
  "#;

static lookup_form : &'static str =
      r#"<form action="define" method="get">
          <input name="word" maxlength=\"100\"/><button>find word</button></form>"#;

fn define_form(word :&str) -> ~str {
       format!("<form action=\"define\" method=\"get\">
               <input name=\"word\" value=\"{word}\" type=\"hidden\"/>
               <input name=\"definition\" maxlength=\"2000\"/>
                   <button>submit definition</button></form>", word=word)
}

enum PageData<'a> {
    NoWord,
    WordAndDef(&'a str, &'a str, Option<&'a str>),
    HomePage(int, int, Vec<~str>),
}

fn construct_html(page_data : PageData) -> ~str {
    let mut result = StrBuf::new();
    result.push_str(format!("{}", header));

    static home_link : &'static str = "home";
    match page_data {
        NoWord => {
            result.push_str("
 that's not a word
");
            result.push_str(lookup_form);
            result.push_str(home_link);
        }
        WordAndDef(word, def_div, err) => {
            result.push_str(format!("
{}
", word));

            result.push_str(def_div);

            match err {
                None => {}
                Some(e) => {
                    result.push_str(format!("
{}
", e));
                }
            }

            result.push_str(define_form(word));
            result.push_str(home_link);
        }
        HomePage(num_defined, total, recent) => {
            result.push_str("
Acronymy
");
            result.push_str("
A user-editable, acronym-only dictionary.
");
            result.push_str(format!("
So far, we have defined {} out of {} words.
",
                                    num_defined, total));
            if recent.len() > 0 {
                result.push_str("
Recently modified words: ");
                let mut idx = 0;
                for w in recent.iter() {
                    if idx != 0 {
                        result.push_str(", ");
                    }
                    result.push_str(format!("{word}", word=*w));
                    idx += 1;
                }
                result.push_str(".
");
            }
            result.push_str(lookup_form);
        }
    }

    result.push_str("");
    result.into_owned()
}

impl WebSession::Server for WebSessionImpl {
    fn get(&mut self, mut context : WebSession::GetContext) {
        println!("GET");
        let (params, results) = context.get();
        let raw_path = params.get_path();
        let content = results.init_content();
        content.set_mime_type("text/html");

        let path = parse_path(raw_path);
        println!("path = {}", raw_path);

        if raw_path == "main.css" {
            content.get_body().set_bytes(main_css.as_bytes())
        } else if path.path.as_slice() == "define" {
            let word : ~str = match path.query.find(&~"word") {
                Some(w) if self.checked(self.is_word(*w)) => {
                    w.clone()
                }
                _ => {
                    content.get_body().set_bytes(construct_html(NoWord).as_bytes());
                    return context.done();
                }
            };

            match path.query.find(&~"definition") {
                None => {
                    let def_div = self.checked(self.get_def(word));

                    content.get_body().set_bytes(construct_html(WordAndDef(word.as_slice(),
                                                                           def_div.as_slice(),
                                                                           None)).as_bytes());
                }
                Some(def_query) => {

                    let definition : ~[&str] = def_query.split('+').collect();

                    if self.checked(self.validate_def(word, definition)) {

                        self.checked(self.write_def(word, definition));
                        let def_div = self.checked(self.get_def(word));
                        content.get_body().set_bytes(
                            construct_html(WordAndDef(word.as_slice(),
                                                      def_div.as_slice(),
                                                      None)).as_bytes());

                    } else {

                        let def_div = self.checked(self.get_def(word));

                        content.get_body().set_bytes(
                            construct_html(WordAndDef(word.as_slice(),
                                                      def_div.as_slice(),
                                                      Some("invalid definition"))).as_bytes());
                    }
                }
            }

        } else {
            let (num_defined, total, recent) = self.checked(self.count_defs());
            content.get_body().set_bytes(construct_html(HomePage(num_defined, total, recent)).as_bytes());
        }
        context.done()
    }
    fn post(&mut self, context : WebSession::PostContext) {
        println!("POST");
        context.done()
    }
    fn put(&mut self, context : WebSession::PutContext) {
        println!("PUT");
        context.done()
    }
    fn delete(&mut self, context : WebSession::DeleteContext) {
        println!("DELETE");
        context.done()
    }
    fn open_web_socket(&mut self, context : WebSession::OpenWebSocketContext) {
        println!("OPEN WEB SOCKET");
        context.done()
    }
}

pub struct FdStream {
    inner : ~::std::rt::rtio::RtioFileStream:Send,
}

impl FdStream {
    pub fn new(fd : ::libc::c_int) -> ::std::io::IoResult<FdStream> {
        ::std::rt::rtio::LocalIo::maybe_raise(|io| {
            Ok (FdStream { inner : io.fs_from_raw_fd(fd, ::std::rt::rtio::DontClose) })
        })
    }
}

impl Reader for FdStream {
    fn read(&mut self, buf : &mut [u8]) -> ::std::io::IoResult<uint> {
        self.inner.read(buf).map(|i| i as uint)
    }
}

impl Writer for FdStream {
    fn write(&mut self, buf : &[u8]) -> ::std::io::IoResult<()> {
        self.inner.write(buf)
    }
}

pub struct Restorer;

impl SturdyRefRestorer for Restorer {
    fn restore(&self, obj_id : AnyPointer::Reader) -> Option<~ClientHook:Send> {
        if obj_id.is_null() {
            let client : UiView::Client = FromServer::new(None::<LocalClient>, ~UiViewImpl);
            Some(client.client.hook)
        } else {
            None
        }
    }
}

pub fn main() -> ::std::io::IoResult<()> {

    let args = ::std::os::args();

    if args.len() == 4 && args[1].as_slice() == "--init" {
        println!("initializing...");
        let initdb_path = ::std::path::Path::new(args[2].as_slice());
        let proddb_path = ::std::path::Path::new(args[3].as_slice());
        println!("copying database from {} to {}", args[2], args[3]);
        try!(::std::io::fs::copy(&initdb_path, &proddb_path));
        println!("success!");
    }

    // sandstorm launches us with a connection file descriptor 3
    let ifs = try!(FdStream::new(3));
    let ofs = try!(FdStream::new(3));

    let connection_state = RpcConnectionState::new();
    connection_state.run(ifs, ofs, Restorer);

    Ok(())
}

<commit_msg>fix maxlength attribute<commit_after>use grain_capnp::{PowerboxCapability, UiView, UiSession};
use web_session_capnp::{WebSession};

use collections::hashmap::HashMap;
use capnp::capability::{ClientHook, FromServer};
use capnp::AnyPointer;
use capnp_rpc::rpc::{RpcConnectionState, SturdyRefRestorer};
use capnp_rpc::capability::{LocalClient};

use sqlite3;

pub struct UiViewImpl;

impl PowerboxCapability::Server for UiViewImpl {
    fn get_powerbox_info(&mut self, context : PowerboxCapability::GetPowerboxInfoContext) {
        context.done()
    }
}

impl UiView::Server for UiViewImpl {
    fn get_view_info(&mut self, context : UiView::GetViewInfoContext) {
        context.done()
    }

    fn new_session(&mut self, mut context : UiView::NewSessionContext) {
        println!("asked for a new session!");
        let (_, results) = context.get();

        let client : WebSession::Client = FromServer::new(None::<LocalClient>, ~WebSessionImpl::new());
        // we need to do this dance to upcast.
        results.set_session(UiSession::Client { client : client.client});

        context.done()
    }
}

pub struct WebSessionImpl {
    db : sqlite3::Database,
}

impl WebSessionImpl {
    pub fn new() -> WebSessionImpl {
        WebSessionImpl {
            db : sqlite3::open("/var/data.db").unwrap(),
        }
    }
}

impl UiSession::Server for WebSessionImpl {

}

struct Path {
    path : ~str,
    query : HashMap<~str, ~str>,
}

impl Path {
    fn new() -> Path {
        Path { path : ~"", query : HashMap::new() }
    }
}

fn parse_path(path : &str) -> Path {
    let mut result = Path::new();

    let v : ~[&str] = path.splitn('?', 2).collect();
    if v.len() == 0 {
        return result;
    }
    result.path = v[0].into_owned();
    if v.len() == 1 {
        return result;
    }
    for attr in v[1].split('&') {
        let a : ~[&str] = attr.splitn('=', 2).collect();
        if a.len() == 2 {
            result.query.insert(a[0].into_owned(),
                                a[1].into_owned());
        }
    }
    return result;
}

impl WebSessionImpl {

    fn checked<T>(&self, result : sqlite3::SqliteResult<T>) -> T {
        match result {
            Err(e) => {fail!("database error: {}, {:?}", e, self.db.get_errmsg()) }
            Ok(v) => { return v; }
        }
    }

    fn is_word(&self, word : &str) -> sqlite3::SqliteResult<bool> {

        if ! word.is_alphanumeric() { return Ok(false); }

        let cursor = try!(self.db.prepare(
            format!("SELECT * FROM Words WHERE Word = \"{}\";", word),
            &None));

        return Ok(try!(cursor.step_row()).is_some());
    }

    fn validate_def(&self, word : &str, definition : &[&str]) -> sqlite3::SqliteResult<bool> {

        if definition.len() != word.len() { return Ok(false); }

        let mut idx = 0;
        for &d in definition.iter() {
            if !(try!(self.is_word(d)) && d.len() > 0 && d[0] == word[idx]) {
                return Ok(false);
            }

            idx += 1;
        }

        return Ok(true);
    }

    fn write_def(&self, word : &str, definition : &[&str]) -> sqlite3::SqliteResult<()> {

        let time : i64 = ::time::get_time().sec;
        let mut query = StrBuf::new();
        query.push_str(format!("BEGIN; DELETE FROM Definitions WHERE Definee =\"{}\"; ", word));
        query.push_str("INSERT INTO Definitions(Definee, Idx, Definer) VALUES");
        let mut idx = 0;
        for &d in definition.iter() {
            if idx != 0 { query.push_str(","); }
            query.push_str(format!("(\"{}\", {}, \"{}\")", word, idx, d));
            idx += 1;
        }
        query.push_str(";");
        query.push_str(format!("INSERT Into Log(Word, Timestamp) VALUES(\"{}\",{});", word, time));
        query.push_str("COMMIT;");

        println!("query: {}", query);

        try!(self.db.exec(query.as_slice()));

        return Ok(());
    }

    fn get_def(&self, word : &str) -> sqlite3::SqliteResult<~str> {

        let cursor = try!(self.db.prepare(
            format!("SELECT * FROM Definitions WHERE Definee = \"{}\";", word),
            &None));

        let mut map = HashMap::<int, ~str>::new();

        loop {
            match try!(cursor.step_row()) {
                None => break,
                Some(row) => {
                    let definer = match row.get(&~"Definer") { &sqlite3::Text(ref t) => t.clone(), _ => fail!(), };
                    let idx = match row.get(&~"Idx") { &sqlite3::Integer(ref i) => i.clone(), _ => fail!(), };

                    map.insert(idx, definer);
                }
            }
        }

        if map.len() != word.len() {
            return Ok(~"
this word has no definition yet
");
        } else {

            let mut result = StrBuf::new();
            result.push_str("
");
            for idx in range::<int>(0, word.len() as int) {
                let definer : &str = map.get(&idx).as_slice();
                result.push_str(format!(" {word} ", word=definer));
            }
            result.push_str("
");
            return Ok(result.into_owned());
        }
    }

    fn count_defs(&self) -> sqlite3::SqliteResult<(int, int, Vec<~str>)> {
        let cursor = try!(self.db.prepare("SELECT COUNT(*) FROM Words;", &None));
        assert!(cursor.step() == sqlite3::SQLITE_ROW);
        let num_words = cursor.get_int(0);

        let cursor = try!(self.db.prepare("SELECT COUNT(*) FROM Definitions WHERE Idx = 0;", &None));
        assert!(cursor.step() == sqlite3::SQLITE_ROW);
        let defined_words = cursor.get_int(0);

        let mut recent_words = Vec::new();
        let cursor = try!(
            self.db.prepare("SELECT Word, Timestamp FROM Log ORDER BY Timestamp DESC LIMIT 5;", &None));
        loop {
            match try!(cursor.step_row()) {
                None => break,
                Some(row) => {
                    let word : ~str = match row.get(&~"Word") {&sqlite3::Text(ref t) => t.clone(), _ => fail!(),};
                    recent_words.push(word);
                }
            }
        }
        Ok((defined_words, num_words, recent_words))

    }

}

static main_css : &'static str =
    "body { font-family: Helvetica, Sans, Arial;
            font-size: medium;
             margin-left: auto;
             margin-right: auto;
             width: 600px;
             text-align: center;
     }
     div {
          padding-bottom: 10pt;
     }
    .word {
        text-align: center;
        font-size: 500%;
     }
     .err {
       font-size: 90%;
       color: #AA0000;
     }
     .title {
       text-align: center;
       font-size:500%;
     }
     ";

static header : &'static str =
  r#"<title> acronymy </title><link rel="stylesheet" type="text/css" href="main.css" >
 <meta http-equiv="Content-Type" content="text/html;charset=utf-8" >
  "#;

static lookup_form : &'static str =
      r#"<form action="define" method="get">
          <input name="word" maxlength="100"/><button>find word</button></form>"#;

fn define_form(word :&str) -> ~str {
       format!("<form action=\"define\" method=\"get\">
               <input name=\"word\" value=\"{word}\" type=\"hidden\"/>
               <input name=\"definition\" maxlength=\"2000\"/>
                   <button>submit definition</button></form>", word=word)
}

enum PageData<'a> {
    NoWord,
    WordAndDef(&'a str, &'a str, Option<&'a str>),
    HomePage(int, int, Vec<~str>),
}

fn construct_html(page_data : PageData) -> ~str {
    let mut result = StrBuf::new();
    result.push_str(format!("{}", header));

    static home_link : &'static str = "home";
    match page_data {
        NoWord => {
            result.push_str("
 that's not a word
");
            result.push_str(lookup_form);
            result.push_str(home_link);
        }
        WordAndDef(word, def_div, err) => {
            result.push_str(format!("
{}
", word));

            result.push_str(def_div);

            match err {
                None => {}
                Some(e) => {
                    result.push_str(format!("
{}
", e));
                }
            }

            result.push_str(define_form(word));
            result.push_str(home_link);
        }
        HomePage(num_defined, total, recent) => {
            result.push_str("
Acronymy
");
            result.push_str("
A user-editable, acronym-only dictionary.
");
            result.push_str(format!("
So far, we have defined {} out of {} words.
",
                                    num_defined, total));
            if recent.len() > 0 {
                result.push_str("
Recently modified words: ");
                let mut idx = 0;
                for w in recent.iter() {
                    if idx != 0 {
                        result.push_str(", ");
                    }
                    result.push_str(format!("{word}", word=*w));
                    idx += 1;
                }
                result.push_str(".
");
            }
            result.push_str(lookup_form);
        }
    }

    result.push_str("");
    result.into_owned()
}

impl WebSession::Server for WebSessionImpl {
    fn get(&mut self, mut context : WebSession::GetContext) {
        println!("GET");
        let (params, results) = context.get();
        let raw_path = params.get_path();
        let content = results.init_content();
        content.set_mime_type("text/html");

        let path = parse_path(raw_path);
        println!("path = {}", raw_path);

        if raw_path == "main.css" {
            content.get_body().set_bytes(main_css.as_bytes())
        } else if path.path.as_slice() == "define" {
            let word : ~str = match path.query.find(&~"word") {
                Some(w) if self.checked(self.is_word(*w)) => {
                    w.clone()
                }
                _ => {
                    content.get_body().set_bytes(construct_html(NoWord).as_bytes());
                    return context.done();
                }
            };

            match path.query.find(&~"definition") {
                None => {
                    let def_div = self.checked(self.get_def(word));

                    content.get_body().set_bytes(construct_html(WordAndDef(word.as_slice(),
                                                                           def_div.as_slice(),
                                                                           None)).as_bytes());
                }
                Some(def_query) => {

                    let definition : ~[&str] = def_query.split('+').collect();

                    if self.checked(self.validate_def(word, definition)) {

                        self.checked(self.write_def(word, definition));
                        let def_div = self.checked(self.get_def(word));
                        content.get_body().set_bytes(
                            construct_html(WordAndDef(word.as_slice(),
                                                      def_div.as_slice(),
                                                      None)).as_bytes());

                    } else {

                        let def_div = self.checked(self.get_def(word));

                        content.get_body().set_bytes(
                            construct_html(WordAndDef(word.as_slice(),
                                                      def_div.as_slice(),
                                                      Some("invalid definition"))).as_bytes());
                    }
                }
            }

        } else {
            let (num_defined, total, recent) = self.checked(self.count_defs());
            content.get_body().set_bytes(construct_html(HomePage(num_defined, total, recent)).as_bytes());
        }
        context.done()
    }
    fn post(&mut self, context : WebSession::PostContext) {
        println!("POST");
        context.done()
    }
    fn put(&mut self, context : WebSession::PutContext) {
        println!("PUT");
        context.done()
    }
    fn delete(&mut self, context : WebSession::DeleteContext) {
        println!("DELETE");
        context.done()
    }
    fn open_web_socket(&mut self, context : WebSession::OpenWebSocketContext) {
        println!("OPEN WEB SOCKET");
        context.done()
    }
}

pub struct FdStream {
    inner : ~::std::rt::rtio::RtioFileStream:Send,
}

impl FdStream {
    pub fn new(fd : ::libc::c_int) -> ::std::io::IoResult<FdStream> {
        ::std::rt::rtio::LocalIo::maybe_raise(|io| {
            Ok (FdStream { inner : io.fs_from_raw_fd(fd, ::std::rt::rtio::DontClose) })
        })
    }
}

impl Reader for FdStream {
    fn read(&mut self, buf : &mut [u8]) -> ::std::io::IoResult<uint> {
        self.inner.read(buf).map(|i| i as uint)
    }
}

impl Writer for FdStream {
    fn write(&mut self, buf : &[u8]) -> ::std::io::IoResult<()> {
        self.inner.write(buf)
    }
}

pub struct Restorer;

impl SturdyRefRestorer for Restorer {
    fn restore(&self, obj_id : AnyPointer::Reader) -> Option<~ClientHook:Send> {
        if obj_id.is_null() {
            let client : UiView::Client = FromServer::new(None::<LocalClient>, ~UiViewImpl);
            Some(client.client.hook)
        } else {
            None
        }
    }
}

pub fn main() -> ::std::io::IoResult<()> {

    let args = ::std::os::args();

    if args.len() == 4 && args[1].as_slice() == "--init" {
        println!("initializing...");
        let initdb_path = ::std::path::Path::new(args[2].as_slice());
        let proddb_path = ::std::path::Path::new(args[3].as_slice());
        println!("copying database from {} to {}", args[2], args[3]);
        try!(::std::io::fs::copy(&initdb_path, &proddb_path));
        println!("success!");
    }

    // sandstorm launches us with a connection file descriptor 3
    let ifs = try!(FdStream::new(3));
    let ofs = try!(FdStream::new(3));

    let connection_state = RpcConnectionState::new();
    connection_state.run(ifs, ofs, Restorer);

    Ok(())
}

<|endoftext|>


<commit_before>use util::io_err;
use std::io::{Write, Read, ErrorKind};
use std::io::Result as IoResult;
use std::net::{IpAddr, TcpStream};
use std::mem;

enum AuthMethod<'s> {
    NoAuth,
    UPass(&'s str, &'s str)
}

pub trait ReadU8 {
    fn read_u8(&mut self) -> IoResult<u8>;
}

impl ReadU8 for TcpStream {
    #[inline]
    fn read_u8(&mut self) -> IoResult<u8> {
        let mut buf = [0u8];
        try!(self.read_exact(&mut buf));
        Ok(buf[0])
    }
}

//pub trait WriteU8 {
//    fn read_u8(&mut self) -> IoResult<u8>;
//}
//
//impl ReadU8 for TcpStream {
//    #[inline]
//    fn read_u8(&mut self) -> IoResult<u8> {
//        let mut buf = [0u8];
//        try!(self.read_exact(&mut buf));
//        Ok(buf[0])
//    }
//}

pub trait Socks5Destination {
    fn write_destination(&self, &mut TcpStream) -> IoResult<()>;
}

impl<'a> Socks5Destination for &'a str {
    fn write_destination(&self, stream: &mut TcpStream) -> IoResult<()> {
        if self.len() > 255 {
            return io_err(ErrorKind::Other, "Domain length is too long");
        }

        try!(stream.write(&[0x03, self.len() as u8]));
        try!(stream.write(&self.as_bytes()));
        Ok(())
    }
}

impl<'a> Socks5Destination for IpAddr {
    fn write_destination(&self, stream: &mut TcpStream) -> IoResult<()> {
        match *self {
            //IpAddr::V4(a, b, c, d) => try!(stream.write_all(&[0x01, a, b, c, d])),
            IpAddr::V4(addr) => match addr.octets() {
                [a,b,c,d] => try!(stream.write_all(&[0x01, a, b, c, d]))
            },
            IpAddr::V6(addr) => match addr.segments() {
                [a,b,c,d,e,f,g,h] => {
                    try!(stream.write_all(&[0x04]));
                    for &pair in [a, b, c, d, e, f, g, h].iter() {
                        let mut aux : u16 = pair;
                        aux = aux.to_be();
                        let array : [u8; 2] = unsafe { mem::transmute(aux) };
                        try!(stream.write_all(&array));
                    }
                }
            }
        }

        Ok(())
    }
}

pub struct Socks5<'a> {
    socks_host: &'a str,
    socks_port: u16,
    socks_auth: AuthMethod<'a>,
}

impl<'a> Socks5<'a> {
    pub fn new(host: &'a str, port: u16) -> Socks5 {
        Socks5 {
            socks_host: host,
            socks_port: port,
            socks_auth: AuthMethod::NoAuth
        }
    }

    pub fn login(&mut self, uname: &'a str, passwd: &'a str) {
        self.socks_auth = AuthMethod::UPass(uname, passwd);
    }

    pub fn connect<T: Socks5Destination>(&mut self, destination: T, port: u16)
                                                        -> IoResult<TcpStream> {
        let mut stream = try!(TcpStream::connect((self.socks_host, self.socks_port)));
        try!(stream.write(&[0x05u8]));
        match self.socks_auth {
            AuthMethod::NoAuth => { try!(stream.write(&[0x01u8, 0x00])); },
            AuthMethod::UPass(..) => { try!(stream.write(&[0x01u8, 0x02])); }
        }

        if try!(stream.read_u8()) != 0x05 {
            return io_err(ErrorKind::Other, "Unexpected SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                match self.socks_auth {
                    AuthMethod::NoAuth => { /* Continue */ },
                    _ => return io_err(ErrorKind::Other,
                        "Wrong authentication method from server")
                }
            }
            0x02 => {
                match self.socks_auth {
                    AuthMethod::UPass(uname, passwd) => {
                        try!(stream.write_all(&[0x01u8, uname.len() as u8]));
                        //try!(stream.write_str(uname));
                        try!(write!(stream, "{}", uname));
                        try!(stream.write_all(&[passwd.len() as u8]));
                        try!(write!(stream, "{}", passwd));
                        //try!(stream.write_str(passwd));

                        if try!(stream.read_u8()) != 0x01 {
                            return io_err(ErrorKind::Other,
                                "Invalid authentication version");
                        }

                        if try!(stream.read_u8()) != 0x00 {
                            return io_err(ErrorKind::ConnectionRefused, "Authentication failed");
                        }
                    }
                    _ => { return io_err(ErrorKind::Other,
                            "Wrong authentication method from server");
                    }
                }
            }
            0xFF => { return io_err(ErrorKind::ConnectionRefused,
                      "Server refused authentication methods"); }
            _ => { return io_err(ErrorKind::Other,
                    "Wrong authentication method from server"); }
        }

        try!(stream.write(&[0x05u8, 0x01, 0x00]));
        try!(destination.write_destination(&mut stream));
        let be_port = port.to_be();
        let array : [u8; 2] = unsafe { mem::transmute(be_port) };
        try!(stream.write_all(&array));

        if try!(stream.read_u8()) != 0x05 {
            return io_err(ErrorKind::Other, "Invalid SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                let _null = try!(stream.read_u8());

                match try!(stream.read_u8()) {
                    0x01 => {
                        let mut _ipv4 = [0; 4];
                        try!(stream.read_exact(&mut _ipv4));
                    }
                    0x03 => {
                        let addrlen = try!(stream.read_u8());
                        let mut _domain = Vec::with_capacity(addrlen as usize);
                        try!(stream.read_exact(&mut _domain));
                    }
                    0x04 => {
                        let mut _ipv6 = [0; 16];
                        try!(stream.read_exact(&mut _ipv6));
                    }
                    _ => return io_err(ErrorKind::Other, "Invalid address type"),
                }

                let mut _port_slice : [u8; 2] = [0u8; 2];
                try!(stream.read_exact(&mut _port_slice));
                let _port_be : u16 = unsafe { mem::transmute(_port_slice) };
                let _port = u16::from_be(_port_be);
                Ok(stream)
            }
            0x01 => io_err(ErrorKind::Other, "General failure"),
            0x02 => io_err(ErrorKind::Other, "Connection not allowed by ruleset"),
            0x03 => io_err(ErrorKind::NotConnected, "Network unreachable"),
            0x04 => io_err(ErrorKind::Other, "Host unreachable"),
            0x05 => io_err(ErrorKind::ConnectionRefused, "Connection refused by destination"),
            0x06 => io_err(ErrorKind::Other, "TTL expired"),
            0x07 => io_err(ErrorKind::Other, "Protocol Error"),
            0x08 => io_err(ErrorKind::Other, "Address type not supported"),
            _ => io_err(ErrorKind::Other, "Unknown error")
        }
    }
}
<commit_msg>minor aesthetic improvements<commit_after>use util::io_err;
use std::io::{Write, Read, ErrorKind};
use std::io::Result as IoResult;
use std::net::{IpAddr, TcpStream};
use std::mem;

enum AuthMethod<'s> {
    NoAuth,
    UPass(&'s str, &'s str)
}

pub trait ReadU8 {
    fn read_u8(&mut self) -> IoResult<u8>;
}

impl ReadU8 for TcpStream {
    #[inline]
    fn read_u8(&mut self) -> IoResult<u8> {
        let mut buf = [0u8];
        try!(self.read_exact(&mut buf));
        Ok(buf[0])
    }
}

//pub trait WriteU8 {
//    fn read_u8(&mut self) -> IoResult<u8>;
//}
//
//impl ReadU8 for TcpStream {
//    #[inline]
//    fn read_u8(&mut self) -> IoResult<u8> {
//        let mut buf = [0u8];
//        try!(self.read_exact(&mut buf));
//        Ok(buf[0])
//    }
//}

pub trait Socks5Destination {
    fn write_destination(&self, &mut TcpStream) -> IoResult<()>;
}

impl<'a> Socks5Destination for &'a str {
    fn write_destination(&self, stream: &mut TcpStream) -> IoResult<()> {
        if self.len() > 255 {
            return io_err(ErrorKind::Other, "Domain length is too long");
        }

        try!(stream.write(&[0x03, self.len() as u8]));
        try!(stream.write(&self.as_bytes()));
        Ok(())
    }
}

impl<'a> Socks5Destination for IpAddr {
    fn write_destination(&self, stream: &mut TcpStream) -> IoResult<()> {
        match *self {
            //IpAddr::V4(a, b, c, d) => try!(stream.write_all(&[0x01, a, b, c, d])),
            IpAddr::V4(addr) => {
                try!(stream.write_all(&[0x01]));
                try!(stream.write_all(&addr.octets()));
            },
            IpAddr::V6(addr) => {
                let segments = addr.segments();
                try!(stream.write_all(&[0x04]));
                for &pair in segments.iter() {
                    let array : [u8; 2] = unsafe {
                        mem::transmute(pair.to_be())
                    };
                    try!(stream.write_all(&array));
                }
            }
        }

        Ok(())
    }
}

pub struct Socks5<'a> {
    socks_host: &'a str,
    socks_port: u16,
    socks_auth: AuthMethod<'a>,
}

impl<'a> Socks5<'a> {
    pub fn new(host: &'a str, port: u16) -> Socks5 {
        Socks5 {
            socks_host: host,
            socks_port: port,
            socks_auth: AuthMethod::NoAuth
        }
    }

    pub fn login(&mut self, uname: &'a str, passwd: &'a str) {
        self.socks_auth = AuthMethod::UPass(uname, passwd);
    }

    pub fn connect<T: Socks5Destination>(&mut self, destination: T, port: u16)
                                                        -> IoResult<TcpStream> {
        let mut stream = try!(TcpStream::connect((self.socks_host, self.socks_port)));
        try!(stream.write(&[0x05u8]));
        match self.socks_auth {
            AuthMethod::NoAuth => { try!(stream.write(&[0x01u8, 0x00])); },
            AuthMethod::UPass(..) => { try!(stream.write(&[0x01u8, 0x02])); }
        }

        if try!(stream.read_u8()) != 0x05 {
            return io_err(ErrorKind::Other, "Unexpected SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                match self.socks_auth {
                    AuthMethod::NoAuth => { /* Continue */ },
                    _ => return io_err(ErrorKind::Other,
                        "Wrong authentication method from server")
                }
            }
            0x02 => {
                match self.socks_auth {
                    AuthMethod::UPass(uname, passwd) => {
                        try!(stream.write_all(&[0x01u8, uname.len() as u8]));
                        //try!(stream.write_str(uname));
                        try!(write!(stream, "{}", uname));
                        try!(stream.write_all(&[passwd.len() as u8]));
                        try!(write!(stream, "{}", passwd));
                        //try!(stream.write_str(passwd));

                        if try!(stream.read_u8()) != 0x01 {
                            return io_err(ErrorKind::Other,
                                "Invalid authentication version");
                        }

                        if try!(stream.read_u8()) != 0x00 {
                            return io_err(ErrorKind::ConnectionRefused, "Authentication failed");
                        }
                    }
                    _ => { return io_err(ErrorKind::Other,
                            "Wrong authentication method from server");
                    }
                }
            }
            0xFF => { return io_err(ErrorKind::ConnectionRefused,
                      "Server refused authentication methods"); }
            _ => { return io_err(ErrorKind::Other,
                    "Wrong authentication method from server"); }
        }

        try!(stream.write(&[0x05u8, 0x01, 0x00]));
        try!(destination.write_destination(&mut stream));
        let be_port = port.to_be();
        let array : [u8; 2] = unsafe { mem::transmute(be_port) };
        try!(stream.write_all(&array));

        if try!(stream.read_u8()) != 0x05 {
            return io_err(ErrorKind::Other, "Invalid SOCKS version number");
        }

        match try!(stream.read_u8()) {
            0x00 => {
                let _null = try!(stream.read_u8());

                match try!(stream.read_u8()) {
                    0x01 => {
                        let mut _ipv4 = [0; 4];
                        try!(stream.read_exact(&mut _ipv4));
                    }
                    0x03 => {
                        let addrlen = try!(stream.read_u8());
                        let mut _domain = Vec::with_capacity(addrlen as usize);
                        try!(stream.read_exact(&mut _domain));
                    }
                    0x04 => {
                        let mut _ipv6 = [0; 16];
                        try!(stream.read_exact(&mut _ipv6));
                    }
                    _ => return io_err(ErrorKind::Other, "Invalid address type"),
                }

                let mut _port_slice : [u8; 2] = [0u8; 2];
                try!(stream.read_exact(&mut _port_slice));
                let _port_be : u16 = unsafe { mem::transmute(_port_slice) };
                let _port = u16::from_be(_port_be);
                Ok(stream)
            }
            0x01 => io_err(ErrorKind::Other, "General failure"),
            0x02 => io_err(ErrorKind::Other, "Connection not allowed by ruleset"),
            0x03 => io_err(ErrorKind::NotConnected, "Network unreachable"),
            0x04 => io_err(ErrorKind::Other, "Host unreachable"),
            0x05 => io_err(ErrorKind::ConnectionRefused, "Connection refused by destination"),
            0x06 => io_err(ErrorKind::Other, "TTL expired"),
            0x07 => io_err(ErrorKind::Other, "Protocol Error"),
            0x08 => io_err(ErrorKind::Other, "Address type not supported"),
            _ => io_err(ErrorKind::Other, "Unknown error")
        }
    }
}
<|endoftext|>