<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);

    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);
    }
}
<commit_msg>use Vec::extend instead of Vec::push in loop<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<_>>();
        source_file_edits.extend(ref_edits);
    }

    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>//! For middleware documentation, see [`NormalizePath`].

use actix_http::http::{PathAndQuery, Uri};
use actix_service::{Service, Transform};
use actix_utils::future::{ready, Ready};
use bytes::Bytes;
use regex::Regex;

use crate::{
    service::{ServiceRequest, ServiceResponse},
    Error,
};

/// Determines the behavior of the [`NormalizePath`] middleware.
///
/// The default is `TrailingSlash::Trim`.
#[non_exhaustive]
#[derive(Debug, Clone, Copy)]
pub enum TrailingSlash {
    /// Trim trailing slashes from the end of the path.
    ///
    /// Using this will require all routes to omit trailing slashes for them to be accessible.
    Trim,

    /// Only merge any present multiple trailing slashes.
    ///
    /// This option provides the best compatibility with behavior in actix-web v2.0.
    MergeOnly,

    /// Always add a trailing slash to the end of the path.
    ///
    /// Using this will require all routes have a trailing slash for them to be accessible.
    Always,
}

impl Default for TrailingSlash {
    fn default() -> Self {
        TrailingSlash::Trim
    }
}

/// Middleware for normalizing a request's path so that routes can be matched more flexibly.
///
/// # Normalization Steps
/// - Merges consecutive slashes into one. (For example, `/path//one` always becomes `/path/one`.)
/// - Appends a trailing slash if one is not present, removes one if present, or keeps trailing
///   slashes as-is, depending on which [`TrailingSlash`] variant is supplied
///   to [`new`](NormalizePath::new()).
///
/// # Default Behavior
/// The default constructor chooses to strip trailing slashes from the end of paths with them
/// ([`TrailingSlash::Trim`]). The implication is that route definitions should be defined without
/// trailing slashes or else they will be inaccessible (or vice versa when using the
/// `TrailingSlash::Always` behavior), as shown in the example tests below.
///
/// # Examples
/// ```
/// use actix_web::{web, middleware, App};
///
/// # actix_web::rt::System::new().block_on(async {
/// let app = App::new()
///     .wrap(middleware::NormalizePath::default())
///     .route("/test", web::get().to(|| async { "test" }))
///     .route("/unmatchable/", web::get().to(|| async { "unmatchable" }));
///
/// use actix_web::http::StatusCode;
/// use actix_web::test::{call_service, init_service, TestRequest};
///
/// let app = init_service(app).await;
///
/// let req = TestRequest::with_uri("/test").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::OK);
///
/// let req = TestRequest::with_uri("/test/").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::OK);
///
/// let req = TestRequest::with_uri("/unmatchable").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::NOT_FOUND);
///
/// let req = TestRequest::with_uri("/unmatchable/").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::NOT_FOUND);
/// # })
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct NormalizePath(TrailingSlash);

impl NormalizePath {
    /// Create new `NormalizePath` middleware with the specified trailing slash style.
    pub fn new(trailing_slash_style: TrailingSlash) -> Self {
        NormalizePath(trailing_slash_style)
    }
}

impl<S, B> Transform<S, ServiceRequest> for NormalizePath
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Transform = NormalizePathNormalization<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(NormalizePathNormalization {
            service,
            merge_slash: Regex::new("//+").unwrap(),
            trailing_slash_behavior: self.0,
        }))
    }
}

pub struct NormalizePathNormalization<S> {
    service: S,
    merge_slash: Regex,
    trailing_slash_behavior: TrailingSlash,
}

impl<S, B> Service<ServiceRequest> for NormalizePathNormalization<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = S::Future;

    actix_service::forward_ready!(service);

    fn call(&self, mut req: ServiceRequest) -> Self::Future {
        let head = req.head_mut();

        let original_path = head.uri.path();

        // Either adds a string to the end (duplicates will be removed anyways) or trims all slashes from the end
        let path = match self.trailing_slash_behavior {
            TrailingSlash::Always => original_path.to_string() + "/",
            TrailingSlash::MergeOnly => original_path.to_string(),
            TrailingSlash::Trim => original_path.trim_end_matches('/').to_string(),
        };

        // normalize multiple /'s to one /
        let path = self.merge_slash.replace_all(&path, "/");

        // Ensure root paths are still resolvable. If resulting path is blank after previous step
        // it means the path was one or more slashes. Reduce to single slash.
        let path = if path.is_empty() { "/" } else { path.as_ref() };

        // Check whether the path has been changed
        //
        // This check was previously implemented as string length comparison
        //
        // That approach fails when a trailing slash is added,
        // and a duplicate slash is removed,
        // since the length of the strings remains the same
        //
        // For example, the path "/v1//s" will be normalized to "/v1/s/"
        // Both of the paths have the same length,
        // so the change can not be deduced from the length comparison
        if path != original_path {
            let mut parts = head.uri.clone().into_parts();
            let query = parts.path_and_query.as_ref().and_then(|pq| pq.query());

            let path = if let Some(q) = query {
                Bytes::from(format!("{}?{}", path, q))
            } else {
                Bytes::copy_from_slice(path.as_bytes())
            };
            parts.path_and_query = Some(PathAndQuery::from_maybe_shared(path).unwrap());

            let uri = Uri::from_parts(parts).unwrap();
            req.match_info_mut().get_mut().update(&uri);
            req.head_mut().uri = uri;
        }

        self.service.call(req)
    }
}

#[cfg(test)]
mod tests {
    use actix_service::IntoService;

    use super::*;
    use crate::{
        dev::ServiceRequest,
        test::{call_service, init_service, TestRequest},
        web, App, HttpResponse,
    };

    #[actix_rt::test]
    async fn test_wrap() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath::default())
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something").to(HttpResponse::Ok)),
        )
        .await;

        let req = TestRequest::with_uri("/").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req = TestRequest::with_uri("/?query=test").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req = TestRequest::with_uri("///").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req = TestRequest::with_uri("/v1//something////").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req2 = TestRequest::with_uri("//v1/something").to_request();
        let res2 = call_service(&app, req2).await;
        assert!(res2.status().is_success());

        let req3 = TestRequest::with_uri("//v1//////something").to_request();
        let res3 = call_service(&app, req3).await;
        assert!(res3.status().is_success());

        let req4 = TestRequest::with_uri("/v1//something").to_request();
        let res4 = call_service(&app, req4).await;
        assert!(res4.status().is_success());
    }

    #[actix_rt::test]
    async fn trim_trailing_slashes() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath(TrailingSlash::Trim))
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something").to(HttpResponse::Ok)),
        )
        .await;

        // root paths should still work
        let req = TestRequest::with_uri("/").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req = TestRequest::with_uri("/?query=test").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req = TestRequest::with_uri("///").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req = TestRequest::with_uri("/v1/something////").to_request();
        let res = call_service(&app, req).await;
        assert!(res.status().is_success());

        let req2 = TestRequest::with_uri("/v1/something/").to_request();
        let res2 = call_service(&app, req2).await;
        assert!(res2.status().is_success());

        let req3 = TestRequest::with_uri("//v1//something//").to_request();
        let res3 = call_service(&app, req3).await;
        assert!(res3.status().is_success());

        let req4 = TestRequest::with_uri("//v1//something").to_request();
        let res4 = call_service(&app, req4).await;
        assert!(res4.status().is_success());
    }

    #[actix_rt::test]
    async fn keep_trailing_slash_unchanged() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath(TrailingSlash::MergeOnly))
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something").to(HttpResponse::Ok))
                .service(web::resource("/v1/").to(HttpResponse::Ok)),
        )
        .await;

        let tests = vec![
            ("/", true), // root paths should still work
            ("/?query=test", true),
            ("///", true),
            ("/v1/something////", false),
            ("/v1/something/", false),
            ("//v1//something", true),
            ("/v1/", true),
            ("/v1", false),
            ("/v1////", true),
            ("//v1//", true),
            ("///v1", false),
        ];

        for (path, success) in tests {
            let req = TestRequest::with_uri(path).to_request();
            let res = call_service(&app, req).await;
            assert_eq!(res.status().is_success(), success);
        }
    }

    #[actix_rt::test]
    async fn test_in_place_normalization() {
        let srv = |req: ServiceRequest| {
            assert_eq!("/v1/something", req.path());
            ready(Ok(req.into_response(HttpResponse::Ok().finish())))
        };

        let normalize = NormalizePath::default()
            .new_transform(srv.into_service())
            .await
            .unwrap();

        let req = TestRequest::with_uri("/v1//something////").to_srv_request();
        let res = normalize.call(req).await.unwrap();
        assert!(res.status().is_success());

        let req2 = TestRequest::with_uri("///v1/something").to_srv_request();
        let res2 = normalize.call(req2).await.unwrap();
        assert!(res2.status().is_success());

        let req3 = TestRequest::with_uri("//v1///something").to_srv_request();
        let res3 = normalize.call(req3).await.unwrap();
        assert!(res3.status().is_success());

        let req4 = TestRequest::with_uri("/v1//something").to_srv_request();
        let res4 = normalize.call(req4).await.unwrap();
        assert!(res4.status().is_success());
    }

    #[actix_rt::test]
    async fn should_normalize_nothing() {
        const URI: &str = "/v1/something";

        let srv = |req: ServiceRequest| {
            assert_eq!(URI, req.path());
            ready(Ok(req.into_response(HttpResponse::Ok().finish())))
        };

        let normalize = NormalizePath::default()
            .new_transform(srv.into_service())
            .await
            .unwrap();

        let req = TestRequest::with_uri(URI).to_srv_request();
        let res = normalize.call(req).await.unwrap();
        assert!(res.status().is_success());
    }

    #[actix_rt::test]
    async fn should_normalize_no_trail() {
        let srv = |req: ServiceRequest| {
            assert_eq!("/v1/something", req.path());
            ready(Ok(req.into_response(HttpResponse::Ok().finish())))
        };

        let normalize = NormalizePath::default()
            .new_transform(srv.into_service())
            .await
            .unwrap();

        let req = TestRequest::with_uri("/v1/something/").to_srv_request();
        let res = normalize.call(req).await.unwrap();
        assert!(res.status().is_success());
    }
}
<commit_msg>Additional test coverage and tidyup (middleware::normalize) (#2243)<commit_after>//! For middleware documentation, see [`NormalizePath`].

use actix_http::http::{PathAndQuery, Uri};
use actix_service::{Service, Transform};
use actix_utils::future::{ready, Ready};
use bytes::Bytes;
use regex::Regex;

use crate::{
    service::{ServiceRequest, ServiceResponse},
    Error,
};

/// Determines the behavior of the [`NormalizePath`] middleware.
///
/// The default is `TrailingSlash::Trim`.
#[non_exhaustive]
#[derive(Debug, Clone, Copy)]
pub enum TrailingSlash {
    /// Trim trailing slashes from the end of the path.
    ///
    /// Using this will require all routes to omit trailing slashes for them to be accessible.
    Trim,

    /// Only merge any present multiple trailing slashes.
    ///
    /// This option provides the best compatibility with behavior in actix-web v2.0.
    MergeOnly,

    /// Always add a trailing slash to the end of the path.
    ///
    /// Using this will require all routes have a trailing slash for them to be accessible.
    Always,
}

impl Default for TrailingSlash {
    fn default() -> Self {
        TrailingSlash::Trim
    }
}

/// Middleware for normalizing a request's path so that routes can be matched more flexibly.
///
/// # Normalization Steps
/// - Merges consecutive slashes into one. (For example, `/path//one` always becomes `/path/one`.)
/// - Appends a trailing slash if one is not present, removes one if present, or keeps trailing
///   slashes as-is, depending on which [`TrailingSlash`] variant is supplied
///   to [`new`](NormalizePath::new()).
///
/// # Default Behavior
/// The default constructor chooses to strip trailing slashes from the end of paths with them
/// ([`TrailingSlash::Trim`]). The implication is that route definitions should be defined without
/// trailing slashes or else they will be inaccessible (or vice versa when using the
/// `TrailingSlash::Always` behavior), as shown in the example tests below.
///
/// # Examples
/// ```
/// use actix_web::{web, middleware, App};
///
/// # actix_web::rt::System::new().block_on(async {
/// let app = App::new()
///     .wrap(middleware::NormalizePath::default())
///     .route("/test", web::get().to(|| async { "test" }))
///     .route("/unmatchable/", web::get().to(|| async { "unmatchable" }));
///
/// use actix_web::http::StatusCode;
/// use actix_web::test::{call_service, init_service, TestRequest};
///
/// let app = init_service(app).await;
///
/// let req = TestRequest::with_uri("/test").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::OK);
///
/// let req = TestRequest::with_uri("/test/").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::OK);
///
/// let req = TestRequest::with_uri("/unmatchable").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::NOT_FOUND);
///
/// let req = TestRequest::with_uri("/unmatchable/").to_request();
/// let res = call_service(&app, req).await;
/// assert_eq!(res.status(), StatusCode::NOT_FOUND);
/// # })
/// ```
#[derive(Debug, Clone, Copy, Default)]
pub struct NormalizePath(TrailingSlash);

impl NormalizePath {
    /// Create new `NormalizePath` middleware with the specified trailing slash style.
    pub fn new(trailing_slash_style: TrailingSlash) -> Self {
        NormalizePath(trailing_slash_style)
    }
}

impl<S, B> Transform<S, ServiceRequest> for NormalizePath
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Transform = NormalizePathNormalization<S>;
    type InitError = ();
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(NormalizePathNormalization {
            service,
            merge_slash: Regex::new("//+").unwrap(),
            trailing_slash_behavior: self.0,
        }))
    }
}

pub struct NormalizePathNormalization<S> {
    service: S,
    merge_slash: Regex,
    trailing_slash_behavior: TrailingSlash,
}

impl<S, B> Service<ServiceRequest> for NormalizePathNormalization<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
    S::Future: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = S::Future;

    actix_service::forward_ready!(service);

    fn call(&self, mut req: ServiceRequest) -> Self::Future {
        let head = req.head_mut();

        let original_path = head.uri.path();

        // Either adds a string to the end (duplicates will be removed anyways) or trims all slashes from the end
        let path = match self.trailing_slash_behavior {
            TrailingSlash::Always => original_path.to_string() + "/",
            TrailingSlash::MergeOnly => original_path.to_string(),
            TrailingSlash::Trim => original_path.trim_end_matches('/').to_string(),
        };

        // normalize multiple /'s to one /
        let path = self.merge_slash.replace_all(&path, "/");

        // Ensure root paths are still resolvable. If resulting path is blank after previous step
        // it means the path was one or more slashes. Reduce to single slash.
        let path = if path.is_empty() { "/" } else { path.as_ref() };

        // Check whether the path has been changed
        //
        // This check was previously implemented as string length comparison
        //
        // That approach fails when a trailing slash is added,
        // and a duplicate slash is removed,
        // since the length of the strings remains the same
        //
        // For example, the path "/v1//s" will be normalized to "/v1/s/"
        // Both of the paths have the same length,
        // so the change can not be deduced from the length comparison
        if path != original_path {
            let mut parts = head.uri.clone().into_parts();
            let query = parts.path_and_query.as_ref().and_then(|pq| pq.query());

            let path = if let Some(q) = query {
                Bytes::from(format!("{}?{}", path, q))
            } else {
                Bytes::copy_from_slice(path.as_bytes())
            };
            parts.path_and_query = Some(PathAndQuery::from_maybe_shared(path).unwrap());

            let uri = Uri::from_parts(parts).unwrap();
            req.match_info_mut().get_mut().update(&uri);
            req.head_mut().uri = uri;
        }

        self.service.call(req)
    }
}

#[cfg(test)]
mod tests {
    use actix_service::IntoService;

    use super::*;
    use crate::{
        dev::ServiceRequest,
        guard::fn_guard,
        test::{call_service, init_service, TestRequest},
        web, App, HttpResponse,
    };

    #[actix_rt::test]
    async fn test_wrap() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath::default())
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something").to(HttpResponse::Ok))
                .service(
                    web::resource("/v2/something")
                        .guard(fn_guard(|req| req.uri.query() == Some("query=test")))
                        .to(HttpResponse::Ok),
                ),
        )
        .await;

        let test_uris = vec![
            "/",
            "/?query=test",
            "///",
            "/v1//something",
            "/v1//something////",
            "//v1/something",
            "//v1//////something",
            "/v2//something?query=test",
            "/v2//something////?query=test",
            "//v2/something?query=test",
            "//v2//////something?query=test",
        ];

        for uri in test_uris {
            let req = TestRequest::with_uri(uri).to_request();
            let res = call_service(&app, req).await;
            assert!(res.status().is_success(), "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn trim_trailing_slashes() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath(TrailingSlash::Trim))
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something").to(HttpResponse::Ok))
                .service(
                    web::resource("/v2/something")
                        .guard(fn_guard(|req| req.uri.query() == Some("query=test")))
                        .to(HttpResponse::Ok),
                ),
        )
        .await;

        let test_uris = vec![
            "/",
            "///",
            "/v1/something",
            "/v1/something/",
            "/v1/something////",
            "//v1//something",
            "//v1//something//",
            "/v2/something?query=test",
            "/v2/something/?query=test",
            "/v2/something////?query=test",
            "//v2//something?query=test",
            "//v2//something//?query=test",
        ];

        for uri in test_uris {
            let req = TestRequest::with_uri(uri).to_request();
            let res = call_service(&app, req).await;
            assert!(res.status().is_success(), "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn trim_root_trailing_slashes_with_query() {
        let app = init_service(
            App::new().wrap(NormalizePath(TrailingSlash::Trim)).service(
                web::resource("/")
                    .guard(fn_guard(|req| req.uri.query() == Some("query=test")))
                    .to(HttpResponse::Ok),
            ),
        )
        .await;

        let test_uris = vec!["/?query=test", "//?query=test", "///?query=test"];

        for uri in test_uris {
            let req = TestRequest::with_uri(uri).to_request();
            let res = call_service(&app, req).await;
            assert!(res.status().is_success(), "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn ensure_trailing_slash() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath(TrailingSlash::Always))
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something/").to(HttpResponse::Ok))
                .service(
                    web::resource("/v2/something/")
                        .guard(fn_guard(|req| req.uri.query() == Some("query=test")))
                        .to(HttpResponse::Ok),
                ),
        )
        .await;

        let test_uris = vec![
            "/",
            "///",
            "/v1/something",
            "/v1/something/",
            "/v1/something////",
            "//v1//something",
            "//v1//something//",
            "/v2/something?query=test",
            "/v2/something/?query=test",
            "/v2/something////?query=test",
            "//v2//something?query=test",
            "//v2//something//?query=test",
        ];

        for uri in test_uris {
            let req = TestRequest::with_uri(uri).to_request();
            let res = call_service(&app, req).await;
            assert!(res.status().is_success(), "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn ensure_root_trailing_slash_with_query() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath(TrailingSlash::Always))
                .service(
                    web::resource("/")
                        .guard(fn_guard(|req| req.uri.query() == Some("query=test")))
                        .to(HttpResponse::Ok),
                ),
        )
        .await;

        let test_uris = vec!["/?query=test", "//?query=test", "///?query=test"];

        for uri in test_uris {
            let req = TestRequest::with_uri(uri).to_request();
            let res = call_service(&app, req).await;
            assert!(res.status().is_success(), "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn keep_trailing_slash_unchanged() {
        let app = init_service(
            App::new()
                .wrap(NormalizePath(TrailingSlash::MergeOnly))
                .service(web::resource("/").to(HttpResponse::Ok))
                .service(web::resource("/v1/something").to(HttpResponse::Ok))
                .service(web::resource("/v1/").to(HttpResponse::Ok))
                .service(
                    web::resource("/v2/something")
                        .guard(fn_guard(|req| req.uri.query() == Some("query=test")))
                        .to(HttpResponse::Ok),
                ),
        )
        .await;

        let tests = vec![
            ("/", true), // root paths should still work
            ("/?query=test", true),
            ("///", true),
            ("/v1/something////", false),
            ("/v1/something/", false),
            ("//v1//something", true),
            ("/v1/", true),
            ("/v1", false),
            ("/v1////", true),
            ("//v1//", true),
            ("///v1", false),
            ("/v2/something?query=test", true),
            ("/v2/something/?query=test", false),
            ("/v2/something//?query=test", false),
            ("//v2//something?query=test", true),
        ];

        for (uri, success) in tests {
            let req = TestRequest::with_uri(uri).to_request();
            let res = call_service(&app, req).await;
            assert_eq!(res.status().is_success(), success, "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn test_in_place_normalization() {
        let srv = |req: ServiceRequest| {
            assert_eq!("/v1/something", req.path());
            ready(Ok(req.into_response(HttpResponse::Ok().finish())))
        };

        let normalize = NormalizePath::default()
            .new_transform(srv.into_service())
            .await
            .unwrap();

        let test_uris = vec![
            "/v1//something////",
            "///v1/something",
            "//v1///something",
            "/v1//something",
        ];

        for uri in test_uris {
            let req = TestRequest::with_uri(uri).to_srv_request();
            let res = normalize.call(req).await.unwrap();
            assert!(res.status().is_success(), "Failed uri: {}", uri);
        }
    }

    #[actix_rt::test]
    async fn should_normalize_nothing() {
        const URI: &str = "/v1/something";

        let srv = |req: ServiceRequest| {
            assert_eq!(URI, req.path());
            ready(Ok(req.into_response(HttpResponse::Ok().finish())))
        };

        let normalize = NormalizePath::default()
            .new_transform(srv.into_service())
            .await
            .unwrap();

        let req = TestRequest::with_uri(URI).to_srv_request();
        let res = normalize.call(req).await.unwrap();
        assert!(res.status().is_success());
    }

    #[actix_rt::test]
    async fn should_normalize_no_trail() {
        let srv = |req: ServiceRequest| {
            assert_eq!("/v1/something", req.path());
            ready(Ok(req.into_response(HttpResponse::Ok().finish())))
        };

        let normalize = NormalizePath::default()
            .new_transform(srv.into_service())
            .await
            .unwrap();

        let req = TestRequest::with_uri("/v1/something/").to_srv_request();
        let res = normalize.call(req).await.unwrap();
        assert!(res.status().is_success());
    }
}
<|endoftext|>


<commit_before>use crate::helpers::{
    check_content_type_for_json, parse_pubkey, publish_attestation_to_network,
    publish_beacon_block_to_network,
};
use crate::response_builder::ResponseBuilder;
use crate::{ApiError, ApiResult, BoxFut, NetworkChannel, UrlQuery};
use beacon_chain::{
    AttestationProcessingOutcome, BeaconChain, BeaconChainTypes, BlockProcessingOutcome,
};
use bls::PublicKey;
use futures::future::Future;
use futures::stream::Stream;
use hyper::{Body, Request};
use serde::{Deserialize, Serialize};
use slog::{info, warn, Logger};
use ssz_derive::{Decode, Encode};
use std::sync::Arc;
use types::beacon_state::EthSpec;
use types::{Attestation, BeaconBlock, CommitteeIndex, Epoch, RelativeEpoch, Slot};

#[derive(PartialEq, Debug, Serialize, Deserialize, Clone)]
pub struct ValidatorDuty {
    /// The validator's BLS public key, uniquely identifying them. _48-bytes, hex encoded with 0x prefix, case insensitive._
    pub validator_pubkey: PublicKey,
    /// The slot at which the validator must attest.
    pub attestation_slot: Option<Slot>,
    /// The index of the committee within `slot` of which the validator is a member.
    pub attestation_committee_index: Option<CommitteeIndex>,
    /// The position of the validator in the committee.
    pub attestation_committee_position: Option<usize>,
    /// The slot in which a validator must propose a block, or `null` if block production is not required.
    pub block_proposal_slot: Option<Slot>,
}

#[derive(PartialEq, Debug, Serialize, Deserialize, Clone, Encode, Decode)]
pub struct BulkValidatorDutiesRequest {
    pub epoch: Epoch,
    pub pubkeys: Vec<PublicKey>,
}

/// HTTP Handler to retrieve a the duties for a set of validators during a particular epoch. This
/// method allows for collecting bulk sets of validator duties without risking exceeding the max
/// URL length with query pairs.
pub fn post_validator_duties<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> BoxFut {
    let response_builder = ResponseBuilder::new(&req);

    let future = req
        .into_body()
        .concat2()
        .map_err(|e| ApiError::ServerError(format!("Unable to get request body: {:?}", e)))
        .and_then(|chunks| {
            serde_json::from_slice::<BulkValidatorDutiesRequest>(&chunks).map_err(|e| {
                ApiError::BadRequest(format!(
                    "Unable to parse JSON into BulkValidatorDutiesRequest: {:?}",
                    e
                ))
            })
        })
        .and_then(|bulk_request| {
            return_validator_duties(beacon_chain, bulk_request.epoch, bulk_request.pubkeys)
        })
        .and_then(|duties| response_builder?.body_no_ssz(&duties));

    Box::new(future)
}

/// HTTP Handler to retrieve a the duties for a set of validators during a particular epoch
///
/// The given `epoch` must be within one epoch of the current epoch.
pub fn get_validator_duties<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> ApiResult {
    let query = UrlQuery::from_request(&req)?;

    let epoch = query.epoch()?;
    let validator_pubkeys = query
        .all_of("validator_pubkeys")?
        .iter()
        .map(|validator_pubkey_str| parse_pubkey(validator_pubkey_str))
        .collect::<Result<_, _>>()?;

    let duties = return_validator_duties(beacon_chain, epoch, validator_pubkeys)?;

    ResponseBuilder::new(&req)?.body_no_ssz(&duties)
}

fn return_validator_duties<T: BeaconChainTypes>(
    beacon_chain: Arc<BeaconChain<T>>,
    epoch: Epoch,
    validator_pubkeys: Vec<PublicKey>,
) -> Result<Vec<ValidatorDuty>, ApiError> {
    let mut state = beacon_chain
        .state_at_slot(epoch.start_slot(T::EthSpec::slots_per_epoch()))
        .map_err(|e| {
            ApiError::ServerError(format!("Unable to load state for epoch {}: {:?}", epoch, e))
        })?;

    let current_epoch = state.current_epoch();
    let relative_epoch = RelativeEpoch::from_epoch(current_epoch, epoch).map_err(|_| {
        ApiError::BadRequest(format!(
            "Epoch must be within one epoch of the current epoch",
        ))
    })?;

    state
        .build_committee_cache(relative_epoch, &beacon_chain.spec)
        .map_err(|e| ApiError::ServerError(format!("Unable to build committee cache: {:?}", e)))?;
    state
        .update_pubkey_cache()
        .map_err(|e| ApiError::ServerError(format!("Unable to build pubkey cache: {:?}", e)))?;

    // Get a list of all validators for this epoch.
    //
    // Used for quickly determining the slot for a proposer.
    let validator_proposers: Vec<(usize, Slot)> = epoch
        .slot_iter(T::EthSpec::slots_per_epoch())
        .map(|slot| {
            state
                .get_beacon_proposer_index(slot, &beacon_chain.spec)
                .map(|i| (i, slot))
                .map_err(|e| {
                    ApiError::ServerError(format!(
                        "Unable to get proposer index for validator: {:?}",
                        e
                    ))
                })
        })
        .collect::<Result<Vec<_>, _>>()?;

    validator_pubkeys
        .into_iter()
        .map(|validator_pubkey| {
            if let Some(validator_index) =
                state.get_validator_index(&validator_pubkey).map_err(|e| {
                    ApiError::ServerError(format!("Unable to read pubkey cache: {:?}", e))
                })?
            {
                let duties = state
                    .get_attestation_duties(validator_index, relative_epoch)
                    .map_err(|e| {
                        ApiError::ServerError(format!(
                            "Unable to obtain attestation duties: {:?}",
                            e
                        ))
                    })?;

                let block_proposal_slot = validator_proposers
                    .iter()
                    .find(|(i, _slot)| validator_index == *i)
                    .map(|(_i, slot)| *slot);

                Ok(ValidatorDuty {
                    validator_pubkey,
                    attestation_slot: duties.map(|d| d.slot),
                    attestation_committee_index: duties.map(|d| d.index),
                    attestation_committee_position: duties.map(|d| d.committee_position),
                    block_proposal_slot,
                })
            } else {
                Ok(ValidatorDuty {
                    validator_pubkey,
                    attestation_slot: None,
                    attestation_committee_index: None,
                    attestation_committee_position: None,
                    block_proposal_slot: None,
                })
            }
        })
        .collect::<Result<Vec<_>, ApiError>>()
}

/// HTTP Handler to produce a new BeaconBlock from the current state, ready to be signed by a validator.
pub fn get_new_beacon_block<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> ApiResult {
    let query = UrlQuery::from_request(&req)?;

    let slot = query.slot()?;
    let randao_reveal = query.randao_reveal()?;

    let (new_block, _state) = beacon_chain
        .produce_block(randao_reveal, slot)
        .map_err(|e| {
            ApiError::ServerError(format!(
                "Beacon node is not able to produce a block: {:?}",
                e
            ))
        })?;

    ResponseBuilder::new(&req)?.body(&new_block)
}

/// HTTP Handler to publish a BeaconBlock, which has been signed by a validator.
pub fn publish_beacon_block<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
    network_chan: NetworkChannel,
    log: Logger,
) -> BoxFut {
    try_future!(check_content_type_for_json(&req));
    let response_builder = ResponseBuilder::new(&req);

    let body = req.into_body();
    Box::new(body
        .concat2()
        .map_err(|e| ApiError::ServerError(format!("Unable to get request body: {:?}",e)))
        .and_then(|chunks| {
            serde_json::from_slice(&chunks).map_err(|e| ApiError::BadRequest(format!("Unable to parse JSON into BeaconBlock: {:?}",e)))
        })
        .and_then(move |block: BeaconBlock<T::EthSpec>| {
            let slot = block.slot;
            match beacon_chain.process_block(block.clone()) {
                Ok(BlockProcessingOutcome::Processed { block_root }) => {
                    // Block was processed, publish via gossipsub
                    info!(log, "Processed valid block from API, transmitting to network."; "block_slot" => slot, "block_root" => format!("{}", block_root));
                    publish_beacon_block_to_network::<T>(network_chan, block)
                }
                Ok(outcome) => {
                    warn!(log, "BeaconBlock could not be processed, but is being sent to the network anyway."; "outcome" => format!("{:?}", outcome));
                    publish_beacon_block_to_network::<T>(network_chan, block)?;
                    Err(ApiError::ProcessingError(format!(
                        "The BeaconBlock could not be processed, but has still been published: {:?}",
                        outcome
                    )))
                }
                Err(e) => {
                    Err(ApiError::ServerError(format!(
                        "Error while processing block: {:?}",
                        e
                    )))
                }
            }
        }).and_then(|_| {
            response_builder?.body_no_ssz(&())
        }))
}

/// HTTP Handler to produce a new Attestation from the current state, ready to be signed by a validator.
pub fn get_new_attestation<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> ApiResult {
    let query = UrlQuery::from_request(&req)?;

    let slot = query.slot()?;
    let index = query.committee_index()?;

    let attestation = beacon_chain
        .produce_attestation(slot, index)
        .map_err(|e| ApiError::BadRequest(format!("Unable to produce attestation: {:?}", e)))?;

    ResponseBuilder::new(&req)?.body(&attestation)
}

/// HTTP Handler to publish an Attestation, which has been signed by a validator.
pub fn publish_attestation<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
    network_chan: NetworkChannel,
    log: Logger,
) -> BoxFut {
    try_future!(check_content_type_for_json(&req));
    let response_builder = ResponseBuilder::new(&req);

    Box::new(req
        .into_body()
        .concat2()
        .map_err(|e| ApiError::ServerError(format!("Unable to get request body: {:?}",e)))
        .map(|chunk| chunk.iter().cloned().collect::<Vec<u8>>())
        .and_then(|chunks| {
            serde_json::from_slice(&chunks.as_slice()).map_err(|e| {
                ApiError::BadRequest(format!(
                    "Unable to deserialize JSON into a BeaconBlock: {:?}",
                    e
                ))
            })
        })
        .and_then(move |attestation: Attestation<T::EthSpec>| {
            match beacon_chain.process_attestation(attestation.clone()) {
                Ok(AttestationProcessingOutcome::Processed) => {
                    // Block was processed, publish via gossipsub
                    info!(log, "Processed valid attestation from API, transmitting to network.");
                    publish_attestation_to_network::<T>(network_chan, attestation)
                }
                Ok(outcome) => {
                    warn!(log, "Attestation could not be processed, but is being sent to the network anyway."; "outcome" => format!("{:?}", outcome));
                    publish_attestation_to_network::<T>(network_chan, attestation)?;
                    Err(ApiError::ProcessingError(format!(
                        "The Attestation could not be processed, but has still been published: {:?}",
                        outcome
                    )))
                }
                Err(e) => {
                    Err(ApiError::ServerError(format!(
                        "Error while processing attestation: {:?}",
                        e
                    )))
                }
            }
        }).and_then(|_| {
        response_builder?.body_no_ssz(&())
    }))
}
<commit_msg>optimize state loading (#636)<commit_after>use crate::helpers::{
    check_content_type_for_json, parse_pubkey, publish_attestation_to_network,
    publish_beacon_block_to_network,
};
use crate::response_builder::ResponseBuilder;
use crate::{ApiError, ApiResult, BoxFut, NetworkChannel, UrlQuery};
use beacon_chain::{
    AttestationProcessingOutcome, BeaconChain, BeaconChainTypes, BlockProcessingOutcome,
};
use bls::PublicKey;
use futures::future::Future;
use futures::stream::Stream;
use hyper::{Body, Request};
use serde::{Deserialize, Serialize};
use slog::{info, warn, Logger};
use ssz_derive::{Decode, Encode};
use std::sync::Arc;
use types::beacon_state::EthSpec;
use types::{Attestation, BeaconBlock, CommitteeIndex, Epoch, RelativeEpoch, Slot};

#[derive(PartialEq, Debug, Serialize, Deserialize, Clone)]
pub struct ValidatorDuty {
    /// The validator's BLS public key, uniquely identifying them. _48-bytes, hex encoded with 0x prefix, case insensitive._
    pub validator_pubkey: PublicKey,
    /// The slot at which the validator must attest.
    pub attestation_slot: Option<Slot>,
    /// The index of the committee within `slot` of which the validator is a member.
    pub attestation_committee_index: Option<CommitteeIndex>,
    /// The position of the validator in the committee.
    pub attestation_committee_position: Option<usize>,
    /// The slot in which a validator must propose a block, or `null` if block production is not required.
    pub block_proposal_slot: Option<Slot>,
}

#[derive(PartialEq, Debug, Serialize, Deserialize, Clone, Encode, Decode)]
pub struct BulkValidatorDutiesRequest {
    pub epoch: Epoch,
    pub pubkeys: Vec<PublicKey>,
}

/// HTTP Handler to retrieve a the duties for a set of validators during a particular epoch. This
/// method allows for collecting bulk sets of validator duties without risking exceeding the max
/// URL length with query pairs.
pub fn post_validator_duties<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> BoxFut {
    let response_builder = ResponseBuilder::new(&req);

    let future = req
        .into_body()
        .concat2()
        .map_err(|e| ApiError::ServerError(format!("Unable to get request body: {:?}", e)))
        .and_then(|chunks| {
            serde_json::from_slice::<BulkValidatorDutiesRequest>(&chunks).map_err(|e| {
                ApiError::BadRequest(format!(
                    "Unable to parse JSON into BulkValidatorDutiesRequest: {:?}",
                    e
                ))
            })
        })
        .and_then(|bulk_request| {
            return_validator_duties(beacon_chain, bulk_request.epoch, bulk_request.pubkeys)
        })
        .and_then(|duties| response_builder?.body_no_ssz(&duties));

    Box::new(future)
}

/// HTTP Handler to retrieve a the duties for a set of validators during a particular epoch
///
/// The given `epoch` must be within one epoch of the current epoch.
pub fn get_validator_duties<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> ApiResult {
    let query = UrlQuery::from_request(&req)?;

    let epoch = query.epoch()?;
    let validator_pubkeys = query
        .all_of("validator_pubkeys")?
        .iter()
        .map(|validator_pubkey_str| parse_pubkey(validator_pubkey_str))
        .collect::<Result<_, _>>()?;

    let duties = return_validator_duties(beacon_chain, epoch, validator_pubkeys)?;

    ResponseBuilder::new(&req)?.body_no_ssz(&duties)
}

fn return_validator_duties<T: BeaconChainTypes>(
    beacon_chain: Arc<BeaconChain<T>>,
    epoch: Epoch,
    validator_pubkeys: Vec<PublicKey>,
) -> Result<Vec<ValidatorDuty>, ApiError> {
    let head_state = beacon_chain.head().beacon_state;
    let head_epoch = head_state.current_epoch();

    let relative_epoch = RelativeEpoch::from_epoch(head_epoch, epoch);
    let mut state = if relative_epoch.is_err() {
        head_state
    } else {
        match beacon_chain.state_at_slot(epoch.start_slot(T::EthSpec::slots_per_epoch())) {
            Ok(state) => state,
            Err(e) => {
                return Err(ApiError::ServerError(format!(
                    "Unable to load state for epoch {}: {:?}",
                    epoch, e
                )))
            }
        }
    };

    let relative_epoch = relative_epoch.or_else(|_| {
        RelativeEpoch::from_epoch(state.current_epoch(), epoch).map_err(|_| {
            ApiError::BadRequest(String::from(
                "Epoch must be within one epoch of the current epoch",
            ))
        })
    })?;

    state
        .build_committee_cache(relative_epoch, &beacon_chain.spec)
        .map_err(|e| ApiError::ServerError(format!("Unable to build committee cache: {:?}", e)))?;
    state
        .update_pubkey_cache()
        .map_err(|e| ApiError::ServerError(format!("Unable to build pubkey cache: {:?}", e)))?;

    // Get a list of all validators for this epoch.
    //
    // Used for quickly determining the slot for a proposer.
    let validator_proposers: Vec<(usize, Slot)> = epoch
        .slot_iter(T::EthSpec::slots_per_epoch())
        .map(|slot| {
            state
                .get_beacon_proposer_index(slot, &beacon_chain.spec)
                .map(|i| (i, slot))
                .map_err(|e| {
                    ApiError::ServerError(format!(
                        "Unable to get proposer index for validator: {:?}",
                        e
                    ))
                })
        })
        .collect::<Result<Vec<_>, _>>()?;

    validator_pubkeys
        .into_iter()
        .map(|validator_pubkey| {
            if let Some(validator_index) =
                state.get_validator_index(&validator_pubkey).map_err(|e| {
                    ApiError::ServerError(format!("Unable to read pubkey cache: {:?}", e))
                })?
            {
                let duties = state
                    .get_attestation_duties(validator_index, relative_epoch)
                    .map_err(|e| {
                        ApiError::ServerError(format!(
                            "Unable to obtain attestation duties: {:?}",
                            e
                        ))
                    })?;

                let block_proposal_slot = validator_proposers
                    .iter()
                    .find(|(i, _slot)| validator_index == *i)
                    .map(|(_i, slot)| *slot);

                Ok(ValidatorDuty {
                    validator_pubkey,
                    attestation_slot: duties.map(|d| d.slot),
                    attestation_committee_index: duties.map(|d| d.index),
                    attestation_committee_position: duties.map(|d| d.committee_position),
                    block_proposal_slot,
                })
            } else {
                Ok(ValidatorDuty {
                    validator_pubkey,
                    attestation_slot: None,
                    attestation_committee_index: None,
                    attestation_committee_position: None,
                    block_proposal_slot: None,
                })
            }
        })
        .collect::<Result<Vec<_>, ApiError>>()
}

/// HTTP Handler to produce a new BeaconBlock from the current state, ready to be signed by a validator.
pub fn get_new_beacon_block<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> ApiResult {
    let query = UrlQuery::from_request(&req)?;

    let slot = query.slot()?;
    let randao_reveal = query.randao_reveal()?;

    let (new_block, _state) = beacon_chain
        .produce_block(randao_reveal, slot)
        .map_err(|e| {
            ApiError::ServerError(format!(
                "Beacon node is not able to produce a block: {:?}",
                e
            ))
        })?;

    ResponseBuilder::new(&req)?.body(&new_block)
}

/// HTTP Handler to publish a BeaconBlock, which has been signed by a validator.
pub fn publish_beacon_block<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
    network_chan: NetworkChannel,
    log: Logger,
) -> BoxFut {
    try_future!(check_content_type_for_json(&req));
    let response_builder = ResponseBuilder::new(&req);

    let body = req.into_body();
    Box::new(body
        .concat2()
        .map_err(|e| ApiError::ServerError(format!("Unable to get request body: {:?}",e)))
        .and_then(|chunks| {
            serde_json::from_slice(&chunks).map_err(|e| ApiError::BadRequest(format!("Unable to parse JSON into BeaconBlock: {:?}",e)))
        })
        .and_then(move |block: BeaconBlock<T::EthSpec>| {
            let slot = block.slot;
            match beacon_chain.process_block(block.clone()) {
                Ok(BlockProcessingOutcome::Processed { block_root }) => {
                    // Block was processed, publish via gossipsub
                    info!(log, "Processed valid block from API, transmitting to network."; "block_slot" => slot, "block_root" => format!("{}", block_root));
                    publish_beacon_block_to_network::<T>(network_chan, block)
                }
                Ok(outcome) => {
                    warn!(log, "BeaconBlock could not be processed, but is being sent to the network anyway."; "outcome" => format!("{:?}", outcome));
                    publish_beacon_block_to_network::<T>(network_chan, block)?;
                    Err(ApiError::ProcessingError(format!(
                        "The BeaconBlock could not be processed, but has still been published: {:?}",
                        outcome
                    )))
                }
                Err(e) => {
                    Err(ApiError::ServerError(format!(
                        "Error while processing block: {:?}",
                        e
                    )))
                }
            }
        }).and_then(|_| {
            response_builder?.body_no_ssz(&())
        }))
}

/// HTTP Handler to produce a new Attestation from the current state, ready to be signed by a validator.
pub fn get_new_attestation<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
) -> ApiResult {
    let query = UrlQuery::from_request(&req)?;

    let slot = query.slot()?;
    let index = query.committee_index()?;

    let attestation = beacon_chain
        .produce_attestation(slot, index)
        .map_err(|e| ApiError::BadRequest(format!("Unable to produce attestation: {:?}", e)))?;

    ResponseBuilder::new(&req)?.body(&attestation)
}

/// HTTP Handler to publish an Attestation, which has been signed by a validator.
pub fn publish_attestation<T: BeaconChainTypes>(
    req: Request<Body>,
    beacon_chain: Arc<BeaconChain<T>>,
    network_chan: NetworkChannel,
    log: Logger,
) -> BoxFut {
    try_future!(check_content_type_for_json(&req));
    let response_builder = ResponseBuilder::new(&req);

    Box::new(req
        .into_body()
        .concat2()
        .map_err(|e| ApiError::ServerError(format!("Unable to get request body: {:?}",e)))
        .map(|chunk| chunk.iter().cloned().collect::<Vec<u8>>())
        .and_then(|chunks| {
            serde_json::from_slice(&chunks.as_slice()).map_err(|e| {
                ApiError::BadRequest(format!(
                    "Unable to deserialize JSON into a BeaconBlock: {:?}",
                    e
                ))
            })
        })
        .and_then(move |attestation: Attestation<T::EthSpec>| {
            match beacon_chain.process_attestation(attestation.clone()) {
                Ok(AttestationProcessingOutcome::Processed) => {
                    // Block was processed, publish via gossipsub
                    info!(log, "Processed valid attestation from API, transmitting to network.");
                    publish_attestation_to_network::<T>(network_chan, attestation)
                }
                Ok(outcome) => {
                    warn!(log, "Attestation could not be processed, but is being sent to the network anyway."; "outcome" => format!("{:?}", outcome));
                    publish_attestation_to_network::<T>(network_chan, attestation)?;
                    Err(ApiError::ProcessingError(format!(
                        "The Attestation could not be processed, but has still been published: {:?}",
                        outcome
                    )))
                }
                Err(e) => {
                    Err(ApiError::ServerError(format!(
                        "Error while processing attestation: {:?}",
                        e
                    )))
                }
            }
        }).and_then(|_| {
        response_builder?.body_no_ssz(&())
    }))
}
<|endoftext|>


<commit_before>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! Node manager.

use super::{encode_contract_name, to_address_vec, to_low_u64_vec};
use super::ContractCallExt;
use libexecutor::executor::Executor;
use rand::{Rng, SeedableRng, StdRng};
use rustc_hex::ToHex;
use std::cmp::Ordering;
use std::iter;
use std::str::FromStr;
use util::{Address, H160};

const LIST_NODE: &'static [u8] = &*b"listNode()";
const LIST_STAKE: &'static [u8] = *&b"listStake()";
/// Each epoch is divided into 1000 slots, each slot represent one opportunity of block proposer
const EPOCH: u64 = 1000;

lazy_static! {
    static ref LIST_NODE_ENCODED: Vec<u8> = encode_contract_name(LIST_NODE);
    static ref LIST_STAKE_ENCODED: Vec<u8> = encode_contract_name(LIST_STAKE);
    static ref CONTRACT_ADDRESS: H160 = H160::from_str("00000000000000000000000000000000013241a2").unwrap();
}

pub fn party_seats<T>(parties: Vec<T>, seats: &[u64]) -> Vec<T>
where
    T: Clone,
{
    let mut party_seats: Vec<T> = Vec::new();
    let seats_len = seats.len();
    for (index, party) in parties.into_iter().enumerate() {
        if index < seats_len {
            party_seats.extend(iter::repeat(party).take(seats[index] as usize));
        }
    }
    party_seats
}

/// https://en.wikipedia.org/wiki/Largest_remainder_method
/// Hare quota
pub fn largest_remainder_electoral(votes: &Vec<u64>, seat_number: u64) -> Vec<u64> {
    let total = votes.iter().fold(0, |acc, &x| acc + x);
    // TODO: PLS fix me!!!
    let hare_quota = total as f64 / seat_number as f64;

    let votes_quota: Vec<f64> = votes
        .iter()
        .map(|vote| (*vote as f64) / hare_quota)
        .collect();

    // calculate automatic seats first
    let mut seats: Vec<u64> = votes_quota.iter().map(|v| v.floor() as u64).collect();

    let mut remainders: Vec<(f64, u64)> = votes_quota
        .iter()
        .enumerate()
        .map(|(i, v)| (v.fract(), i as u64))
        .collect();

    let remaining_seat: u64 = seat_number - seats.iter().fold(0, |acc, &x| acc + x);

    remainders.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal));
    remainders.reverse();

    let highest_remainder_seats: Vec<u64> = remainders
        .iter()
        .take(remaining_seat as usize)
        .map(|v| v.1)
        .collect();

    for index in highest_remainder_seats {
        seats[index as usize] += 1;
    }

    seats
}

pub struct NodeManager;

impl NodeManager {
    pub fn nodes(executor: &Executor) -> Vec<Address> {
        let output = executor.call_method_latest(&*CONTRACT_ADDRESS, &*LIST_NODE_ENCODED.as_slice());
        trace!(
            "node manager output: {:?}",
            ToHex::to_hex(output.as_slice())
        );

        let nodes: Vec<Address> = to_address_vec(&output);
        trace!("node manager nodes: {:?}", nodes);
        nodes
    }

    pub fn shuffle_node<T>(node_vec: &mut Vec<T>, rng_seed: u64) {
        let seed: &[_] = &[rng_seed as usize];
        let mut rng: StdRng = SeedableRng::from_seed(seed);

        for i in 0..node_vec.len() {
            let j: usize = rng.gen::<usize>() % (i + 1);
            node_vec.swap(i, j);
        }
    }

    pub fn stakes(executor: &Executor) -> Vec<u64> {
        let output = executor.call_method_latest(&*CONTRACT_ADDRESS, &*LIST_STAKE_ENCODED.as_slice());
        trace!("stakes output: {:?}", ToHex::to_hex(output.as_slice()));

        let stakes: Vec<u64> = to_low_u64_vec(&output);
        trace!("node manager stakes: {:?}", stakes);
        stakes
    }

    pub fn stake_nodes(exeuctor: &Executor) -> Vec<Address> {
        let nodes = NodeManager::nodes(&exeuctor);
        let stakes = NodeManager::stakes(&exeuctor);
        let total = stakes.iter().fold(0, |acc, &x| acc + x);
        let stake_nodes: Vec<Address>;

        if total == 0 {
            stake_nodes = nodes;
        } else {
            let total_seats = largest_remainder_electoral(&stakes, EPOCH);
            stake_nodes = party_seats(nodes, &total_seats);
        }
        stake_nodes
    }
}

#[cfg(test)]
mod tests {
    extern crate logger;
    extern crate mktemp;

    use super::{largest_remainder_electoral, party_seats, NodeManager};
    use std::str::FromStr;
    use tests::helpers::init_executor;
    use util::H160;

    #[test]
    fn test_node_manager_contract() {
        let executor = init_executor();
        let nodes = NodeManager::nodes(&executor);

        assert_eq!(
            nodes,
            vec![
                H160::from_str("50ad2b9d6946d9c75ae978534043e3021ee1bfb1").unwrap(),
                H160::from_str("eeb3a71c4046f63a941013f826fccc503be26b77").unwrap(),
                H160::from_str("a2bbb65d4f8c3ada29f7471abe416e18061127f3").unwrap(),
                H160::from_str("72eb1e258c9cdccebb7b62930a35cfb6ef4cd24b").unwrap(),
            ]
        )
    }

    #[test]
    fn test_largest_remainder_electoral() {
        let votes = vec![47_000, 16_000, 15_800, 12_000, 6_100, 3_100];
        let seats = 10;
        let total_seats = largest_remainder_electoral(&votes, seats);
        assert_eq!(vec![5, 2, 1, 1, 1, 0], total_seats);
    }

    #[test]
    fn test_largest_remainder_electoral_again() {
        let votes = vec![1500, 1500, 900, 500, 500, 200];
        let seats = 25;
        let total_seats = largest_remainder_electoral(&votes, seats);
        assert_eq!(vec![7, 7, 4, 3, 3, 1], total_seats);

        let new_seats = 26;
        let new_total_seats = largest_remainder_electoral(&votes, new_seats);
        assert_eq!(vec![8, 8, 5, 2, 2, 1], new_total_seats);
    }

    #[test]
    fn test_party_seats() {
        let parties = vec!["a", "b", "c"];
        let seats = vec![3, 5, 2, 2, 1];
        assert_eq!(
            party_seats(parties, &seats),
            vec!["a", "a", "a", "b", "b", "b", "b", "b", "c", "c"]
        );

        let parties = vec!["a", "b"];
        let seats = vec![2, 1];
        assert_eq!(party_seats(parties, &seats), vec!["a", "a", "b"]);

        let parties = vec!["a", "b", "c"];
        let seats = vec![2, 2];
        assert_eq!(party_seats(parties, &seats), vec!["a", "a", "b", "b"]);
    }
}
<commit_msg>[skip ci] skip ci<commit_after>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! Node manager.

use super::{encode_contract_name, to_address_vec, to_low_u64_vec};
use super::ContractCallExt;
use libexecutor::executor::Executor;
use rand::{Rng, SeedableRng, StdRng};
use rustc_hex::ToHex;
use std::cmp::Ordering;
use std::iter;
use std::str::FromStr;
use util::{Address, H160};

const LIST_NODE: &'static [u8] = &*b"listNode()";
const LIST_STAKE: &'static [u8] = *&b"listStake()";
/// Each epoch is divided into 1000 slots, each slot represent one opportunity of block proposer
const EPOCH: u64 = 1000;

lazy_static! {
    static ref LIST_NODE_ENCODED: Vec<u8> = encode_contract_name(LIST_NODE);
    static ref LIST_STAKE_ENCODED: Vec<u8> = encode_contract_name(LIST_STAKE);
    static ref CONTRACT_ADDRESS: H160 = H160::from_str("00000000000000000000000000000000013241a2").unwrap();
}

pub fn party_seats<T>(parties: Vec<T>, seats: &[u64]) -> Vec<T>
where
    T: Clone,
{
    let mut party_seats: Vec<T> = Vec::new();
    let seats_len = seats.len();
    for (index, party) in parties.into_iter().enumerate() {
        if index < seats_len {
            party_seats.extend(iter::repeat(party).take(seats[index] as usize));
        }
    }
    party_seats
}

/// https://en.wikipedia.org/wiki/Largest_remainder_method
/// Hare quota
pub fn largest_remainder_electoral(votes: &Vec<u64>, seat_number: u64) -> Vec<u64> {
    let total = votes.iter().fold(0, |acc, &x| acc + x);
    // TODO: PLS fix me!!!
    let hare_quota = total as f64 / seat_number as f64;

    let votes_quota: Vec<f64> = votes
        .iter()
        .map(|vote| (*vote as f64) / hare_quota)
        .collect();

    // calculate automatic seats first
    let mut seats: Vec<u64> = votes_quota.iter().map(|v| v.floor() as u64).collect();

    let mut remainders: Vec<(f64, u64)> = votes_quota
        .iter()
        .enumerate()
        .map(|(i, v)| (v.fract(), i as u64))
        .collect();

    let remaining_seat: u64 = seat_number - seats.iter().fold(0, |acc, &x| acc + x);

    remainders.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal));
    remainders.reverse();

    let highest_remainder_seats: Vec<u64> = remainders
        .iter()
        .take(remaining_seat as usize)
        .map(|v| v.1)
        .collect();

    for index in highest_remainder_seats {
        seats[index as usize] += 1;
    }

    seats
}

pub struct NodeManager;

impl NodeManager {
    pub fn nodes(executor: &Executor) -> Vec<Address> {
        let output = executor.call_method_latest(&*CONTRACT_ADDRESS, &*LIST_NODE_ENCODED.as_slice());
        trace!(
            "node manager output: {:?}",
            ToHex::to_hex(output.as_slice())
        );

        let nodes: Vec<Address> = to_address_vec(&output);
        trace!("node manager nodes: {:?}", nodes);
        nodes
    }

    pub fn shuffle_node<T>(node_vec: &mut Vec<T>, rng_seed: u64) {
        let seed: &[_] = &[rng_seed as usize];
        let mut rng: StdRng = SeedableRng::from_seed(seed);

        for i in 0..node_vec.len() {
            let j: usize = rng.gen::<usize>() % (i + 1);
            node_vec.swap(i, j);
        }
    }

    pub fn stakes(executor: &Executor) -> Vec<u64> {
        let output = executor.call_method_latest(&*CONTRACT_ADDRESS, &*LIST_STAKE_ENCODED.as_slice());
        trace!("stakes output: {:?}", ToHex::to_hex(output.as_slice()));

        let stakes: Vec<u64> = to_low_u64_vec(&output);
        trace!("node manager stakes: {:?}", stakes);
        stakes
    }

    pub fn stake_nodes(exeuctor: &Executor) -> Vec<Address> {
        let nodes = NodeManager::nodes(&exeuctor);
        let stakes = NodeManager::stakes(&exeuctor);
        let total = stakes.iter().fold(0, |acc, &x| acc + x);
        let stake_nodes: Vec<Address>;

        if total == 0 {
            stake_nodes = nodes;
        } else {
            let total_seats = largest_remainder_electoral(&stakes, EPOCH);
            stake_nodes = party_seats(nodes, &total_seats);
        }

        stake_nodes
    }
}

#[cfg(test)]
mod tests {
    extern crate logger;
    extern crate mktemp;

    use super::{largest_remainder_electoral, party_seats, NodeManager};
    use std::str::FromStr;
    use tests::helpers::init_executor;
    use util::H160;

    #[test]
    fn test_node_manager_contract() {
        let executor = init_executor();
        let nodes = NodeManager::nodes(&executor);

        assert_eq!(
            nodes,
            vec![
                H160::from_str("50ad2b9d6946d9c75ae978534043e3021ee1bfb1").unwrap(),
                H160::from_str("eeb3a71c4046f63a941013f826fccc503be26b77").unwrap(),
                H160::from_str("a2bbb65d4f8c3ada29f7471abe416e18061127f3").unwrap(),
                H160::from_str("72eb1e258c9cdccebb7b62930a35cfb6ef4cd24b").unwrap(),
            ]
        )
    }

    #[test]
    fn test_largest_remainder_electoral() {
        let votes = vec![47_000, 16_000, 15_800, 12_000, 6_100, 3_100];
        let seats = 10;
        let total_seats = largest_remainder_electoral(&votes, seats);
        assert_eq!(vec![5, 2, 1, 1, 1, 0], total_seats);
    }

    #[test]
    fn test_largest_remainder_electoral_again() {
        let votes = vec![1500, 1500, 900, 500, 500, 200];
        let seats = 25;
        let total_seats = largest_remainder_electoral(&votes, seats);
        assert_eq!(vec![7, 7, 4, 3, 3, 1], total_seats);

        let new_seats = 26;
        let new_total_seats = largest_remainder_electoral(&votes, new_seats);
        assert_eq!(vec![8, 8, 5, 2, 2, 1], new_total_seats);
    }

    #[test]
    fn test_party_seats() {
        let parties = vec!["a", "b", "c"];
        let seats = vec![3, 5, 2, 2, 1];
        assert_eq!(
            party_seats(parties, &seats),
            vec!["a", "a", "a", "b", "b", "b", "b", "b", "c", "c"]
        );

        let parties = vec!["a", "b"];
        let seats = vec![2, 1];
        assert_eq!(party_seats(parties, &seats), vec!["a", "a", "b"]);

        let parties = vec!["a", "b", "c"];
        let seats = vec![2, 2];
        assert_eq!(party_seats(parties, &seats), vec!["a", "a", "b", "b"]);
    }
}
<|endoftext|>


<commit_before>//! # The Rust MessagePack Library
//!
//! RMP is a pure Rust [MessagePack](http://msgpack.org) implementation.
//!
//! MessagePack is an efficient binary serialization format.
//!
//! **Warning** this library is still in rapid development and everything may change until 1.0 comes.
//!
//! ## Usage
//!
//! To use `rmp`, first add this to your `Cargo.toml`:
//!
//! ```toml
//! [dependencies.rmp]
//! rmp = "0.7.2"
//! ```
//!
//! Then, add this to your crate root:
//!
//! ```rust
//! extern crate rmp as msgpack; // Or just `rmp`.
//! ```
//!
//! ## Features
//!
//! - **Convenient API**
//!
//!   RMP is designed to be lightweight and straightforward. There are low-level API, which gives you
//!   full control on data encoding/decoding process and makes no heap allocations. On the other hand
//!   there are high-level API, which provides you convenient interface using Rust standard library and
//!   compiler reflection, allowing to encode/decode structures using `derive` attribute.
//!
//! - **Zero-copy value decoding**
//!
//!   RMP allows to decode bytes from a buffer in a zero-copy manner easily and blazingly fast, while Rust
//!   static checks guarantees that the data will be valid until buffer lives.
//!
//! - **Clear error handling**
//!
//!   RMP's error system guarantees that you never receive an error enum with unreachable variant.
//!
//! - **Robust and tested**
//!
//!   This project is developed using TDD and CI, so any found bugs will be fixed without breaking
//!   existing functionality.

extern crate byteorder;

pub mod encode;
pub mod decode;

mod init;

pub mod value;

pub const MSGPACK_VERSION : u32 = 5;

pub use init::Marker;

pub use value::{Value, ValueRef};
<commit_msg>Doc(rmp): minor update.<commit_after>//! # The Rust MessagePack Library
//!
//! RMP is a pure Rust [MessagePack](http://msgpack.org) implementation.
//!
//! MessagePack is an efficient binary serialization format.
//!
//! **Warning** this library is still in rapid development and everything may change until 1.0 comes.
//!
//! ## Usage
//!
//! To use `rmp`, first add this to your `Cargo.toml`:
//!
//! ```toml
//! [dependencies.rmp]
//! rmp = "0.7.2"
//! ```
//!
//! Then, add this to your crate root:
//!
//! ```rust
//! extern crate rmp as msgpack; // Or just `rmp`.
//! ```
//!
//! ## Features
//!
//! - **Convenient API**
//!
//!   RMP is designed to be lightweight and straightforward. There are low-level API, which gives you
//!   full control on data encoding/decoding process and makes no heap allocations. On the other hand
//!   there are high-level API, which provides you convenient interface using Rust standard library and
//!   compiler reflection, allowing to encode/decode structures using `derive` attribute.
//!
//! - **Zero-copy value decoding**
//!
//!   RMP allows to decode bytes from a buffer in a zero-copy manner easily and blazingly fast, while Rust
//!   static checks guarantees that the data will be valid until buffer lives.
//!
//! - **Clear error handling**
//!
//!   RMP's error system guarantees that you never receive an error enum with unreachable variant.
//!
//! - **Robust and tested**
//!
//!   This project is developed using TDD and CI, so any found bugs will be fixed without breaking
//!   existing functionality.
//!
//! ## Detailed
//!

extern crate byteorder;

pub mod encode;
pub mod decode;

mod init;

pub mod value;

pub const MSGPACK_VERSION : u32 = 5;

pub use init::Marker;

pub use value::{Value, ValueRef};
<|endoftext|>