<commit_before>use nom::{
  be_u8, be_u16, be_u32, be_u64,
  le_u32,
  IResult,
  ErrorCode, Err,
  Needed,
};

use std::str::from_utf8;

use metadata::{
  Block, BlockData,
  StreamInfo, Application, VorbisComment, CueSheet, Picture,
  SeekPoint, CueSheetTrack, CueSheetTrackIndex, PictureType,
};

use utility::to_u32;

macro_rules! skip_bytes(
  ($input: expr, $length: expr) => (
    {
      match take!($input, $length) {
        IResult::Done(i, bytes)   => {
          let sum = bytes.iter().fold(0, |result, byte| result + byte);

          if sum == 0 {
            IResult::Done(i, bytes)
          } else {
            IResult::Error(Err::Position(ErrorCode::Digit as u32, $input))
          }
        }
        IResult::Error(error)     => IResult::Error(error),
        IResult::Incomplete(need) => IResult::Incomplete(need),
      }
    }
  );
);

named!(pub stream_info <&[u8], BlockData>,
  chain!(
    min_block_size: be_u16 ~
    max_block_size: be_u16 ~
    min_frame_size: map!(take!(3), to_u32) ~
    max_frame_size: map!(take!(3), to_u32) ~
    bytes: take!(8) ~
    md5_sum: take!(16),
    || {
      let sample_rate     = ((bytes[0] as u32) << 12) +
                            ((bytes[1] as u32) << 4)  +
                            (bytes[2] as u32) >> 4;
      let channels        = (bytes[2] >> 1) & 0b0111;
      let bits_per_sample = ((bytes[2] & 0b01) << 4) +
                            bytes[3] >> 4;
      let total_samples   = (((bytes[3] as u64) & 0x0f) << 32) +
                            ((bytes[4] as u64) << 24) +
                            ((bytes[5] as u64) << 16) +
                            ((bytes[6] as u64) << 8) +
                            (bytes[7] as u64);

      BlockData::StreamInfo(StreamInfo {
        min_block_size: min_block_size,
        max_block_size: max_block_size,
        min_frame_size: min_frame_size,
        max_frame_size: max_frame_size,
        sample_rate: sample_rate,
        channels: channels + 1,
        bits_per_sample: bits_per_sample + 1,
        total_samples: total_samples,
        md5_sum: md5_sum,
      })
    }
  )
);

fn padding(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  map!(input, skip_bytes!(length), |_| BlockData::Padding(0))
}

fn application(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  chain!(input,
    id: take_str!(4) ~
    data: take!(length - 4),
    || {
      BlockData::Application(Application {
        id: id,
        data: data,
      })
    }
  )
}

named!(seek_point <&[u8], SeekPoint>,
  chain!(
    sample_number: be_u64 ~
    stream_offset: be_u64 ~
    frame_samples: be_u16,
    || {
      SeekPoint {
        sample_number: sample_number,
        stream_offset: stream_offset,
        frame_samples: frame_samples,
      }
    }
  )
);

fn seek_table(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  let seek_count = (length / 18) as usize;

  map!(input, count!(seek_point, seek_count), BlockData::SeekTable)
}

named!(vorbis_comment <&[u8], BlockData>,
  chain!(
    vendor_string_length: le_u32 ~
    vendor_string: take_str!(vendor_string_length)  ~
    number_of_comments: le_u32 ~
    comments: count!(comment_field, number_of_comments as usize),
    || {
      BlockData::VorbisComment(VorbisComment {
        vendor_string: vendor_string,
        comments: comments,
      })
    }
  )
);

named!(comment_field <&[u8], &str>,
  chain!(
    comment_length: le_u32 ~
    comment: take_str!(comment_length),
    || { comment }
  )
);

named!(cue_sheet <&[u8], BlockData>,
  chain!(
    media_catalog_number: take_str!(128) ~
    lead_in: be_u64 ~
    bytes: take!(259) ~ // TODO: last (7 + 258 * 8) bits must be 0
    num_tracks: be_u8 ~
    tracks: count!(cue_sheet_track, num_tracks as usize),
    || {
      let is_cd = ((bytes[0] >> 7) & 0b01) == 1;

      BlockData::CueSheet(CueSheet {
        media_catalog_number: media_catalog_number,
        lead_in: lead_in,
        is_cd: is_cd,
        tracks: tracks,
      })
    }
  )
);

named!(cue_sheet_track <&[u8], CueSheetTrack>,
  chain!(
    offset: be_u64 ~
    number: be_u8 ~
    isrc: take_str!(12) ~
    bytes: take!(14) ~ // TODO: last (6 + 13 * 8) bits must be 0
    num_indices: be_u8 ~
    indices: cond!(
      num_indices != 0,
      count!(cue_sheet_track_index, num_indices as usize)
    ),
    || {
      let isnt_audio      = ((bytes[0] >> 7) & 0b01) == 1;
      let is_pre_emphasis = ((bytes[0] >> 6) & 0b01) == 1;

      CueSheetTrack {
        offset: offset,
        number: number,
        isrc: isrc,
        isnt_audio: isnt_audio,
        is_pre_emphasis: is_pre_emphasis,
        indices: indices.unwrap_or(Vec::new()),
      }
    }
  )
);

named!(cue_sheet_track_index <&[u8], CueSheetTrackIndex>,
  chain!(
    offset: be_u64 ~
    number: be_u8 ~
    skip_bytes!(3),
    || {
      CueSheetTrackIndex {
        offset: offset,
        number: number,
      }
    }
  )
);

named!(picture <&[u8], BlockData>,
  chain!(
    picture_type_num: be_u32 ~
    mime_type_length:  be_u32 ~
    mime_type: take_str!(mime_type_length) ~
    description_length: be_u32 ~
    description: take_str!(description_length) ~
    width: be_u32 ~
    height: be_u32 ~
    depth: be_u32 ~
    colors: be_u32 ~
    data_length: be_u32 ~
    data: take!(data_length),
    || {
      let picture_type = match picture_type_num {
        0  => PictureType::Other,
        1  => PictureType::FileIconStandard,
        2  => PictureType::FileIcon,
        3  => PictureType::FrontCover,
        4  => PictureType::BackCover,
        5  => PictureType::LeafletPage,
        6  => PictureType::Media,
        7  => PictureType::LeadArtist,
        8  => PictureType::Artist,
        9  => PictureType::Conductor,
        10 => PictureType::Band,
        11 => PictureType::Composer,
        12 => PictureType::Lyricist,
        13 => PictureType::RecordingLocation,
        14 => PictureType::DuringRecording,
        15 => PictureType::DuringPerformace,
        16 => PictureType::VideoScreenCapture,
        17 => PictureType::Fish,
        18 => PictureType::Illustration,
        19 => PictureType::BandLogoType,
        20 => PictureType::PublisherLogoType,
        _  => PictureType::Other,
      };

      BlockData::Picture(Picture {
        picture_type: picture_type,
        mime_type: mime_type,
        description: description,
        width: width,
        height: height,
        depth: depth,
        colors: colors,
        data: data,
      })
    }
  )
);

// As of FLAC v1.3.1, there is support for up to 127 different metadata
// `Block`s but actually 7 that are implemented. When the `Block` type isn't
// recognised, this block gets skipped over with this parser.
fn unknown(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  map!(input, take!(length), BlockData::Unknown)
}

named!(pub header <&[u8], (bool, u8, u32)>,
  chain!(
    block_byte: be_u8 ~
    length: map!(take!(3), to_u32),
    || {
      let is_last    = (block_byte >> 7) == 1;
      let block_type = block_byte & 0b01111111;

      (is_last, block_type, length)
    }
  )
);

fn block_data(input: &[u8], block_type: u8, length: u32)
              -> IResult<&[u8], BlockData> {
  match block_type {
    0       => stream_info(input),
    1       => padding(input, length),
    2       => application(input, length),
    3       => seek_table(input, length),
    4       => vorbis_comment(input),
    5       => cue_sheet(input),
    6       => picture(input),
    7...126 => unknown(input, length),
    _       => IResult::Error(Err::Position(ErrorCode::Alt as u32, input)),
  }
}

named!(block <&[u8], Block>,
  chain!(
    block_header: header ~
    data: apply!(block_data, block_header.1, block_header.2),
    || {
      Block {
        is_last: block_header.0,
        length: block_header.2,
        data: data
      }
    }
  )
);

/// Parses one or more metadata `Block`s
///
/// The first block should always be `StreamInfo` since that is the only
/// required `Block`. At the moment `many_blocks` doesn't check that
/// requirement.
pub fn many_blocks(input: &[u8]) -> IResult<&[u8], Vec<Block>> {
  let mut is_last   = false;
  let mut blocks    = Vec::new();
  let mut start     = 0;
  let mut remaining = input.len();

  while !is_last {
    match block(&input[start..]) {
      IResult::Done(i, block) => {
        let result_len = i.len();

        if result_len == input[start..].len() {
          break;
        }

        start    += remaining - result_len;
        remaining = result_len;
        is_last   = block.is_last;

        blocks.push(block);
      }
      _                       => break,
    }
  }

  if blocks.len() == 0 {
    IResult::Error(Err::Position(ErrorCode::Many1 as u32, input))
  } else if is_last {
    IResult::Done(&input[start..], blocks)
  } else {
    IResult::Incomplete(Needed::Unknown)
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use nom::IResult;

  #[test]
  fn test_header() {
    let inputs = [ [0x80, 0x00, 0x00, 0x22]
                 , [0x01, 0x00, 0x04, 0x00]
                 , [0x84, 0x00, 0x00, 0xf8]
                 ];

    assert!(header(&inputs[0]) == IResult::Done(&[], (true, 0, 34)),
            "Header Test #1");
    assert!(header(&inputs[1]) == IResult::Done(&[], (false, 1, 1024)),
            "Header Test #2");
    assert!(header(&inputs[2]) == IResult::Done(&[], (true, 4, 248)),
            "Header Test #3");
  }
}
<commit_msg>Add failing test for stream_info parser<commit_after>use nom::{
  be_u8, be_u16, be_u32, be_u64,
  le_u32,
  IResult,
  ErrorCode, Err,
  Needed,
};

use std::str::from_utf8;

use metadata::{
  Block, BlockData,
  StreamInfo, Application, VorbisComment, CueSheet, Picture,
  SeekPoint, CueSheetTrack, CueSheetTrackIndex, PictureType,
};

use utility::to_u32;

macro_rules! skip_bytes(
  ($input: expr, $length: expr) => (
    {
      match take!($input, $length) {
        IResult::Done(i, bytes)   => {
          let sum = bytes.iter().fold(0, |result, byte| result + byte);

          if sum == 0 {
            IResult::Done(i, bytes)
          } else {
            IResult::Error(Err::Position(ErrorCode::Digit as u32, $input))
          }
        }
        IResult::Error(error)     => IResult::Error(error),
        IResult::Incomplete(need) => IResult::Incomplete(need),
      }
    }
  );
);

named!(pub stream_info <&[u8], BlockData>,
  chain!(
    min_block_size: be_u16 ~
    max_block_size: be_u16 ~
    min_frame_size: map!(take!(3), to_u32) ~
    max_frame_size: map!(take!(3), to_u32) ~
    bytes: take!(8) ~
    md5_sum: take!(16),
    || {
      let sample_rate     = ((bytes[0] as u32) << 12) +
                            ((bytes[1] as u32) << 4)  +
                            (bytes[2] as u32) >> 4;
      let channels        = (bytes[2] >> 1) & 0b0111;
      let bits_per_sample = ((bytes[2] & 0b01) << 4) +
                            bytes[3] >> 4;
      let total_samples   = (((bytes[3] as u64) & 0x0f) << 32) +
                            ((bytes[4] as u64) << 24) +
                            ((bytes[5] as u64) << 16) +
                            ((bytes[6] as u64) << 8) +
                            (bytes[7] as u64);

      BlockData::StreamInfo(StreamInfo {
        min_block_size: min_block_size,
        max_block_size: max_block_size,
        min_frame_size: min_frame_size,
        max_frame_size: max_frame_size,
        sample_rate: sample_rate,
        channels: channels + 1,
        bits_per_sample: bits_per_sample + 1,
        total_samples: total_samples,
        md5_sum: md5_sum,
      })
    }
  )
);

fn padding(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  map!(input, skip_bytes!(length), |_| BlockData::Padding(0))
}

fn application(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  chain!(input,
    id: take_str!(4) ~
    data: take!(length - 4),
    || {
      BlockData::Application(Application {
        id: id,
        data: data,
      })
    }
  )
}

named!(seek_point <&[u8], SeekPoint>,
  chain!(
    sample_number: be_u64 ~
    stream_offset: be_u64 ~
    frame_samples: be_u16,
    || {
      SeekPoint {
        sample_number: sample_number,
        stream_offset: stream_offset,
        frame_samples: frame_samples,
      }
    }
  )
);

fn seek_table(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  let seek_count = (length / 18) as usize;

  map!(input, count!(seek_point, seek_count), BlockData::SeekTable)
}

named!(vorbis_comment <&[u8], BlockData>,
  chain!(
    vendor_string_length: le_u32 ~
    vendor_string: take_str!(vendor_string_length)  ~
    number_of_comments: le_u32 ~
    comments: count!(comment_field, number_of_comments as usize),
    || {
      BlockData::VorbisComment(VorbisComment {
        vendor_string: vendor_string,
        comments: comments,
      })
    }
  )
);

named!(comment_field <&[u8], &str>,
  chain!(
    comment_length: le_u32 ~
    comment: take_str!(comment_length),
    || { comment }
  )
);

named!(cue_sheet <&[u8], BlockData>,
  chain!(
    media_catalog_number: take_str!(128) ~
    lead_in: be_u64 ~
    bytes: take!(259) ~ // TODO: last (7 + 258 * 8) bits must be 0
    num_tracks: be_u8 ~
    tracks: count!(cue_sheet_track, num_tracks as usize),
    || {
      let is_cd = ((bytes[0] >> 7) & 0b01) == 1;

      BlockData::CueSheet(CueSheet {
        media_catalog_number: media_catalog_number,
        lead_in: lead_in,
        is_cd: is_cd,
        tracks: tracks,
      })
    }
  )
);

named!(cue_sheet_track <&[u8], CueSheetTrack>,
  chain!(
    offset: be_u64 ~
    number: be_u8 ~
    isrc: take_str!(12) ~
    bytes: take!(14) ~ // TODO: last (6 + 13 * 8) bits must be 0
    num_indices: be_u8 ~
    indices: cond!(
      num_indices != 0,
      count!(cue_sheet_track_index, num_indices as usize)
    ),
    || {
      let isnt_audio      = ((bytes[0] >> 7) & 0b01) == 1;
      let is_pre_emphasis = ((bytes[0] >> 6) & 0b01) == 1;

      CueSheetTrack {
        offset: offset,
        number: number,
        isrc: isrc,
        isnt_audio: isnt_audio,
        is_pre_emphasis: is_pre_emphasis,
        indices: indices.unwrap_or(Vec::new()),
      }
    }
  )
);

named!(cue_sheet_track_index <&[u8], CueSheetTrackIndex>,
  chain!(
    offset: be_u64 ~
    number: be_u8 ~
    skip_bytes!(3),
    || {
      CueSheetTrackIndex {
        offset: offset,
        number: number,
      }
    }
  )
);

named!(picture <&[u8], BlockData>,
  chain!(
    picture_type_num: be_u32 ~
    mime_type_length:  be_u32 ~
    mime_type: take_str!(mime_type_length) ~
    description_length: be_u32 ~
    description: take_str!(description_length) ~
    width: be_u32 ~
    height: be_u32 ~
    depth: be_u32 ~
    colors: be_u32 ~
    data_length: be_u32 ~
    data: take!(data_length),
    || {
      let picture_type = match picture_type_num {
        0  => PictureType::Other,
        1  => PictureType::FileIconStandard,
        2  => PictureType::FileIcon,
        3  => PictureType::FrontCover,
        4  => PictureType::BackCover,
        5  => PictureType::LeafletPage,
        6  => PictureType::Media,
        7  => PictureType::LeadArtist,
        8  => PictureType::Artist,
        9  => PictureType::Conductor,
        10 => PictureType::Band,
        11 => PictureType::Composer,
        12 => PictureType::Lyricist,
        13 => PictureType::RecordingLocation,
        14 => PictureType::DuringRecording,
        15 => PictureType::DuringPerformace,
        16 => PictureType::VideoScreenCapture,
        17 => PictureType::Fish,
        18 => PictureType::Illustration,
        19 => PictureType::BandLogoType,
        20 => PictureType::PublisherLogoType,
        _  => PictureType::Other,
      };

      BlockData::Picture(Picture {
        picture_type: picture_type,
        mime_type: mime_type,
        description: description,
        width: width,
        height: height,
        depth: depth,
        colors: colors,
        data: data,
      })
    }
  )
);

// As of FLAC v1.3.1, there is support for up to 127 different metadata
// `Block`s but actually 7 that are implemented. When the `Block` type isn't
// recognised, this block gets skipped over with this parser.
fn unknown(input: &[u8], length: u32) -> IResult<&[u8], BlockData> {
  map!(input, take!(length), BlockData::Unknown)
}

named!(pub header <&[u8], (bool, u8, u32)>,
  chain!(
    block_byte: be_u8 ~
    length: map!(take!(3), to_u32),
    || {
      let is_last    = (block_byte >> 7) == 1;
      let block_type = block_byte & 0b01111111;

      (is_last, block_type, length)
    }
  )
);

fn block_data(input: &[u8], block_type: u8, length: u32)
              -> IResult<&[u8], BlockData> {
  match block_type {
    0       => stream_info(input),
    1       => padding(input, length),
    2       => application(input, length),
    3       => seek_table(input, length),
    4       => vorbis_comment(input),
    5       => cue_sheet(input),
    6       => picture(input),
    7...126 => unknown(input, length),
    _       => IResult::Error(Err::Position(ErrorCode::Alt as u32, input)),
  }
}

named!(block <&[u8], Block>,
  chain!(
    block_header: header ~
    data: apply!(block_data, block_header.1, block_header.2),
    || {
      Block {
        is_last: block_header.0,
        length: block_header.2,
        data: data
      }
    }
  )
);

/// Parses one or more metadata `Block`s
///
/// The first block should always be `StreamInfo` since that is the only
/// required `Block`. At the moment `many_blocks` doesn't check that
/// requirement.
pub fn many_blocks(input: &[u8]) -> IResult<&[u8], Vec<Block>> {
  let mut is_last   = false;
  let mut blocks    = Vec::new();
  let mut start     = 0;
  let mut remaining = input.len();

  while !is_last {
    match block(&input[start..]) {
      IResult::Done(i, block) => {
        let result_len = i.len();

        if result_len == input[start..].len() {
          break;
        }

        start    += remaining - result_len;
        remaining = result_len;
        is_last   = block.is_last;

        blocks.push(block);
      }
      _                       => break,
    }
  }

  if blocks.len() == 0 {
    IResult::Error(Err::Position(ErrorCode::Many1 as u32, input))
  } else if is_last {
    IResult::Done(&input[start..], blocks)
  } else {
    IResult::Incomplete(Needed::Unknown)
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use metadata::{
    BlockData,
    StreamInfo,
  };
  use nom::IResult;

  #[test]
  fn test_header() {
    let inputs = [ [0x80, 0x00, 0x00, 0x22]
                 , [0x01, 0x00, 0x04, 0x00]
                 , [0x84, 0x00, 0x00, 0xf8]
                 ];

    assert!(header(&inputs[0]) == IResult::Done(&[], (true, 0, 34)),
            "Header Test #1");
    assert!(header(&inputs[1]) == IResult::Done(&[], (false, 1, 1024)),
            "Header Test #2");
    assert!(header(&inputs[2]) == IResult::Done(&[], (true, 4, 248)),
            "Header Test #3");
  }

  #[test]
  fn test_stream_info() {
    let input   = [ 0x12, 0x00, 0x12, 0x00, 0x00, 0x00, 0x0e, 0x00, 0x00, 0x10
                  , 0x01, 0xf4, 0x02, 0x70, 0x00, 0x01, 0x38, 0x80, 0xa0, 0x42
                  , 0x23, 0x7c, 0x54, 0x93, 0xfd, 0xb9, 0x65, 0x6b, 0x94, 0xa8
                  , 0x36, 0x08, 0xd1, 0x1a
                  ];

    let md5_sum = [ 0xa0, 0x42, 0x23, 0x7c, 0x54, 0x93, 0xfd, 0xb9, 0x65, 0x6b
                  , 0x94, 0xa8, 0x36, 0x08, 0xd1, 0x1a
                  ];

    let result = BlockData::StreamInfo(StreamInfo {
      min_block_size: 4608,
      max_block_size: 4608,
      min_frame_size: 14,
      max_frame_size: 16,
      sample_rate: 8000,
      channels: 2,
      bits_per_sample: 8,
      total_samples: 80000,
      md5_sum: &md5_sum,
    });

    assert_eq!(stream_info(&input), IResult::Done(&[][..], result));
  }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use anyhow::Result;
use libra_logger::info;
use libradb::LibraDB;
use std::path::PathBuf;
use storage_interface::DbReader;
use transaction_builder::get_transaction_name;

use libra_types::{account_address::AccountAddress, account_config::AccountResource};
use std::convert::TryFrom;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    #[structopt(long, parse(from_os_str))]
    db: PathBuf,

    #[structopt(subcommand)] // Note that we mark a field as a subcommand
    cmd: Option<Command>,
}

#[derive(Debug, StructOpt)]
enum Command {
    #[structopt(name = "list-txns")]
    ListTXNs,
    #[structopt(name = "print-txn")]
    PrintTXN { version: u64 },
    #[structopt(name = "print-account")]
    PrintAccount {
        #[structopt(parse(try_from_str))]
        address: AccountAddress,
    },
}

/// Print out latest information stored in the DB.
fn print_head(db: &LibraDB) -> Result<()> {
    let si = db
        .get_startup_info()
        .expect("Can't get startup info")
        .expect("StartupInfo is empty, database is empty.");

    let version = si.latest_ledger_info.ledger_info().version();
    info!("Version: {}", version);

    info!(
        "The latest ledger info: {}",
        si.latest_ledger_info.ledger_info()
    );

    info!("Signatures: {:?}", si.latest_ledger_info.signatures());

    info!(
        "Current Validator Set: {}",
        si.latest_validator_set
            .as_ref()
            .expect("Unable to determine validator set, DB incorrect")
    );

    let backup = db.get_backup_handler();
    let iter = backup.get_account_iter(version)?;
    let num_account_state = iter.count();
    info!("Total Accounts: {}", num_account_state);

    print_txn(db, version);

    Ok(())
}

fn print_txn(db: &LibraDB, version: u64) {
    let tx = db
        .get_transaction_with_proof(version, version, false)
        .expect("Unable to load latest TXN");
    println!(
        "Transaction {}: {}",
        version,
        tx.transaction.format_for_client(get_transaction_name)
    );
}

fn print_account(db: &LibraDB, addr: AccountAddress) {
    let maybe_blob = db
        .get_latest_account_state(addr)
        .expect("Unable to read AccountState");
    if let Some(blob) = maybe_blob {
        match AccountResource::try_from(&blob) {
            Ok(r) => {
                println!("Account {}: {:?}", addr, r);
            }
            Err(e) => {
                info!(
                    "Account {} exists, but have no AccountResource: {}.",
                    addr, e
                );
            }
        }
    } else {
        info!("Account {} doesn't exists", addr);
    }
}

fn list_txns(db: &LibraDB) {
    let version = db
        .get_latest_version()
        .expect("Unable to get latest version");
    let backup = db.get_backup_handler();
    let iter = backup
        .get_transaction_iter(0, version)
        .expect("Unable to get txn iter");
    for (v, tx) in iter.enumerate() {
        println!(
            "TXN {}: {}",
            v,
            tx.expect("Unable to read TX")
                .format_for_client(get_transaction_name)
        );
    }
}

fn main() {
    ::libra_logger::Logger::new().init();

    let opt = Opt::from_args();

    let p = opt.db.as_path();

    if !p.is_dir() {
        info!("Invalid Directory {:?}!", p);
        std::process::exit(-1);
    }

    let log_dir = tempfile::tempdir().expect("Unable to get temp dir");
    info!("Opening DB at: {:?}, log at {:?}", p, log_dir.path());

    let db = LibraDB::open(p, true, Some(log_dir.path())).expect("Unable to open LibraDB");
    info!("DB opened successfully.");

    if let Some(cmd) = opt.cmd {
        match cmd {
            Command::ListTXNs => {
                list_txns(&db);
            }
            Command::PrintTXN { version } => {
                print_txn(&db, version);
            }
            Command::PrintAccount { address } => {
                print_account(&db, address);
            }
        }
    } else {
        print_head(&db).expect("Unable to read information from DB");

        Opt::clap().print_help().unwrap();
        println!();
    }
}
<commit_msg>[storage] add list-accounts option for LibraDB storage inspector<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use anyhow::Result;
use libra_logger::info;
use libradb::LibraDB;
use std::path::PathBuf;
use storage_interface::DbReader;
use transaction_builder::get_transaction_name;

use libra_types::{
    account_address::AccountAddress, account_config::AccountResource, account_state::AccountState,
};
use std::convert::TryFrom;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
struct Opt {
    #[structopt(long, parse(from_os_str))]
    db: PathBuf,

    #[structopt(subcommand)] // Note that we mark a field as a subcommand
    cmd: Option<Command>,
}

#[derive(Debug, StructOpt)]
enum Command {
    #[structopt(name = "list-txns")]
    ListTXNs,
    #[structopt(name = "print-txn")]
    PrintTXN { version: u64 },
    #[structopt(name = "print-account")]
    PrintAccount {
        #[structopt(parse(try_from_str))]
        address: AccountAddress,
    },
    #[structopt(name = "list-accounts")]
    ListAccounts,
}

/// Print out latest information stored in the DB.
fn print_head(db: &LibraDB) -> Result<()> {
    let si = db
        .get_startup_info()
        .expect("Can't get startup info")
        .expect("StartupInfo is empty, database is empty.");

    let version = si.latest_ledger_info.ledger_info().version();
    info!("Version: {}", version);

    info!(
        "The latest ledger info: {}",
        si.latest_ledger_info.ledger_info()
    );

    info!("Signatures: {:?}", si.latest_ledger_info.signatures());

    info!(
        "Current Validator Set: {}",
        si.latest_validator_set
            .as_ref()
            .expect("Unable to determine validator set, DB incorrect")
    );

    let backup = db.get_backup_handler();
    let iter = backup.get_account_iter(version)?;
    let num_account_state = iter.count();
    info!("Total Accounts: {}", num_account_state);

    print_txn(db, version);

    Ok(())
}

fn print_txn(db: &LibraDB, version: u64) {
    let tx = db
        .get_transaction_with_proof(version, version, false)
        .expect("Unable to load latest TXN");
    println!(
        "Transaction {}: {}",
        version,
        tx.transaction.format_for_client(get_transaction_name)
    );
}

fn print_account(db: &LibraDB, addr: AccountAddress) {
    let maybe_blob = db
        .get_latest_account_state(addr)
        .expect("Unable to read AccountState");
    if let Some(blob) = maybe_blob {
        match AccountResource::try_from(&blob) {
            Ok(r) => {
                println!("Account {}: {:?}", addr, r);
            }
            Err(e) => {
                info!(
                    "Account {} exists, but have no AccountResource: {}.",
                    addr, e
                );
            }
        }
    } else {
        info!("Account {} doesn't exists", addr);
    }
}

fn list_txns(db: &LibraDB) {
    let version = db
        .get_latest_version()
        .expect("Unable to get latest version");
    let backup = db.get_backup_handler();
    let iter = backup
        .get_transaction_iter(0, version)
        .expect("Unable to get txn iter");
    for (v, tx) in iter.enumerate() {
        println!(
            "TXN {}: {}",
            v,
            tx.expect("Unable to read TX")
                .format_for_client(get_transaction_name)
        );
    }
}

fn list_accounts(db: &LibraDB) {
    let version = db
        .get_latest_version()
        .expect("Unable to get latest version");
    let backup = db.get_backup_handler();
    let iter = backup
        .get_account_iter(version)
        .expect("Unagle to get account iter");
    let mut num_account = 0;
    for res in iter {
        match res {
            Ok((_, blob)) => {
                let accs = AccountState::try_from(&blob).expect("Failed to read AccountState");
                let addr = accs
                    .get_account_address()
                    .expect("Could not get address from state");
                match addr {
                    Some(x) => {
                        num_account += 1;
                        println!("Address: {:?}", x);
                    }
                    None => println!("Skipping: No address for AccountState: {:?}", accs),
                }
            }
            Err(x) => println!("Got err iterating through AccountStateBlobs {:?}", x),
        }
    }
    info!("Total Accounts: {}", num_account);
}

fn main() {
    ::libra_logger::Logger::new().init();

    let opt = Opt::from_args();

    let p = opt.db.as_path();

    if !p.is_dir() {
        info!("Invalid Directory {:?}!", p);
        std::process::exit(-1);
    }

    let log_dir = tempfile::tempdir().expect("Unable to get temp dir");
    info!("Opening DB at: {:?}, log at {:?}", p, log_dir.path());

    let db = LibraDB::open(p, true, Some(log_dir.path())).expect("Unable to open LibraDB");
    info!("DB opened successfully.");

    if let Some(cmd) = opt.cmd {
        match cmd {
            Command::ListTXNs => {
                list_txns(&db);
            }
            Command::PrintTXN { version } => {
                print_txn(&db, version);
            }
            Command::PrintAccount { address } => {
                print_account(&db, address);
            }
            Command::ListAccounts => {
                list_accounts(&db);
            }
        }
    } else {
        print_head(&db).expect("Unable to read information from DB");

        Opt::clap().print_help().unwrap();
        println!();
    }
}
<|endoftext|>


<commit_before><commit_msg>- added remainign url params, documentation - updated test case to account for extra url params<commit_after><|endoftext|>


<commit_before><commit_msg>Return error when Le calculation overflows<commit_after><|endoftext|>


<commit_before>// Rust Bitcoin Library
// Written in 2014 by
//   Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Network message
//!
//! This module defines the `Message` traits which are used
//! for (de)serializing Bitcoin objects for transmission on the network. It
//! also defines (de)serialization routines for many primitives.
//!

use collections::Vec;
use std::io::{IoError, IoResult, OtherIoError};
use std::io::MemReader;

use blockdata::block;
use network::address::Address;
use network::message_network;
use network::message_blockdata;
use network::encodable::{ConsensusDecodable, ConsensusEncodable};
use network::encodable::CheckedData;
use network::serialize::{serialize, RawDecoder, SimpleEncoder, SimpleDecoder};
use util::misc::prepend_err;

/// Serializer for command string
#[deriving(PartialEq, Clone, Show)]
pub struct CommandString(pub String);

impl<S:SimpleEncoder<E>, E> ConsensusEncodable<S, E> for CommandString {
  #[inline]
  fn consensus_encode(&self, s: &mut S) -> Result<(), E> {
    let &CommandString(ref inner_str) = self;
    let mut rawbytes = [0u8, ..12];
    rawbytes.clone_from_slice(inner_str.as_bytes().as_slice());
    rawbytes.consensus_encode(s)
  }
}

impl<D:SimpleDecoder<E>, E> ConsensusDecodable<D, E> for CommandString {
  #[inline]
  fn consensus_decode(d: &mut D) -> Result<CommandString, E> {
    let rawbytes: [u8, ..12] = try!(ConsensusDecodable::consensus_decode(d));
    let rv: String = FromIterator::from_iter(rawbytes.iter().filter_map(|&u| if u > 0 { Some(u as char) } else { None }));
    Ok(CommandString(rv))
  }
}

/// A Network message
pub struct RawNetworkMessage {
  /// Magic bytes to identify the network these messages are meant for
  pub magic: u32,
  /// The actual message data
  pub payload: NetworkMessage
}

#[deriving(Show)]
/// A Network message payload. Proper documentation is available on the Bitcoin
/// wiki https://en.bitcoin.it/wiki/Protocol_specification
pub enum NetworkMessage{
  /// `version`
  Version(message_network::VersionMessage),
  /// `verack`
  Verack,
  /// `addr`
  Addr(Vec<(u32, Address)>),
  /// `inv`
  Inv(Vec<message_blockdata::Inventory>),
  /// `getdata`
  GetData(Vec<message_blockdata::Inventory>),
  /// `notfound`
  NotFound(Vec<message_blockdata::Inventory>),
  /// `getblocks`
  GetBlocks(message_blockdata::GetBlocksMessage),
  /// `getheaders`
  GetHeaders(message_blockdata::GetHeadersMessage),
  // TODO: tx,
  /// `block`
  Block(block::Block),
  /// `headers`
  Headers(Vec<block::LoneBlockHeader>),
  // TODO: getaddr,
  // TODO: mempool,
  // TODO: checkorder,
  // TODO: submitorder,
  // TODO: reply,
  /// `ping`
  Ping(u64),
  /// `pong`
  Pong(u64),
  // TODO: reject,
  // TODO: bloom filtering
  // TODO: alert
}

impl RawNetworkMessage {
  fn command(&self) -> String {
    match self.payload {
      Version(_) => "version",
      Verack     => "verack",
      Addr(_)    => "addr",
      Inv(_)     => "inv",
      GetData(_) => "getdata",
      NotFound(_) => "notfound",
      GetBlocks(_) => "getblocks",
      GetHeaders(_) => "getheaders",
      Block(_)   => "block",
      Headers(_) => "headers",
      Ping(_)    => "ping",
      Pong(_)    => "pong"
    }.to_string()
  }
}

impl<S:SimpleEncoder<E>, E> ConsensusEncodable<S, E> for RawNetworkMessage {
  fn consensus_encode(&self, s: &mut S) -> Result<(), E> {
    try!(self.magic.consensus_encode(s));
    try!(CommandString(self.command()).consensus_encode(s));
    try!(CheckedData(match self.payload {
      Version(ref dat) => serialize(dat),
      Verack           => Ok(vec![]),
      Addr(ref dat)    => serialize(dat),
      Inv(ref dat)     => serialize(dat),
      GetData(ref dat) => serialize(dat),
      NotFound(ref dat) => serialize(dat),
      GetBlocks(ref dat) => serialize(dat),
      GetHeaders(ref dat) => serialize(dat),
      Block(ref dat)   => serialize(dat),
      Headers(ref dat) => serialize(dat),
      Ping(ref dat)    => serialize(dat),
      Pong(ref dat)    => serialize(dat),
    }.unwrap()).consensus_encode(s));
    Ok(())
  }
}

impl<D:SimpleDecoder<IoError>> ConsensusDecodable<D, IoError> for RawNetworkMessage {
  fn consensus_decode(d: &mut D) -> IoResult<RawNetworkMessage> {
    let magic = try!(ConsensusDecodable::consensus_decode(d));
    let CommandString(cmd): CommandString= try!(ConsensusDecodable::consensus_decode(d));
    let CheckedData(raw_payload): CheckedData = try!(ConsensusDecodable::consensus_decode(d));

    let mut mem_d = RawDecoder::new(MemReader::new(raw_payload));
    let payload = match cmd.as_slice() {
      "version" => Version(try!(prepend_err("version", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "verack"  => Verack,
      "addr"    => Addr(try!(prepend_err("addr", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "inv"     => Inv(try!(prepend_err("inv", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "getdata" => GetData(try!(prepend_err("getdata", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "notfound" => NotFound(try!(prepend_err("notfound", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "getblocks" => GetBlocks(try!(prepend_err("getblocks", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "getheaders" => GetHeaders(try!(prepend_err("getheaders", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "block"   => Block(try!(prepend_err("block", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "headers" => Headers(try!(prepend_err("headers", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "ping"    => Ping(try!(prepend_err("ping", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "pong"    => Ping(try!(prepend_err("pong", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      cmd => {
        return Err(IoError {
                     kind: OtherIoError,
                     desc: "unknown message type",
                     detail: Some(format!("`{}` not recognized", cmd))
                  });
      }
    };
    Ok(RawNetworkMessage {
      magic: magic,
      payload: payload
    })
  }
}

#[cfg(test)]
mod test {
  use super::{RawNetworkMessage, CommandString, Verack, Ping};

  use std::io::IoResult;

  use network::serialize::{deserialize, serialize};

  #[test]
  fn serialize_commandstring_test() {
    let cs = CommandString(String::from_str("Andrew"));
    assert_eq!(serialize(&cs), Ok(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]));
  }

  #[test]
  fn deserialize_commandstring_test() {
    let cs: IoResult<CommandString> = deserialize(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]);
    assert!(cs.is_ok());
    assert_eq!(cs.unwrap(), CommandString(String::from_str("Andrew")));

    let short_cs: IoResult<CommandString> = deserialize(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0]);
    assert!(short_cs.is_err());
  }

  #[test]
  fn serialize_verack_test() {
    assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: Verack }),
               Ok(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x76, 0x65, 0x72, 0x61,
                       0x63, 0x6B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                       0x00, 0x00, 0x00, 0x00, 0x5d, 0xf6, 0xe0, 0xe2]));
  }

  #[test]
  fn serialize_ping_test() {
    assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: Ping(100) }),
               Ok(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x70, 0x69, 0x6e, 0x67,
                       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                       0x08, 0x00, 0x00, 0x00, 0x24, 0x67, 0xf1, 0x1d,
                       0x64, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]));
  }
}

<commit_msg>Add support for `tx` messages<commit_after>// Rust Bitcoin Library
// Written in 2014 by
//   Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Network message
//!
//! This module defines the `Message` traits which are used
//! for (de)serializing Bitcoin objects for transmission on the network. It
//! also defines (de)serialization routines for many primitives.
//!

use collections::Vec;
use std::io::{IoError, IoResult, OtherIoError};
use std::io::MemReader;

use blockdata::block;
use blockdata::transaction;
use network::address::Address;
use network::message_network;
use network::message_blockdata;
use network::encodable::{ConsensusDecodable, ConsensusEncodable};
use network::encodable::CheckedData;
use network::serialize::{serialize, RawDecoder, SimpleEncoder, SimpleDecoder};
use util::misc::prepend_err;

/// Serializer for command string
#[deriving(PartialEq, Clone, Show)]
pub struct CommandString(pub String);

impl<S:SimpleEncoder<E>, E> ConsensusEncodable<S, E> for CommandString {
  #[inline]
  fn consensus_encode(&self, s: &mut S) -> Result<(), E> {
    let &CommandString(ref inner_str) = self;
    let mut rawbytes = [0u8, ..12];
    rawbytes.clone_from_slice(inner_str.as_bytes().as_slice());
    rawbytes.consensus_encode(s)
  }
}

impl<D:SimpleDecoder<E>, E> ConsensusDecodable<D, E> for CommandString {
  #[inline]
  fn consensus_decode(d: &mut D) -> Result<CommandString, E> {
    let rawbytes: [u8, ..12] = try!(ConsensusDecodable::consensus_decode(d));
    let rv: String = FromIterator::from_iter(rawbytes.iter().filter_map(|&u| if u > 0 { Some(u as char) } else { None }));
    Ok(CommandString(rv))
  }
}

/// A Network message
pub struct RawNetworkMessage {
  /// Magic bytes to identify the network these messages are meant for
  pub magic: u32,
  /// The actual message data
  pub payload: NetworkMessage
}

#[deriving(Show)]
/// A Network message payload. Proper documentation is available on the Bitcoin
/// wiki https://en.bitcoin.it/wiki/Protocol_specification
pub enum NetworkMessage{
  /// `version`
  Version(message_network::VersionMessage),
  /// `verack`
  Verack,
  /// `addr`
  Addr(Vec<(u32, Address)>),
  /// `inv`
  Inv(Vec<message_blockdata::Inventory>),
  /// `getdata`
  GetData(Vec<message_blockdata::Inventory>),
  /// `notfound`
  NotFound(Vec<message_blockdata::Inventory>),
  /// `getblocks`
  GetBlocks(message_blockdata::GetBlocksMessage),
  /// `getheaders`
  GetHeaders(message_blockdata::GetHeadersMessage),
  /// tx
  Tx(transaction::Transaction),
  /// `block`
  Block(block::Block),
  /// `headers`
  Headers(Vec<block::LoneBlockHeader>),
  // TODO: getaddr,
  // TODO: mempool,
  // TODO: checkorder,
  // TODO: submitorder,
  // TODO: reply,
  /// `ping`
  Ping(u64),
  /// `pong`
  Pong(u64),
  // TODO: reject,
  // TODO: bloom filtering
  // TODO: alert
}

impl RawNetworkMessage {
  fn command(&self) -> String {
    match self.payload {
      Version(_) => "version",
      Verack     => "verack",
      Addr(_)    => "addr",
      Inv(_)     => "inv",
      GetData(_) => "getdata",
      NotFound(_) => "notfound",
      GetBlocks(_) => "getblocks",
      GetHeaders(_) => "getheaders",
      Tx(_)      => "tx",
      Block(_)   => "block",
      Headers(_) => "headers",
      Ping(_)    => "ping",
      Pong(_)    => "pong"
    }.to_string()
  }
}

impl<S:SimpleEncoder<E>, E> ConsensusEncodable<S, E> for RawNetworkMessage {
  fn consensus_encode(&self, s: &mut S) -> Result<(), E> {
    try!(self.magic.consensus_encode(s));
    try!(CommandString(self.command()).consensus_encode(s));
    try!(CheckedData(match self.payload {
      Version(ref dat) => serialize(dat),
      Verack           => Ok(vec![]),
      Addr(ref dat)    => serialize(dat),
      Inv(ref dat)     => serialize(dat),
      GetData(ref dat) => serialize(dat),
      NotFound(ref dat) => serialize(dat),
      GetBlocks(ref dat) => serialize(dat),
      GetHeaders(ref dat) => serialize(dat),
      Tx(ref dat)      => serialize(dat),
      Block(ref dat)   => serialize(dat),
      Headers(ref dat) => serialize(dat),
      Ping(ref dat)    => serialize(dat),
      Pong(ref dat)    => serialize(dat),
    }.unwrap()).consensus_encode(s));
    Ok(())
  }
}

impl<D:SimpleDecoder<IoError>> ConsensusDecodable<D, IoError> for RawNetworkMessage {
  fn consensus_decode(d: &mut D) -> IoResult<RawNetworkMessage> {
    let magic = try!(ConsensusDecodable::consensus_decode(d));
    let CommandString(cmd): CommandString= try!(ConsensusDecodable::consensus_decode(d));
    let CheckedData(raw_payload): CheckedData = try!(ConsensusDecodable::consensus_decode(d));

    let mut mem_d = RawDecoder::new(MemReader::new(raw_payload));
    let payload = match cmd.as_slice() {
      "version" => Version(try!(prepend_err("version", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "verack"  => Verack,
      "addr"    => Addr(try!(prepend_err("addr", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "inv"     => Inv(try!(prepend_err("inv", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "getdata" => GetData(try!(prepend_err("getdata", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "notfound" => NotFound(try!(prepend_err("notfound", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "getblocks" => GetBlocks(try!(prepend_err("getblocks", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "getheaders" => GetHeaders(try!(prepend_err("getheaders", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "block"   => Block(try!(prepend_err("block", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "headers" => Headers(try!(prepend_err("headers", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "ping"    => Ping(try!(prepend_err("ping", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "pong"    => Ping(try!(prepend_err("pong", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      "tx"      => Tx(try!(prepend_err("tx", ConsensusDecodable::consensus_decode(&mut mem_d)))),
      cmd => {
        return Err(IoError {
                     kind: OtherIoError,
                     desc: "unknown message type",
                     detail: Some(format!("`{}` not recognized", cmd))
                  });
      }
    };
    Ok(RawNetworkMessage {
      magic: magic,
      payload: payload
    })
  }
}

#[cfg(test)]
mod test {
  use super::{RawNetworkMessage, CommandString, Verack, Ping};

  use std::io::IoResult;

  use network::serialize::{deserialize, serialize};

  #[test]
  fn serialize_commandstring_test() {
    let cs = CommandString(String::from_str("Andrew"));
    assert_eq!(serialize(&cs), Ok(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]));
  }

  #[test]
  fn deserialize_commandstring_test() {
    let cs: IoResult<CommandString> = deserialize(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]);
    assert!(cs.is_ok());
    assert_eq!(cs.unwrap(), CommandString(String::from_str("Andrew")));

    let short_cs: IoResult<CommandString> = deserialize(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0]);
    assert!(short_cs.is_err());
  }

  #[test]
  fn serialize_verack_test() {
    assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: Verack }),
               Ok(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x76, 0x65, 0x72, 0x61,
                       0x63, 0x6B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                       0x00, 0x00, 0x00, 0x00, 0x5d, 0xf6, 0xe0, 0xe2]));
  }

  #[test]
  fn serialize_ping_test() {
    assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: Ping(100) }),
               Ok(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x70, 0x69, 0x6e, 0x67,
                       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                       0x08, 0x00, 0x00, 0x00, 0x24, 0x67, 0xf1, 0x1d,
                       0x64, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]));
  }
}

<|endoftext|>


<commit_before>

#[abi = "cdecl"]
native mod rustrt {
    fn unsupervise();
}

#[abi = "cdecl"]
#[nolink]
native mod bar { }

#[abi = "cdecl"]
#[nolink]
native mod zed { }

#[abi = "cdecl"]
#[nolink]
native mod libc {
    fn write(fd: int, buf: *u8, count: uint) -> int;
}

#[abi = "cdecl"]
#[nolink]
native mod baz { }

fn main(args: [str]) { }
<commit_msg>rewrite to use ctypes types<commit_after>

#[abi = "cdecl"]
native mod rustrt {
    fn unsupervise();
}

#[abi = "cdecl"]
#[nolink]
native mod bar { }

#[abi = "cdecl"]
#[nolink]
native mod zed { }

#[abi = "cdecl"]
#[nolink]
native mod libc {
    fn write(fd: int, buf: *u8, count: ctypes::size_t) -> ctypes::ssize_t;
}

#[abi = "cdecl"]
#[nolink]
native mod baz { }

fn main(args: [str]) { }
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use hir::ModuleSource;
use ra_db::{RelativePath, RelativePathBuf, SourceDatabase, SourceDatabaseExt};
use ra_ide_db::RootDatabase;
use ra_syntax::{
    algo::find_node_at_offset, ast, lex_single_valid_syntax_kind, AstNode, SyntaxKind, SyntaxNode,
};
use ra_text_edit::TextEdit;

use crate::{
    FileId, FilePosition, FileSystemEdit, RangeInfo, SourceChange, SourceFileEdit, TextRange,
};

use super::find_all_refs;

pub(crate) fn rename(
    db: &RootDatabase,
    position: FilePosition,
    new_name: &str,
) -> Option<RangeInfo<SourceChange>> {
    match lex_single_valid_syntax_kind(new_name)? {
        SyntaxKind::IDENT | SyntaxKind::UNDERSCORE => (),
        _ => return None,
    }

    let parse = db.parse(position.file_id);
    if let Some((ast_name, ast_module)) =
        find_name_and_module_at_offset(parse.tree().syntax(), position)
    {
        let range = ast_name.syntax().text_range();
        rename_mod(db, &ast_name, &ast_module, position, new_name)
            .map(|info| RangeInfo::new(range, info))
    } else {
        rename_reference(db, position, new_name)
    }
}

fn find_name_and_module_at_offset(
    syntax: &SyntaxNode,
    position: FilePosition,
) -> Option<(ast::Name, ast::Module)> {
    let ast_name = find_node_at_offset::<ast::Name>(syntax, position.offset)?;
    let ast_module = ast::Module::cast(ast_name.syntax().parent()?)?;
    Some((ast_name, ast_module))
}

fn source_edit_from_file_id_range(
    file_id: FileId,
    range: TextRange,
    new_name: &str,
) -> SourceFileEdit {
    SourceFileEdit { file_id, edit: TextEdit::replace(range, new_name.into()) }
}

fn rename_mod(
    db: &RootDatabase,
    ast_name: &ast::Name,
    ast_module: &ast::Module,
    position: FilePosition,
    new_name: &str,
) -> Option<SourceChange> {
    let mut source_file_edits = Vec::new();
    let mut file_system_edits = Vec::new();
    let module_src = hir::InFile { file_id: position.file_id.into(), value: ast_module.clone() };
    if let Some(module) = hir::SourceBinder::new(db).to_def(module_src) {
        let src = module.definition_source(db);
        let file_id = src.file_id.original_file(db);
        match src.value {
            ModuleSource::SourceFile(..) => {
                let mod_path: RelativePathBuf = db.file_relative_path(file_id);
                // mod is defined in path/to/dir/mod.rs
                let dst_path = if mod_path.file_stem() == Some("mod") {
                    mod_path
                        .parent()
                        .and_then(|p| p.parent())
                        .or_else(|| Some(RelativePath::new("")))
                        .map(|p| p.join(new_name).join("mod.rs"))
                } else {
                    Some(mod_path.with_file_name(new_name).with_extension("rs"))
                };
                if let Some(path) = dst_path {
                    let move_file = FileSystemEdit::MoveFile {
                        src: file_id,
                        dst_source_root: db.file_source_root(position.file_id),
                        dst_path: path,
                    };
                    file_system_edits.push(move_file);
                }
            }
            ModuleSource::Module(..) => {}
        }
    }

    let edit = SourceFileEdit {
        file_id: position.file_id,
        edit: TextEdit::replace(ast_name.syntax().text_range(), new_name.into()),
    };
    source_file_edits.push(edit);

    Some(SourceChange::from_edits("rename", source_file_edits, file_system_edits))
}

fn rename_reference(
    db: &RootDatabase,
    position: FilePosition,
    new_name: &str,
) -> Option<RangeInfo<SourceChange>> {
    let RangeInfo { range, info: refs } = find_all_refs(db, position, None)?;

    let edit = refs
        .into_iter()
        .map(|reference| {
            source_edit_from_file_id_range(
                reference.file_range.file_id,
                reference.file_range.range,
                new_name,
            )
        })
        .collect::<Vec<_>>();

    if edit.is_empty() {
        return None;
    }

    Some(RangeInfo::new(range, SourceChange::source_file_edits("rename", edit)))
}

#[cfg(test)]
mod tests {
    use insta::assert_debug_snapshot;
    use ra_text_edit::TextEditBuilder;
    use test_utils::assert_eq_text;

    use crate::{
        mock_analysis::analysis_and_position, mock_analysis::single_file_with_position, FileId,
    };

    #[test]
    fn test_rename_to_underscore() {
        test_rename(
            r#"
    fn main() {
        let i<|> = 1;
    }"#,
            "_",
            r#"
    fn main() {
        let _ = 1;
    }"#,
        );
    }

    #[test]
    fn test_rename_to_raw_identifier() {
        test_rename(
            r#"
    fn main() {
        let i<|> = 1;
    }"#,
            "r#fn",
            r#"
    fn main() {
        let r#fn = 1;
    }"#,
        );
    }

    #[test]
    fn test_rename_to_invalid_identifier() {
        let (analysis, position) = single_file_with_position(
            "
    fn main() {
        let i<|> = 1;
    }",
        );
        let new_name = "invalid!";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert!(source_change.is_none());
    }

    #[test]
    fn test_rename_for_local() {
        test_rename(
            r#"
    fn main() {
        let mut i = 1;
        let j = 1;
        i = i<|> + j;

        {
            i = 0;
        }

        i = 5;
    }"#,
            "k",
            r#"
    fn main() {
        let mut k = 1;
        let j = 1;
        k = k + j;

        {
            k = 0;
        }

        k = 5;
    }"#,
        );
    }

    #[test]
    fn test_rename_for_macro_args() {
        test_rename(
            r#"
    macro_rules! foo {($i:ident) => {$i} }
    fn main() {
        let a<|> = "test";
        foo!(a);
    }"#,
            "b",
            r#"
    macro_rules! foo {($i:ident) => {$i} }
    fn main() {
        let b = "test";
        foo!(b);
    }"#,
        );
    }

    #[test]
    fn test_rename_for_param_inside() {
        test_rename(
            r#"
    fn foo(i : u32) -> u32 {
        i<|>
    }"#,
            "j",
            r#"
    fn foo(j : u32) -> u32 {
        j
    }"#,
        );
    }

    #[test]
    fn test_rename_refs_for_fn_param() {
        test_rename(
            r#"
    fn foo(i<|> : u32) -> u32 {
        i
    }"#,
            "new_name",
            r#"
    fn foo(new_name : u32) -> u32 {
        new_name
    }"#,
        );
    }

    #[test]
    fn test_rename_for_mut_param() {
        test_rename(
            r#"
    fn foo(mut i<|> : u32) -> u32 {
        i
    }"#,
            "new_name",
            r#"
    fn foo(mut new_name : u32) -> u32 {
        new_name
    }"#,
        );
    }

    #[test]
    fn test_rename_mod() {
        let (analysis, position) = analysis_and_position(
            "
            //- /lib.rs
            mod bar;

            //- /bar.rs
            mod foo<|>;

            //- /bar/foo.rs
            // emtpy
            ",
        );
        let new_name = "foo2";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert_debug_snapshot!(&source_change,
@r###"
        Some(
            RangeInfo {
                range: [4; 7),
                info: SourceChange {
                    label: "rename",
                    source_file_edits: [
                        SourceFileEdit {
                            file_id: FileId(
                                2,
                            ),
                            edit: TextEdit {
                                atoms: [
                                    AtomTextEdit {
                                        delete: [4; 7),
                                        insert: "foo2",
                                    },
                                ],
                            },
                        },
                    ],
                    file_system_edits: [
                        MoveFile {
                            src: FileId(
                                3,
                            ),
                            dst_source_root: SourceRootId(
                                0,
                            ),
                            dst_path: "bar/foo2.rs",
                        },
                    ],
                    cursor_position: None,
                },
            },
        )
        "###);
    }

    #[test]
    fn test_rename_mod_in_dir() {
        let (analysis, position) = analysis_and_position(
            "
            //- /lib.rs
            mod fo<|>o;
            //- /foo/mod.rs
            // emtpy
            ",
        );
        let new_name = "foo2";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert_debug_snapshot!(&source_change,
        @r###"
        Some(
            RangeInfo {
                range: [4; 7),
                info: SourceChange {
                    label: "rename",
                    source_file_edits: [
                        SourceFileEdit {
                            file_id: FileId(
                                1,
                            ),
                            edit: TextEdit {
                                atoms: [
                                    AtomTextEdit {
                                        delete: [4; 7),
                                        insert: "foo2",
                                    },
                                ],
                            },
                        },
                    ],
                    file_system_edits: [
                        MoveFile {
                            src: FileId(
                                2,
                            ),
                            dst_source_root: SourceRootId(
                                0,
                            ),
                            dst_path: "foo2/mod.rs",
                        },
                    ],
                    cursor_position: None,
                },
            },
        )
        "###
               );
    }

    fn test_rename(text: &str, new_name: &str, expected: &str) {
        let (analysis, position) = single_file_with_position(text);
        let source_change = analysis.rename(position, new_name).unwrap();
        let mut text_edit_builder = TextEditBuilder::default();
        let mut file_id: Option<FileId> = None;
        if let Some(change) = source_change {
            for edit in change.info.source_file_edits {
                file_id = Some(edit.file_id);
                for atom in edit.edit.as_atoms() {
                    text_edit_builder.replace(atom.delete, atom.insert.clone());
                }
            }
        }
        let result =
            text_edit_builder.finish().apply(&*analysis.file_text(file_id.unwrap()).unwrap());
        assert_eq_text!(expected, &*result);
    }
}
<commit_msg>rename module references<commit_after>//! FIXME: write short doc here

use hir::ModuleSource;
use ra_db::{RelativePath, RelativePathBuf, SourceDatabase, SourceDatabaseExt};
use ra_ide_db::RootDatabase;
use ra_syntax::{
    algo::find_node_at_offset, ast, lex_single_valid_syntax_kind, AstNode, SyntaxKind, SyntaxNode,
};
use ra_text_edit::TextEdit;

use crate::{
    FileId, FilePosition, FileSystemEdit, RangeInfo, SourceChange, SourceFileEdit, TextRange,
};

use super::find_all_refs;

pub(crate) fn rename(
    db: &RootDatabase,
    position: FilePosition,
    new_name: &str,
) -> Option<RangeInfo<SourceChange>> {
    match lex_single_valid_syntax_kind(new_name)? {
        SyntaxKind::IDENT | SyntaxKind::UNDERSCORE => (),
        _ => return None,
    }

    let parse = db.parse(position.file_id);
    if let Some((ast_name, ast_module)) =
        find_name_and_module_at_offset(parse.tree().syntax(), position)
    {
        let range = ast_name.syntax().text_range();
        rename_mod(db, &ast_name, &ast_module, position, new_name)
            .map(|info| RangeInfo::new(range, info))
    } else {
        rename_reference(db, position, new_name)
    }
}

fn find_name_and_module_at_offset(
    syntax: &SyntaxNode,
    position: FilePosition,
) -> Option<(ast::Name, ast::Module)> {
    let ast_name = find_node_at_offset::<ast::Name>(syntax, position.offset)?;
    let ast_module = ast::Module::cast(ast_name.syntax().parent()?)?;
    Some((ast_name, ast_module))
}

fn source_edit_from_file_id_range(
    file_id: FileId,
    range: TextRange,
    new_name: &str,
) -> SourceFileEdit {
    SourceFileEdit { file_id, edit: TextEdit::replace(range, new_name.into()) }
}

fn rename_mod(
    db: &RootDatabase,
    ast_name: &ast::Name,
    ast_module: &ast::Module,
    position: FilePosition,
    new_name: &str,
) -> Option<SourceChange> {
    let mut source_file_edits = Vec::new();
    let mut file_system_edits = Vec::new();
    let module_src = hir::InFile { file_id: position.file_id.into(), value: ast_module.clone() };
    if let Some(module) = hir::SourceBinder::new(db).to_def(module_src) {
        let src = module.definition_source(db);
        let file_id = src.file_id.original_file(db);
        match src.value {
            ModuleSource::SourceFile(..) => {
                let mod_path: RelativePathBuf = db.file_relative_path(file_id);
                // mod is defined in path/to/dir/mod.rs
                let dst_path = if mod_path.file_stem() == Some("mod") {
                    mod_path
                        .parent()
                        .and_then(|p| p.parent())
                        .or_else(|| Some(RelativePath::new("")))
                        .map(|p| p.join(new_name).join("mod.rs"))
                } else {
                    Some(mod_path.with_file_name(new_name).with_extension("rs"))
                };
                if let Some(path) = dst_path {
                    let move_file = FileSystemEdit::MoveFile {
                        src: file_id,
                        dst_source_root: db.file_source_root(position.file_id),
                        dst_path: path,
                    };
                    file_system_edits.push(move_file);
                }
            }
            ModuleSource::Module(..) => {}
        }
    }

    let edit = SourceFileEdit {
        file_id: position.file_id,
        edit: TextEdit::replace(ast_name.syntax().text_range(), new_name.into()),
    };
    source_file_edits.push(edit);

    if let Some(RangeInfo { range: _, info: refs }) = find_all_refs(db, position, None) {
        let ref_edits = refs
            .references
            .into_iter()
            .map(|reference| {
                source_edit_from_file_id_range(
                    reference.file_range.file_id,
                    reference.file_range.range,
                    new_name,
                )
            })
            .collect::<Vec<_>>();
        for ref_edit in ref_edits {
            source_file_edits.push(ref_edit);
        }
    }

    Some(SourceChange::from_edits("rename", source_file_edits, file_system_edits))
}

fn rename_reference(
    db: &RootDatabase,
    position: FilePosition,
    new_name: &str,
) -> Option<RangeInfo<SourceChange>> {
    let RangeInfo { range, info: refs } = find_all_refs(db, position, None)?;

    let edit = refs
        .into_iter()
        .map(|reference| {
            source_edit_from_file_id_range(
                reference.file_range.file_id,
                reference.file_range.range,
                new_name,
            )
        })
        .collect::<Vec<_>>();

    if edit.is_empty() {
        return None;
    }

    Some(RangeInfo::new(range, SourceChange::source_file_edits("rename", edit)))
}

#[cfg(test)]
mod tests {
    use insta::assert_debug_snapshot;
    use ra_text_edit::TextEditBuilder;
    use test_utils::assert_eq_text;

    use crate::{
        mock_analysis::analysis_and_position, mock_analysis::single_file_with_position, FileId,
    };

    #[test]
    fn test_rename_to_underscore() {
        test_rename(
            r#"
    fn main() {
        let i<|> = 1;
    }"#,
            "_",
            r#"
    fn main() {
        let _ = 1;
    }"#,
        );
    }

    #[test]
    fn test_rename_to_raw_identifier() {
        test_rename(
            r#"
    fn main() {
        let i<|> = 1;
    }"#,
            "r#fn",
            r#"
    fn main() {
        let r#fn = 1;
    }"#,
        );
    }

    #[test]
    fn test_rename_to_invalid_identifier() {
        let (analysis, position) = single_file_with_position(
            "
    fn main() {
        let i<|> = 1;
    }",
        );
        let new_name = "invalid!";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert!(source_change.is_none());
    }

    #[test]
    fn test_rename_for_local() {
        test_rename(
            r#"
    fn main() {
        let mut i = 1;
        let j = 1;
        i = i<|> + j;

        {
            i = 0;
        }

        i = 5;
    }"#,
            "k",
            r#"
    fn main() {
        let mut k = 1;
        let j = 1;
        k = k + j;

        {
            k = 0;
        }

        k = 5;
    }"#,
        );
    }

    #[test]
    fn test_rename_for_macro_args() {
        test_rename(
            r#"
    macro_rules! foo {($i:ident) => {$i} }
    fn main() {
        let a<|> = "test";
        foo!(a);
    }"#,
            "b",
            r#"
    macro_rules! foo {($i:ident) => {$i} }
    fn main() {
        let b = "test";
        foo!(b);
    }"#,
        );
    }

    #[test]
    fn test_rename_for_param_inside() {
        test_rename(
            r#"
    fn foo(i : u32) -> u32 {
        i<|>
    }"#,
            "j",
            r#"
    fn foo(j : u32) -> u32 {
        j
    }"#,
        );
    }

    #[test]
    fn test_rename_refs_for_fn_param() {
        test_rename(
            r#"
    fn foo(i<|> : u32) -> u32 {
        i
    }"#,
            "new_name",
            r#"
    fn foo(new_name : u32) -> u32 {
        new_name
    }"#,
        );
    }

    #[test]
    fn test_rename_for_mut_param() {
        test_rename(
            r#"
    fn foo(mut i<|> : u32) -> u32 {
        i
    }"#,
            "new_name",
            r#"
    fn foo(mut new_name : u32) -> u32 {
        new_name
    }"#,
        );
    }

    #[test]
    fn test_rename_mod() {
        let (analysis, position) = analysis_and_position(
            "
            //- /lib.rs
            mod bar;

            //- /bar.rs
            mod foo<|>;

            //- /bar/foo.rs
            // emtpy
            ",
        );
        let new_name = "foo2";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert_debug_snapshot!(&source_change,
@r###"
        Some(
            RangeInfo {
                range: [4; 7),
                info: SourceChange {
                    label: "rename",
                    source_file_edits: [
                        SourceFileEdit {
                            file_id: FileId(
                                2,
                            ),
                            edit: TextEdit {
                                atoms: [
                                    AtomTextEdit {
                                        delete: [4; 7),
                                        insert: "foo2",
                                    },
                                ],
                            },
                        },
                    ],
                    file_system_edits: [
                        MoveFile {
                            src: FileId(
                                3,
                            ),
                            dst_source_root: SourceRootId(
                                0,
                            ),
                            dst_path: "bar/foo2.rs",
                        },
                    ],
                    cursor_position: None,
                },
            },
        )
        "###);
    }

    #[test]
    fn test_rename_mod_in_dir() {
        let (analysis, position) = analysis_and_position(
            "
            //- /lib.rs
            mod fo<|>o;
            //- /foo/mod.rs
            // emtpy
            ",
        );
        let new_name = "foo2";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert_debug_snapshot!(&source_change,
        @r###"
        Some(
            RangeInfo {
                range: [4; 7),
                info: SourceChange {
                    label: "rename",
                    source_file_edits: [
                        SourceFileEdit {
                            file_id: FileId(
                                1,
                            ),
                            edit: TextEdit {
                                atoms: [
                                    AtomTextEdit {
                                        delete: [4; 7),
                                        insert: "foo2",
                                    },
                                ],
                            },
                        },
                    ],
                    file_system_edits: [
                        MoveFile {
                            src: FileId(
                                2,
                            ),
                            dst_source_root: SourceRootId(
                                0,
                            ),
                            dst_path: "foo2/mod.rs",
                        },
                    ],
                    cursor_position: None,
                },
            },
        )
        "###
               );
    }

    #[test]
    fn test_module_rename_in_path() {
        test_rename(
            r#"
    mod <|>foo {
        pub fn bar() {}
    }

    fn main() {
        foo::bar();
    }"#,
            "baz",
            r#"
    mod baz {
        pub fn bar() {}
    }

    fn main() {
        baz::bar();
    }"#,
        );
    }

    #[test]
    fn test_rename_mod_filename_and_path() {
        let (analysis, position) = analysis_and_position(
            "
            //- /lib.rs
            mod bar;
            fn f() {
                bar::foo::fun()
            }

            //- /bar.rs
            pub mod foo<|>;

            //- /bar/foo.rs
            // pub fn fun() {}
            ",
        );
        let new_name = "foo2";
        let source_change = analysis.rename(position, new_name).unwrap();
        assert_debug_snapshot!(&source_change,
@r###"
        Some(
            RangeInfo {
                range: [8; 11),
                info: SourceChange {
                    label: "rename",
                    source_file_edits: [
                        SourceFileEdit {
                            file_id: FileId(
                                2,
                            ),
                            edit: TextEdit {
                                atoms: [
                                    AtomTextEdit {
                                        delete: [8; 11),
                                        insert: "foo2",
                                    },
                                ],
                            },
                        },
                        SourceFileEdit {
                            file_id: FileId(
                                1,
                            ),
                            edit: TextEdit {
                                atoms: [
                                    AtomTextEdit {
                                        delete: [27; 30),
                                        insert: "foo2",
                                    },
                                ],
                            },
                        },
                    ],
                    file_system_edits: [
                        MoveFile {
                            src: FileId(
                                3,
                            ),
                            dst_source_root: SourceRootId(
                                0,
                            ),
                            dst_path: "bar/foo2.rs",
                        },
                    ],
                    cursor_position: None,
                },
            },
        )
        "###);
    }

    fn test_rename(text: &str, new_name: &str, expected: &str) {
        let (analysis, position) = single_file_with_position(text);
        let source_change = analysis.rename(position, new_name).unwrap();
        let mut text_edit_builder = TextEditBuilder::default();
        let mut file_id: Option<FileId> = None;
        if let Some(change) = source_change {
            for edit in change.info.source_file_edits {
                file_id = Some(edit.file_id);
                for atom in edit.edit.as_atoms() {
                    text_edit_builder.replace(atom.delete, atom.insert.clone());
                }
            }
        }
        let result =
            text_edit_builder.finish().apply(&*analysis.file_text(file_id.unwrap()).unwrap());
        assert_eq_text!(expected, &*result);
    }
}
<|endoftext|>


<commit_before>/// Construct a `serde_json::Value` from a JSON literal.
///
/// ```rust
/// # #![allow(unused_variables)]
/// # #[macro_use] extern crate serde_json;
/// # fn main() {
/// let value = json!({
///     "code": 200,
///     "success": true,
///     "payload": {
///         "features": [
///             "serde",
///             "json"
///         ]
///     }
/// });
/// # }
/// ```
///
/// Variables or expressions can be interpolated into the JSON literal. Any type
/// interpolated into an array element or object value must implement Serde's
/// `Serialize` trait, while any type interpolated into a object key must
/// implement `Into<String>`.
///
/// ```rust
/// # #![allow(unused_variables)]
/// # #[macro_use] extern crate serde_json;
/// # fn main() {
/// let code = 200;
/// let features = vec!["serde", "json"];
///
/// let value = json!({
///    "code": code,
///    "success": code == 200,
///    "payload": {
///        features[0]: features[1]
///    }
/// });
/// # }
/// ```
///
/// Trailing commas are allowed inside both arrays and objects.
///
/// ```rust
/// # #![allow(unused_variables)]
/// # #[macro_use] extern crate serde_json;
/// # fn main() {
/// let value = json!([
///     "notice",
///     "the",
///     "trailing",
///     "comma -->",
/// ]);
/// # }
/// ```
#[macro_export]
macro_rules! json {
    // Hide distracting implementation details from the generated rustdoc.
    ($($json:tt)+) => {
        json_internal!($($json)+)
    };
}

// Rocket relies on this because they export their own `json!` with a different
// doc comment than ours, and various Rust bugs prevent them from calling our
// `json!` from their `json!` so they call `json_internal!` directly. Check with
// @SergioBenitez before making breaking changes to this macro.
//
// Changes are fine as long as `json_internal!` does not call any new helper
// macros and can still be invoked as `json_internal!($($json)+)`.
#[macro_export]
#[doc(hidden)]
macro_rules! json_internal {
    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an array [...]. Produces a vec![...]
    // of the elements.
    //
    // Must be invoked as: json_internal!(@array [] $($tt)*)
    //////////////////////////////////////////////////////////////////////////

    // Done with trailing comma.
    (@array [$($elems:expr,)*]) => {
        vec![$($elems,)*]
    };

    // Done without trailing comma.
    (@array [$($elems:expr),*]) => {
        vec![$($elems),*]
    };

    // Next element is `null`.
    (@array [$($elems:expr,)*] null $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(null)] $($rest)*)
    };

    // Next element is `true`.
    (@array [$($elems:expr,)*] true $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(true)] $($rest)*)
    };

    // Next element is `false`.
    (@array [$($elems:expr,)*] false $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(false)] $($rest)*)
    };

    // Next element is an array.
    (@array [$($elems:expr,)*] [$($array:tt)*] $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!([$($array)*])] $($rest)*)
    };

    // Next element is a map.
    (@array [$($elems:expr,)*] {$($map:tt)*} $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!({$($map)*})] $($rest)*)
    };

    // Next element is an expression followed by comma.
    (@array [$($elems:expr,)*] $next:expr, $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!($next),] $($rest)*)
    };

    // Last element is an expression with no trailing comma.
    (@array [$($elems:expr,)*] $last:expr) => {
        json_internal!(@array [$($elems,)* json_internal!($last)])
    };

    // Comma after the most recent element.
    (@array [$($elems:expr),*] , $($rest:tt)*) => {
        json_internal!(@array [$($elems,)*] $($rest)*)
    };

    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an object {...}. Each entry is
    // inserted into the given map variable.
    //
    // Must be invoked as: json_internal!(@object $map () $($tt)*)
    //////////////////////////////////////////////////////////////////////////

    // Done.
    (@object $object:ident ()) => {};

    // Insert the current entry followed by trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr) , $($rest:tt)*) => {
        $object.insert(($($key)+).into(), $value);
        json_internal!(@object $object () $($rest)*);
    };

    // Insert the last entry without trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr)) => {
        $object.insert(($($key)+).into(), $value);
    };

    // Next value is `null`.
    (@object $object:ident ($($key:tt)+) : null $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!(null)) $($rest)*);
    };

    // Next value is `true`.
    (@object $object:ident ($($key:tt)+) : true $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!(true)) $($rest)*);
    };

    // Next value is `false`.
    (@object $object:ident ($($key:tt)+) : false $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!(false)) $($rest)*);
    };

    // Next value is an array.
    (@object $object:ident ($($key:tt)+) : [$($array:tt)*] $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!([$($array)*])) $($rest)*);
    };

    // Next value is a map.
    (@object $object:ident ($($key:tt)+) : {$($map:tt)*} $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!({$($map)*})) $($rest)*);
    };

    // Next value is an expression followed by comma.
    (@object $object:ident ($($key:tt)+) : $value:expr , $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)) , $($rest)*);
    };

    // Last value is an expression with no trailing comma.
    (@object $object:ident ($($key:tt)+) : $value:expr) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)));
    };

    // Missing value for last entry. Trigger a reasonable error message
    // referring to the unexpected end of macro invocation.
    (@object $object:ident ($($key:tt)+) :) => {
        json_internal!();
    };

    // Misplaced colon. Trigger a reasonable error message by failing to match
    // the colon in the recursive call.
    (@object $object:ident () : $($rest:tt)*) => {
        json_internal!(@object $object [:] (:));
    };

    // Found a comma inside a key. Trigger a reasonable error message by failing
    // to match the comma in the recursive call.
    (@object $object:ident ($($key:tt)*) , $($rest:tt)*) => {
        json_internal!(@object $object [,] (,));
    };

    // Key is fully parenthesized. This avoids clippy double_parens false
    // positives because the parenthesization may be necessary here.
    (@object $object:ident () ($key:expr) : $($rest:tt)*) => {
        json_internal!(@object $object ($key) : $($rest)*);
    };

    // Munch a token into the current key.
    (@object $object:ident ($($key:tt)*) $tt:tt $($rest:tt)*) => {
        json_internal!(@object $object ($($key)* $tt) $($rest)*);
    };

    //////////////////////////////////////////////////////////////////////////
    // The main implementation.
    //
    // Must be invoked as: json_internal!($($json)+)
    //////////////////////////////////////////////////////////////////////////

    (null) => {
        $crate::Value::Null
    };

    (true) => {
        $crate::Value::Bool(true)
    };

    (false) => {
        $crate::Value::Bool(false)
    };

    ([]) => {
        $crate::Value::Array(vec![])
    };

    ([ $($tt:tt)+ ]) => {
        $crate::Value::Array(json_internal!(@array [] $($tt)+))
    };

    ({}) => {
        $crate::Value::Object($crate::Map::new())
    };

    ({ $($tt:tt)+ }) => {
        $crate::Value::Object({
            let mut object = $crate::Map::new();
            json_internal!(@object object () $($tt)+);
            object
        })
    };

    // Any Serialize type: numbers, strings, struct literals, variables etc.
    // Must be below every other rule.
    ($other:expr) => {
        $crate::to_value(&$other).unwrap()
    };
}
<commit_msg>Simplify json macro error cases<commit_after>/// Construct a `serde_json::Value` from a JSON literal.
///
/// ```rust
/// # #![allow(unused_variables)]
/// # #[macro_use] extern crate serde_json;
/// # fn main() {
/// let value = json!({
///     "code": 200,
///     "success": true,
///     "payload": {
///         "features": [
///             "serde",
///             "json"
///         ]
///     }
/// });
/// # }
/// ```
///
/// Variables or expressions can be interpolated into the JSON literal. Any type
/// interpolated into an array element or object value must implement Serde's
/// `Serialize` trait, while any type interpolated into a object key must
/// implement `Into<String>`.
///
/// ```rust
/// # #![allow(unused_variables)]
/// # #[macro_use] extern crate serde_json;
/// # fn main() {
/// let code = 200;
/// let features = vec!["serde", "json"];
///
/// let value = json!({
///    "code": code,
///    "success": code == 200,
///    "payload": {
///        features[0]: features[1]
///    }
/// });
/// # }
/// ```
///
/// Trailing commas are allowed inside both arrays and objects.
///
/// ```rust
/// # #![allow(unused_variables)]
/// # #[macro_use] extern crate serde_json;
/// # fn main() {
/// let value = json!([
///     "notice",
///     "the",
///     "trailing",
///     "comma -->",
/// ]);
/// # }
/// ```
#[macro_export]
macro_rules! json {
    // Hide distracting implementation details from the generated rustdoc.
    ($($json:tt)+) => {
        json_internal!($($json)+)
    };
}

// Rocket relies on this because they export their own `json!` with a different
// doc comment than ours, and various Rust bugs prevent them from calling our
// `json!` from their `json!` so they call `json_internal!` directly. Check with
// @SergioBenitez before making breaking changes to this macro.
//
// Changes are fine as long as `json_internal!` does not call any new helper
// macros and can still be invoked as `json_internal!($($json)+)`.
#[macro_export]
#[doc(hidden)]
macro_rules! json_internal {
    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an array [...]. Produces a vec![...]
    // of the elements.
    //
    // Must be invoked as: json_internal!(@array [] $($tt)*)
    //////////////////////////////////////////////////////////////////////////

    // Done with trailing comma.
    (@array [$($elems:expr,)*]) => {
        vec![$($elems,)*]
    };

    // Done without trailing comma.
    (@array [$($elems:expr),*]) => {
        vec![$($elems),*]
    };

    // Next element is `null`.
    (@array [$($elems:expr,)*] null $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(null)] $($rest)*)
    };

    // Next element is `true`.
    (@array [$($elems:expr,)*] true $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(true)] $($rest)*)
    };

    // Next element is `false`.
    (@array [$($elems:expr,)*] false $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!(false)] $($rest)*)
    };

    // Next element is an array.
    (@array [$($elems:expr,)*] [$($array:tt)*] $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!([$($array)*])] $($rest)*)
    };

    // Next element is a map.
    (@array [$($elems:expr,)*] {$($map:tt)*} $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!({$($map)*})] $($rest)*)
    };

    // Next element is an expression followed by comma.
    (@array [$($elems:expr,)*] $next:expr, $($rest:tt)*) => {
        json_internal!(@array [$($elems,)* json_internal!($next),] $($rest)*)
    };

    // Last element is an expression with no trailing comma.
    (@array [$($elems:expr,)*] $last:expr) => {
        json_internal!(@array [$($elems,)* json_internal!($last)])
    };

    // Comma after the most recent element.
    (@array [$($elems:expr),*] , $($rest:tt)*) => {
        json_internal!(@array [$($elems,)*] $($rest)*)
    };

    //////////////////////////////////////////////////////////////////////////
    // TT muncher for parsing the inside of an object {...}. Each entry is
    // inserted into the given map variable.
    //
    // Must be invoked as: json_internal!(@object $map () $($tt)*)
    //////////////////////////////////////////////////////////////////////////

    // Done.
    (@object $object:ident ()) => {};

    // Insert the current entry followed by trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr) , $($rest:tt)*) => {
        $object.insert(($($key)+).into(), $value);
        json_internal!(@object $object () $($rest)*);
    };

    // Insert the last entry without trailing comma.
    (@object $object:ident [$($key:tt)+] ($value:expr)) => {
        $object.insert(($($key)+).into(), $value);
    };

    // Next value is `null`.
    (@object $object:ident ($($key:tt)+) : null $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!(null)) $($rest)*);
    };

    // Next value is `true`.
    (@object $object:ident ($($key:tt)+) : true $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!(true)) $($rest)*);
    };

    // Next value is `false`.
    (@object $object:ident ($($key:tt)+) : false $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!(false)) $($rest)*);
    };

    // Next value is an array.
    (@object $object:ident ($($key:tt)+) : [$($array:tt)*] $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!([$($array)*])) $($rest)*);
    };

    // Next value is a map.
    (@object $object:ident ($($key:tt)+) : {$($map:tt)*} $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!({$($map)*})) $($rest)*);
    };

    // Next value is an expression followed by comma.
    (@object $object:ident ($($key:tt)+) : $value:expr , $($rest:tt)*) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)) , $($rest)*);
    };

    // Last value is an expression with no trailing comma.
    (@object $object:ident ($($key:tt)+) : $value:expr) => {
        json_internal!(@object $object [$($key)+] (json_internal!($value)));
    };

    // Missing value for last entry. Trigger a reasonable error message
    // referring to the unexpected end of macro invocation.
    (@object $object:ident ($($key:tt)+) :) => {
        json_internal!();
    };

    // Misplaced colon. Trigger a reasonable error message by failing to match
    // the colon in the recursive call.
    (@object $object:ident () : $($rest:tt)*) => {
        json_internal!(:);
    };

    // Found a comma inside a key. Trigger a reasonable error message by failing
    // to match the comma in the recursive call.
    (@object $object:ident ($($key:tt)*) , $($rest:tt)*) => {
        json_internal!(,);
    };

    // Key is fully parenthesized. This avoids clippy double_parens false
    // positives because the parenthesization may be necessary here.
    (@object $object:ident () ($key:expr) : $($rest:tt)*) => {
        json_internal!(@object $object ($key) : $($rest)*);
    };

    // Munch a token into the current key.
    (@object $object:ident ($($key:tt)*) $tt:tt $($rest:tt)*) => {
        json_internal!(@object $object ($($key)* $tt) $($rest)*);
    };

    //////////////////////////////////////////////////////////////////////////
    // The main implementation.
    //
    // Must be invoked as: json_internal!($($json)+)
    //////////////////////////////////////////////////////////////////////////

    (null) => {
        $crate::Value::Null
    };

    (true) => {
        $crate::Value::Bool(true)
    };

    (false) => {
        $crate::Value::Bool(false)
    };

    ([]) => {
        $crate::Value::Array(vec![])
    };

    ([ $($tt:tt)+ ]) => {
        $crate::Value::Array(json_internal!(@array [] $($tt)+))
    };

    ({}) => {
        $crate::Value::Object($crate::Map::new())
    };

    ({ $($tt:tt)+ }) => {
        $crate::Value::Object({
            let mut object = $crate::Map::new();
            json_internal!(@object object () $($tt)+);
            object
        })
    };

    // Any Serialize type: numbers, strings, struct literals, variables etc.
    // Must be below every other rule.
    ($other:expr) => {
        $crate::to_value(&$other).unwrap()
    };
}
<|endoftext|>


<commit_before>//! Connection helper.
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;

use tungstenite::client::uri_mode;
use tungstenite::handshake::client::Response;
use tungstenite::protocol::WebSocketConfig;
use tungstenite::Error;

use super::{client_async_with_config, IntoClientRequest, Request, WebSocketStream};

#[cfg(feature = "tls")]
pub(crate) mod encryption {
    use native_tls::TlsConnector;
    use tokio_tls::{TlsConnector as TokioTlsConnector, TlsStream};

    use tokio::io::{AsyncRead, AsyncWrite};

    use tungstenite::stream::Mode;
    use tungstenite::Error;

    use crate::stream::Stream as StreamSwitcher;

    /// A stream that might be protected with TLS.
    pub type MaybeTlsStream<S> = StreamSwitcher<S, TlsStream<S>>;

    pub type AutoStream<S> = MaybeTlsStream<S>;

    pub async fn wrap_stream<S>(
        socket: S,
        domain: String,
        mode: Mode,
    ) -> Result<AutoStream<S>, Error>
    where
        S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(socket)),
            Mode::Tls => {
                let try_connector = TlsConnector::new();
                let connector = try_connector.map_err(Error::Tls)?;
                let stream = TokioTlsConnector::from(connector);
                let connected = stream.connect(&domain, socket).await;
                match connected {
                    Err(e) => Err(Error::Tls(e)),
                    Ok(s) => Ok(StreamSwitcher::Tls(s)),
                }
            }
        }
    }
}

#[cfg(feature = "tls")]
pub use self::encryption::MaybeTlsStream;

#[cfg(not(feature = "tls"))]
pub(crate) mod encryption {
    use tokio::io::{AsyncRead, AsyncWrite};

    use tungstenite::stream::Mode;
    use tungstenite::Error;

    pub type AutoStream<S> = S;

    pub async fn wrap_stream<S>(
        socket: S,
        _domain: String,
        mode: Mode,
    ) -> Result<AutoStream<S>, Error>
    where
        S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    {
        match mode {
            Mode::Plain => Ok(socket),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

use self::encryption::{wrap_stream, AutoStream};

/// Get a domain from an URL.
#[inline]
fn domain(request: &Request) -> Result<String, Error> {
    match request.uri().host() {
        Some(d) => Ok(d.to_string()),
        None => Err(Error::Url("no host name in the url".into())),
    }
}

/// Creates a WebSocket handshake from a request and a stream,
/// upgrading the stream to TLS if required.
pub async fn client_async_tls<R, S>(
    request: R,
    stream: S,
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
    S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    AutoStream<S>: Unpin,
{
    client_async_tls_with_config(request, stream, None).await
}

/// The same as `client_async_tls()` but the one can specify a websocket configuration.
/// Please refer to `client_async_tls()` for more details.

pub async fn client_async_tls_with_config<R, S>(
    request: R,
    stream: S,
    config: Option<WebSocketConfig>
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
    S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    AutoStream<S>: Unpin,
{
    let request = request.into_client_request()?;

    let domain = domain(&request)?;

    // Make sure we check domain and mode first. URL must be valid.
    let mode = uri_mode(&request.uri())?;

    let stream = wrap_stream(stream, domain, mode).await?;
    client_async_with_config(request, stream, config).await
}

/// Connect to a given URL.
pub async fn connect_async<R>(
    request: R
) -> Result<(WebSocketStream<AutoStream<TcpStream>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
{
    connect_async_with_config(request, None).await
}

/// The same as `connect_async()` but the one can specify a websocket configuration.
/// Please refer to `connect_async()` for more details.

pub async fn connect_async_with_config<R>(
    request: R,
    config: Option<WebSocketConfig>
) -> Result<(WebSocketStream<AutoStream<TcpStream>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
{
    let request = request.into_client_request()?;

    let domain = domain(&request)?;
    let port = request
        .uri()
        .port_u16()
        .or_else(|| match request.uri().scheme_str() {
            Some("wss") => Some(443),
            Some("ws") => Some(80),
            _ => None,
        })
        .ok_or_else(|| Error::Url("Url scheme not supported".into()))?;

    let addr = format!("{}:{}", domain, port);
    let try_socket = TcpStream::connect(addr).await;
    let socket = try_socket.map_err(Error::Io)?;
    client_async_tls_with_config(request, socket, config).await
}
<commit_msg>Fix formatting of comments<commit_after>//! Connection helper.
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;

use tungstenite::client::uri_mode;
use tungstenite::handshake::client::Response;
use tungstenite::protocol::WebSocketConfig;
use tungstenite::Error;

use super::{client_async_with_config, IntoClientRequest, Request, WebSocketStream};

#[cfg(feature = "tls")]
pub(crate) mod encryption {
    use native_tls::TlsConnector;
    use tokio_tls::{TlsConnector as TokioTlsConnector, TlsStream};

    use tokio::io::{AsyncRead, AsyncWrite};

    use tungstenite::stream::Mode;
    use tungstenite::Error;

    use crate::stream::Stream as StreamSwitcher;

    /// A stream that might be protected with TLS.
    pub type MaybeTlsStream<S> = StreamSwitcher<S, TlsStream<S>>;

    pub type AutoStream<S> = MaybeTlsStream<S>;

    pub async fn wrap_stream<S>(
        socket: S,
        domain: String,
        mode: Mode,
    ) -> Result<AutoStream<S>, Error>
    where
        S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(socket)),
            Mode::Tls => {
                let try_connector = TlsConnector::new();
                let connector = try_connector.map_err(Error::Tls)?;
                let stream = TokioTlsConnector::from(connector);
                let connected = stream.connect(&domain, socket).await;
                match connected {
                    Err(e) => Err(Error::Tls(e)),
                    Ok(s) => Ok(StreamSwitcher::Tls(s)),
                }
            }
        }
    }
}

#[cfg(feature = "tls")]
pub use self::encryption::MaybeTlsStream;

#[cfg(not(feature = "tls"))]
pub(crate) mod encryption {
    use tokio::io::{AsyncRead, AsyncWrite};

    use tungstenite::stream::Mode;
    use tungstenite::Error;

    pub type AutoStream<S> = S;

    pub async fn wrap_stream<S>(
        socket: S,
        _domain: String,
        mode: Mode,
    ) -> Result<AutoStream<S>, Error>
    where
        S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    {
        match mode {
            Mode::Plain => Ok(socket),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

use self::encryption::{wrap_stream, AutoStream};

/// Get a domain from an URL.
#[inline]
fn domain(request: &Request) -> Result<String, Error> {
    match request.uri().host() {
        Some(d) => Ok(d.to_string()),
        None => Err(Error::Url("no host name in the url".into())),
    }
}

/// Creates a WebSocket handshake from a request and a stream,
/// upgrading the stream to TLS if required.
pub async fn client_async_tls<R, S>(
    request: R,
    stream: S,
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
    S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    AutoStream<S>: Unpin,
{
    client_async_tls_with_config(request, stream, None).await
}

/// The same as `client_async_tls()` but the one can specify a websocket configuration.
/// Please refer to `client_async_tls()` for more details.
pub async fn client_async_tls_with_config<R, S>(
    request: R,
    stream: S,
    config: Option<WebSocketConfig>
) -> Result<(WebSocketStream<AutoStream<S>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
    S: 'static + AsyncRead + AsyncWrite + Send + Unpin,
    AutoStream<S>: Unpin,
{
    let request = request.into_client_request()?;

    let domain = domain(&request)?;

    // Make sure we check domain and mode first. URL must be valid.
    let mode = uri_mode(&request.uri())?;

    let stream = wrap_stream(stream, domain, mode).await?;
    client_async_with_config(request, stream, config).await
}

/// Connect to a given URL.
pub async fn connect_async<R>(
    request: R
) -> Result<(WebSocketStream<AutoStream<TcpStream>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
{
    connect_async_with_config(request, None).await
}

/// The same as `connect_async()` but the one can specify a websocket configuration.
/// Please refer to `connect_async()` for more details.
pub async fn connect_async_with_config<R>(
    request: R,
    config: Option<WebSocketConfig>
) -> Result<(WebSocketStream<AutoStream<TcpStream>>, Response), Error>
where
    R: IntoClientRequest + Unpin,
{
    let request = request.into_client_request()?;

    let domain = domain(&request)?;
    let port = request
        .uri()
        .port_u16()
        .or_else(|| match request.uri().scheme_str() {
            Some("wss") => Some(443),
            Some("ws") => Some(80),
            _ => None,
        })
        .ok_or_else(|| Error::Url("Url scheme not supported".into()))?;

    let addr = format!("{}:{}", domain, port);
    let try_socket = TcpStream::connect(addr).await;
    let socket = try_socket.map_err(Error::Io)?;
    client_async_tls_with_config(request, socket, config).await
}
<|endoftext|>