<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use proc_macro2::Span;
use syn::{
    parse::{Parse, ParseStream},
    LitStr, Token,
};
use syn::{Ident, Result as ParseResult};

#[derive(PartialEq, Clone, Debug, Hash)]
pub enum UnsafePolicy {
    AllFunctionsSafe,
    AllFunctionsUnsafe,
}

impl Parse for UnsafePolicy {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        if input.parse::<Option<Token![unsafe]>>()?.is_some() {
            return Ok(UnsafePolicy::AllFunctionsSafe);
        }
        let r = match input.parse::<Option<syn::Ident>>()? {
            Some(id) => {
                if id == "unsafe_ffi" {
                    Ok(UnsafePolicy::AllFunctionsSafe)
                } else {
                    Err(syn::Error::new(id.span(), "expected unsafe_ffi"))
                }
            }
            None => Ok(UnsafePolicy::AllFunctionsUnsafe),
        };
        if !input.is_empty() {
            return Err(syn::Error::new(
                Span::call_site(),
                "unexpected tokens within safety directive",
            ));
        }
        r
    }
}

/// Allowlist configuration.
#[derive(Hash, Debug)]
pub enum Allowlist {
    Unspecified(Vec<String>),
    All,
    Specific(Vec<String>),
}

impl Allowlist {
    pub(crate) fn push(&mut self, item: LitStr) -> ParseResult<()> {
        match self {
            Allowlist::Unspecified(ref mut uncommitted_list) => {
                let new_list = uncommitted_list
                    .drain(..)
                    .chain(std::iter::once(item.value()))
                    .collect();
                *self = Allowlist::Specific(new_list);
            }
            Allowlist::All => {
                return Err(syn::Error::new(
                    item.span(),
                    "use either generate!/generate_pod! or generate_all!, not both.",
                ))
            }
            Allowlist::Specific(list) => list.push(item.value()),
        };
        Ok(())
    }

    pub(crate) fn set_all(&mut self, ident: &Ident) -> ParseResult<()> {
        if matches!(self, Allowlist::Specific(..)) {
            return Err(syn::Error::new(
                ident.span(),
                "use either generate!/generate_pod! or generate_all!, not both.",
            ));
        }
        *self = Allowlist::All;
        Ok(())
    }
}

impl Default for Allowlist {
    fn default() -> Self {
        Allowlist::Unspecified(Vec::new())
    }
}

#[derive(Hash, Debug)]
pub struct Subclass {
    pub superclass: String,
    pub subclass: Ident,
}

#[derive(Hash, Debug)]
pub struct IncludeCppConfig {
    pub inclusions: Vec<String>,
    pub unsafe_policy: UnsafePolicy,
    pub parse_only: bool,
    pub exclude_impls: bool,
    pod_requests: Vec<String>,
    allowlist: Allowlist,
    blocklist: Vec<String>,
    exclude_utilities: bool,
    mod_name: Option<Ident>,
    pub rust_types: Vec<Ident>,
    pub subclasses: Vec<Subclass>,
}

impl Parse for IncludeCppConfig {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut parse_only = false;
        let mut exclude_impls = false;
        let mut unsafe_policy = UnsafePolicy::AllFunctionsUnsafe;
        let mut allowlist = Allowlist::default();
        let mut blocklist = Vec::new();
        let mut pod_requests = Vec::new();
        let mut rust_types = Vec::new();
        let mut exclude_utilities = false;
        let mut mod_name = None;
        let mut subclasses = Vec::new();

        while !input.is_empty() {
            let has_hexathorpe = input.parse::<Option<syn::token::Pound>>()?.is_some();
            let ident: syn::Ident = input.parse()?;
            if has_hexathorpe {
                if ident != "include" {
                    return Err(syn::Error::new(ident.span(), "expected include"));
                }
                let hdr: syn::LitStr = input.parse()?;
                inclusions.push(hdr.value());
            } else {
                input.parse::<Option<syn::token::Bang>>()?;
                if ident == "generate" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    allowlist.push(generate)?;
                } else if ident == "generate_pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate_pod: syn::LitStr = args.parse()?;
                    pod_requests.push(generate_pod.value());
                    allowlist.push(generate_pod)?;
                } else if ident == "pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let pod: syn::LitStr = args.parse()?;
                    pod_requests.push(pod.value());
                } else if ident == "block" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    blocklist.push(generate.value());
                } else if ident == "rust_type" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    rust_types.push(ident);
                } else if ident == "subclass" {
                    let args;
                    syn::parenthesized!(args in input);
                    let superclass: syn::LitStr = args.parse()?;
                    if superclass.value().contains("::") {
                        return Err(syn::Error::new(
                            superclass.span(),
                            "at this time, superclasses can't be in namespaces - see https://github.com/google/autocxx/issues/599",
                        ));
                    }
                    args.parse::<syn::token::Comma>()?;
                    let subclass: syn::Ident = args.parse()?;
                    subclasses.push(Subclass {
                        superclass: superclass.value(),
                        subclass,
                    });
                } else if ident == "parse_only" {
                    parse_only = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "exclude_impls" {
                    exclude_impls = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "generate_all" {
                    allowlist.set_all(&ident)?;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "name" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    mod_name = Some(ident);
                } else if ident == "exclude_utilities" {
                    exclude_utilities = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "safety" {
                    let args;
                    syn::parenthesized!(args in input);
                    unsafe_policy = args.parse()?;
                } else {
                    return Err(syn::Error::new(
                        ident.span(),
                        "expected generate, generate_pod, nested_type, safety or exclude_utilities",
                    ));
                }
            }
            if input.is_empty() {
                break;
            }
        }

        Ok(IncludeCppConfig {
            inclusions,
            unsafe_policy,
            parse_only,
            exclude_impls,
            pod_requests,
            rust_types,
            allowlist,
            blocklist,
            exclude_utilities,
            mod_name,
            subclasses,
        })
    }
}

fn swallow_parentheses(input: &ParseStream, latest_ident: &Ident) -> ParseResult<()> {
    let args;
    syn::parenthesized!(args in input);
    if args.is_empty() {
        Ok(())
    } else {
        Err(syn::Error::new(
            latest_ident.span(),
            "expected no arguments to directive",
        ))
    }
}

impl IncludeCppConfig {
    pub fn get_pod_requests(&self) -> &[String] {
        &self.pod_requests
    }

    pub fn get_mod_name(&self) -> Ident {
        self.mod_name
            .as_ref()
            .cloned()
            .unwrap_or_else(|| Ident::new("ffi", Span::call_site()))
    }

    /// Whether to avoid generating the standard helpful utility
    /// functions which we normally include in every mod.
    pub fn exclude_utilities(&self) -> bool {
        self.exclude_utilities
    }

    /// Items which the user has explicitly asked us to generate;
    /// we should raise an error if we weren't able to do so.
    pub fn must_generate_list(&self) -> Box<dyn Iterator<Item = String> + '_> {
        if let Allowlist::Specific(items) = &self.allowlist {
            Box::new(items.iter().chain(self.pod_requests.iter()).cloned())
        } else {
            Box::new(self.pod_requests.iter().cloned())
        }
    }

    /// The allowlist of items to be passed into bindgen, if any.
    pub fn bindgen_allowlist(&self) -> Option<Box<dyn Iterator<Item = String> + '_>> {
        match &self.allowlist {
            Allowlist::All => None,
            Allowlist::Specific(items) => Some(Box::new(
                items
                    .iter()
                    .chain(self.pod_requests.iter())
                    .cloned()
                    .chain(self.active_utilities())
                    .chain(
                        self.subclasses
                            .iter()
                            .map(|sc| {
                                [
                                    format!("{}Cpp", sc.subclass),
                                    sc.subclass.to_string(),
                                    sc.superclass.clone(),
                                ]
                            })
                            .flatten(),
                    ),
            )),
            Allowlist::Unspecified(_) => unreachable!(),
        }
    }

    fn active_utilities(&self) -> Vec<String> {
        if self.exclude_utilities {
            Vec::new()
        } else {
            vec![self.get_makestring_name()]
        }
    }

    /// Whether this type is on the allowlist specified by the user.
    ///
    /// A note on the allowlist handling in general. It's used in two places:
    /// 1) As directives to bindgen
    /// 2) After bindgen has generated code, to filter the APIs which
    ///    we pass to cxx.
    /// This second pass may seem redundant. But sometimes bindgen generates
    /// unnecessary stuff.
    pub fn is_on_allowlist(&self, cpp_name: &str) -> bool {
        match self.bindgen_allowlist() {
            None => true,
            Some(mut items) => {
                items.any(|item| item == cpp_name)
                    || self.active_utilities().iter().any(|item| *item == cpp_name)
            }
        }
    }

    pub fn is_on_blocklist(&self, cpp_name: &str) -> bool {
        self.blocklist.contains(&cpp_name.to_string())
    }

    pub fn get_blocklist(&self) -> impl Iterator<Item = &String> {
        self.blocklist.iter()
    }

    pub fn get_makestring_name(&self) -> String {
        format!(
            "autocxx_make_string_{}",
            self.mod_name
                .as_ref()
                .map(|i| i.to_string())
                .unwrap_or_else(|| "default".into())
        )
    }

    pub fn is_rust_type(&self, id: &Ident) -> bool {
        self.rust_types.contains(id)
    }

    pub fn superclasses(&self) -> impl Iterator<Item = &String> {
        let mut uniquified = HashSet::new();
        uniquified.extend(self.subclasses.iter().map(|sc| &sc.superclass));
        uniquified.into_iter()
    }

    pub fn is_subclass_holder(&self, id: &str) -> bool {
        self.subclasses
            .iter()
            .any(|sc| format!("{}Holder", sc.subclass.to_string()) == id)
    }

    /// Return the filename to which generated .rs should be written.
    pub fn get_rs_filename(&self) -> String {
        "autocxx-gen.rs".into()
    }
}

#[cfg(test)]
mod parse_tests {
    use crate::config::UnsafePolicy;
    use syn::parse_quote;
    #[test]
    fn test_safety_unsafe() {
        let us: UnsafePolicy = parse_quote! {
            unsafe
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_unsafe_ffi() {
        let us: UnsafePolicy = parse_quote! {
            unsafe_ffi
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_safe() {
        let us: UnsafePolicy = parse_quote! {};
        assert_eq!(us, UnsafePolicy::AllFunctionsUnsafe)
    }
}
<commit_msg>Supporting multiple mods.<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashSet;

use proc_macro2::Span;
use syn::{
    parse::{Parse, ParseStream},
    LitStr, Token,
};
use syn::{Ident, Result as ParseResult};

#[derive(PartialEq, Clone, Debug, Hash)]
pub enum UnsafePolicy {
    AllFunctionsSafe,
    AllFunctionsUnsafe,
}

impl Parse for UnsafePolicy {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        if input.parse::<Option<Token![unsafe]>>()?.is_some() {
            return Ok(UnsafePolicy::AllFunctionsSafe);
        }
        let r = match input.parse::<Option<syn::Ident>>()? {
            Some(id) => {
                if id == "unsafe_ffi" {
                    Ok(UnsafePolicy::AllFunctionsSafe)
                } else {
                    Err(syn::Error::new(id.span(), "expected unsafe_ffi"))
                }
            }
            None => Ok(UnsafePolicy::AllFunctionsUnsafe),
        };
        if !input.is_empty() {
            return Err(syn::Error::new(
                Span::call_site(),
                "unexpected tokens within safety directive",
            ));
        }
        r
    }
}

/// Allowlist configuration.
#[derive(Hash, Debug)]
pub enum Allowlist {
    Unspecified(Vec<String>),
    All,
    Specific(Vec<String>),
}

impl Allowlist {
    pub(crate) fn push(&mut self, item: LitStr) -> ParseResult<()> {
        match self {
            Allowlist::Unspecified(ref mut uncommitted_list) => {
                let new_list = uncommitted_list
                    .drain(..)
                    .chain(std::iter::once(item.value()))
                    .collect();
                *self = Allowlist::Specific(new_list);
            }
            Allowlist::All => {
                return Err(syn::Error::new(
                    item.span(),
                    "use either generate!/generate_pod! or generate_all!, not both.",
                ))
            }
            Allowlist::Specific(list) => list.push(item.value()),
        };
        Ok(())
    }

    pub(crate) fn set_all(&mut self, ident: &Ident) -> ParseResult<()> {
        if matches!(self, Allowlist::Specific(..)) {
            return Err(syn::Error::new(
                ident.span(),
                "use either generate!/generate_pod! or generate_all!, not both.",
            ));
        }
        *self = Allowlist::All;
        Ok(())
    }
}

impl Default for Allowlist {
    fn default() -> Self {
        Allowlist::Unspecified(Vec::new())
    }
}

#[derive(Hash, Debug)]
pub struct Subclass {
    pub superclass: String,
    pub subclass: Ident,
}

#[derive(Hash, Debug)]
pub struct IncludeCppConfig {
    pub inclusions: Vec<String>,
    pub unsafe_policy: UnsafePolicy,
    pub parse_only: bool,
    pub exclude_impls: bool,
    pod_requests: Vec<String>,
    allowlist: Allowlist,
    blocklist: Vec<String>,
    exclude_utilities: bool,
    mod_name: Option<Ident>,
    pub rust_types: Vec<Ident>,
    pub subclasses: Vec<Subclass>,
}

impl Parse for IncludeCppConfig {
    fn parse(input: ParseStream) -> ParseResult<Self> {
        // Takes as inputs:
        // 1. List of headers to include
        // 2. List of #defines to include
        // 3. Allowlist

        let mut inclusions = Vec::new();
        let mut parse_only = false;
        let mut exclude_impls = false;
        let mut unsafe_policy = UnsafePolicy::AllFunctionsUnsafe;
        let mut allowlist = Allowlist::default();
        let mut blocklist = Vec::new();
        let mut pod_requests = Vec::new();
        let mut rust_types = Vec::new();
        let mut exclude_utilities = false;
        let mut mod_name = None;
        let mut subclasses = Vec::new();

        while !input.is_empty() {
            let has_hexathorpe = input.parse::<Option<syn::token::Pound>>()?.is_some();
            let ident: syn::Ident = input.parse()?;
            if has_hexathorpe {
                if ident != "include" {
                    return Err(syn::Error::new(ident.span(), "expected include"));
                }
                let hdr: syn::LitStr = input.parse()?;
                inclusions.push(hdr.value());
            } else {
                input.parse::<Option<syn::token::Bang>>()?;
                if ident == "generate" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    allowlist.push(generate)?;
                } else if ident == "generate_pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate_pod: syn::LitStr = args.parse()?;
                    pod_requests.push(generate_pod.value());
                    allowlist.push(generate_pod)?;
                } else if ident == "pod" {
                    let args;
                    syn::parenthesized!(args in input);
                    let pod: syn::LitStr = args.parse()?;
                    pod_requests.push(pod.value());
                } else if ident == "block" {
                    let args;
                    syn::parenthesized!(args in input);
                    let generate: syn::LitStr = args.parse()?;
                    blocklist.push(generate.value());
                } else if ident == "rust_type" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    rust_types.push(ident);
                } else if ident == "subclass" {
                    let args;
                    syn::parenthesized!(args in input);
                    let superclass: syn::LitStr = args.parse()?;
                    if superclass.value().contains("::") {
                        return Err(syn::Error::new(
                            superclass.span(),
                            "at this time, superclasses can't be in namespaces - see https://github.com/google/autocxx/issues/599",
                        ));
                    }
                    args.parse::<syn::token::Comma>()?;
                    let subclass: syn::Ident = args.parse()?;
                    subclasses.push(Subclass {
                        superclass: superclass.value(),
                        subclass,
                    });
                } else if ident == "parse_only" {
                    parse_only = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "exclude_impls" {
                    exclude_impls = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "generate_all" {
                    allowlist.set_all(&ident)?;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "name" {
                    let args;
                    syn::parenthesized!(args in input);
                    let ident: syn::Ident = args.parse()?;
                    mod_name = Some(ident);
                } else if ident == "exclude_utilities" {
                    exclude_utilities = true;
                    swallow_parentheses(&input, &ident)?;
                } else if ident == "safety" {
                    let args;
                    syn::parenthesized!(args in input);
                    unsafe_policy = args.parse()?;
                } else {
                    return Err(syn::Error::new(
                        ident.span(),
                        "expected generate, generate_pod, nested_type, safety or exclude_utilities",
                    ));
                }
            }
            if input.is_empty() {
                break;
            }
        }

        Ok(IncludeCppConfig {
            inclusions,
            unsafe_policy,
            parse_only,
            exclude_impls,
            pod_requests,
            rust_types,
            allowlist,
            blocklist,
            exclude_utilities,
            mod_name,
            subclasses,
        })
    }
}

fn swallow_parentheses(input: &ParseStream, latest_ident: &Ident) -> ParseResult<()> {
    let args;
    syn::parenthesized!(args in input);
    if args.is_empty() {
        Ok(())
    } else {
        Err(syn::Error::new(
            latest_ident.span(),
            "expected no arguments to directive",
        ))
    }
}

impl IncludeCppConfig {
    pub fn get_pod_requests(&self) -> &[String] {
        &self.pod_requests
    }

    pub fn get_mod_name(&self) -> Ident {
        self.mod_name
            .as_ref()
            .cloned()
            .unwrap_or_else(|| Ident::new("ffi", Span::call_site()))
    }

    /// Whether to avoid generating the standard helpful utility
    /// functions which we normally include in every mod.
    pub fn exclude_utilities(&self) -> bool {
        self.exclude_utilities
    }

    /// Items which the user has explicitly asked us to generate;
    /// we should raise an error if we weren't able to do so.
    pub fn must_generate_list(&self) -> Box<dyn Iterator<Item = String> + '_> {
        if let Allowlist::Specific(items) = &self.allowlist {
            Box::new(items.iter().chain(self.pod_requests.iter()).cloned())
        } else {
            Box::new(self.pod_requests.iter().cloned())
        }
    }

    /// The allowlist of items to be passed into bindgen, if any.
    pub fn bindgen_allowlist(&self) -> Option<Box<dyn Iterator<Item = String> + '_>> {
        match &self.allowlist {
            Allowlist::All => None,
            Allowlist::Specific(items) => Some(Box::new(
                items
                    .iter()
                    .chain(self.pod_requests.iter())
                    .cloned()
                    .chain(self.active_utilities())
                    .chain(
                        self.subclasses
                            .iter()
                            .map(|sc| {
                                [
                                    format!("{}Cpp", sc.subclass),
                                    sc.subclass.to_string(),
                                    sc.superclass.clone(),
                                ]
                            })
                            .flatten(),
                    ),
            )),
            Allowlist::Unspecified(_) => unreachable!(),
        }
    }

    fn active_utilities(&self) -> Vec<String> {
        if self.exclude_utilities {
            Vec::new()
        } else {
            vec![self.get_makestring_name()]
        }
    }

    /// Whether this type is on the allowlist specified by the user.
    ///
    /// A note on the allowlist handling in general. It's used in two places:
    /// 1) As directives to bindgen
    /// 2) After bindgen has generated code, to filter the APIs which
    ///    we pass to cxx.
    /// This second pass may seem redundant. But sometimes bindgen generates
    /// unnecessary stuff.
    pub fn is_on_allowlist(&self, cpp_name: &str) -> bool {
        match self.bindgen_allowlist() {
            None => true,
            Some(mut items) => {
                items.any(|item| item == cpp_name)
                    || self.active_utilities().iter().any(|item| *item == cpp_name)
            }
        }
    }

    pub fn is_on_blocklist(&self, cpp_name: &str) -> bool {
        self.blocklist.contains(&cpp_name.to_string())
    }

    pub fn get_blocklist(&self) -> impl Iterator<Item = &String> {
        self.blocklist.iter()
    }

    pub fn get_makestring_name(&self) -> String {
        format!(
            "autocxx_make_string_{}",
            self.mod_name
                .as_ref()
                .map(|i| i.to_string())
                .unwrap_or_else(|| "default".into())
        )
    }

    pub fn is_rust_type(&self, id: &Ident) -> bool {
        self.rust_types.contains(id)
    }

    pub fn superclasses(&self) -> impl Iterator<Item = &String> {
        let mut uniquified = HashSet::new();
        uniquified.extend(self.subclasses.iter().map(|sc| &sc.superclass));
        uniquified.into_iter()
    }

    pub fn is_subclass_holder(&self, id: &str) -> bool {
        self.subclasses
            .iter()
            .any(|sc| format!("{}Holder", sc.subclass.to_string()) == id)
    }

    /// Return the filename to which generated .rs should be written.
    pub fn get_rs_filename(&self) -> String {
        format!(
            "autocxx-{}-gen.rs",
            self.mod_name
                .as_ref()
                .map(|id| id.to_string())
                .unwrap_or("ffi-default".into())
        )
    }
}

#[cfg(test)]
mod parse_tests {
    use crate::config::UnsafePolicy;
    use syn::parse_quote;
    #[test]
    fn test_safety_unsafe() {
        let us: UnsafePolicy = parse_quote! {
            unsafe
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_unsafe_ffi() {
        let us: UnsafePolicy = parse_quote! {
            unsafe_ffi
        };
        assert_eq!(us, UnsafePolicy::AllFunctionsSafe)
    }

    #[test]
    fn test_safety_safe() {
        let us: UnsafePolicy = parse_quote! {};
        assert_eq!(us, UnsafePolicy::AllFunctionsUnsafe)
    }
}
<|endoftext|>


<commit_before>//! Types and traits for working with the Matrix protocol.
//!
//! This crate re-exports things from all of the other ruma crates so you don't
//! have to manually keep all the versions in sync.
//!
//! Which crates are re-exported can be configured through cargo features.
//! Depending on which parts of Matrix are relevant to you, activate the
//! following features:
//!
//! * `client-api` for the client-server API
//! * `federation-api` for the server-server (federation) API
//! * `appservice-api` for the application service API

#![deny(missing_docs)]

pub use ruma_common::*;
#[doc(inline)]
pub use ruma_identifiers as identifiers;

pub use ruma_identifiers::{
    device_id, device_key_id, event_id, room_alias_id, room_id, room_version_id, server_key_id,
    server_name, user_id, DeviceId, DeviceKeyAlgorithm, DeviceKeyId, EventId, RoomAliasId, RoomId,
    RoomIdOrAliasId, RoomVersionId, ServerKeyAlgorithm, ServerKeyId, ServerName, UserId,
};

#[cfg(feature = "ruma-events")]
#[doc(inline)]
pub use ruma_events as events;
#[cfg(feature = "ruma-signatures")]
#[doc(inline)]
pub use ruma_signatures as signatures;

/// Rust types for various Matrix APIs requests and responses and abstractions for them.
#[cfg(feature = "ruma-api")]
#[doc(inline)]
pub mod api {
    // Everything except the `ruma_api!` macro
    pub use ruma_api::{error, Endpoint, EndpointError, Metadata, Outgoing};

    #[cfg(feature = "ruma-appservice-api")]
    #[doc(inline)]
    pub use ruma_appservice_api as appservice;
    #[cfg(feature = "ruma-client-api")]
    #[doc(inline)]
    pub use ruma_client_api as client;
    #[cfg(feature = "ruma-federation-api")]
    #[doc(inline)]
    pub use ruma_federation_api as federation;
    #[cfg(feature = "identity-service-api")]
    #[doc(inline)]
    pub use ruma_identity_service_api as identity_service;
    #[cfg(feature = "identity-service-api")]
    #[doc(inline)]
    pub use ruma_push_gateway_api as push_gateway;
}
<commit_msg>Re-export ruma_api's pub api in the ruma crate<commit_after>//! Types and traits for working with the Matrix protocol.
//!
//! This crate re-exports things from all of the other ruma crates so you don't
//! have to manually keep all the versions in sync.
//!
//! Which crates are re-exported can be configured through cargo features.
//! Depending on which parts of Matrix are relevant to you, activate the
//! following features:
//!
//! * `client-api` for the client-server API
//! * `federation-api` for the server-server (federation) API
//! * `appservice-api` for the application service API

#![deny(missing_docs)]

pub use ruma_common::*;
#[doc(inline)]
pub use ruma_identifiers as identifiers;

pub use ruma_identifiers::{
    device_id, device_key_id, event_id, room_alias_id, room_id, room_version_id, server_key_id,
    server_name, user_id, DeviceId, DeviceKeyAlgorithm, DeviceKeyId, EventId, RoomAliasId, RoomId,
    RoomIdOrAliasId, RoomVersionId, ServerKeyAlgorithm, ServerKeyId, ServerName, UserId,
};

#[cfg(feature = "ruma-events")]
#[doc(inline)]
pub use ruma_events as events;
#[cfg(feature = "ruma-signatures")]
#[doc(inline)]
pub use ruma_signatures as signatures;

/// Rust types for various Matrix APIs requests and responses and abstractions for them.
#[cfg(feature = "ruma-api")]
#[doc(inline)]
pub mod api {
    pub use ruma_api::*;

    #[cfg(feature = "ruma-appservice-api")]
    #[doc(inline)]
    pub use ruma_appservice_api as appservice;
    #[cfg(feature = "ruma-client-api")]
    #[doc(inline)]
    pub use ruma_client_api as client;
    #[cfg(feature = "ruma-federation-api")]
    #[doc(inline)]
    pub use ruma_federation_api as federation;
    #[cfg(feature = "identity-service-api")]
    #[doc(inline)]
    pub use ruma_identity_service_api as identity_service;
    #[cfg(feature = "identity-service-api")]
    #[doc(inline)]
    pub use ruma_push_gateway_api as push_gateway;
}
<|endoftext|>


<commit_before>#![no_std]
#![feature(alloc_error_handler)]

extern crate alloc;

use core::alloc::Layout;
use core::cell::Cell;
use core::convert::TryInto;
use core::fmt::{self, Debug, Error, Formatter};
use core::time::Duration;
use ctap2_authenticator::credentials::{RpCredential, UserCredential};
use ctap2_authenticator::usbhid::{
    CtapHidCapabilities, CtapHidPlatform, KeepaliveResponse, TransactionProcessor,
};
use ctap2_authenticator::Authenticator;
use ctap2_authenticator::{
    AuthenticatorPlatform, CredentialDescriptorList, CtapOptions, PublicKey, Signature,
};
use generic_array::GenericArray;
use libtock::ctap::{CtapRecvBuffer, CtapSendBuffer};
use libtock::hmac::{HmacDataBuffer, HmacDestBuffer, HmacDriverFactory, HmacKeyBuffer};
use libtock::println;
use libtock::result::TockResult;
use libtock::syscalls;
use p256::ecdsa::{signature::Signer, SigningKey};
use p256::elliptic_curve::ff::PrimeField;
use p256::{Scalar, SecretKey};
use subtle::{Choice, ConditionallySelectable};

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
pub struct PrivateKey(Scalar);

impl ConditionallySelectable for PrivateKey {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self(ConditionallySelectable::conditional_select(
            &a.0, &b.0, choice,
        ))
    }
}

impl Default for PrivateKey {
    fn default() -> Self {
        Self(Scalar::default())
    }
}

impl PrivateKey {
    pub fn from_bytes(bytes: &[u8; 32]) -> Option<Self> {
        Scalar::from_repr(GenericArray::clone_from_slice(bytes)).map(|s| Self(s))
    }
}

/// This is the provided implementation of `CtapHidPlatform` to be used in the `UsbContext`
pub(crate) struct UsbKeyHidPlatform {}

impl UsbKeyHidPlatform {
    pub fn new() -> Self {
        Self {}
    }
}

impl CtapHidPlatform for UsbKeyHidPlatform {
    // Should be kept in sync with the Crates version if possible
    const MAJOR_VERSION: u8 = 0;
    const MINOR_VERSION: u8 = 0;
    const BUILD_VERSION: u8 = 0;
    const CAPABILITIES: CtapHidCapabilities = CtapHidCapabilities {
        wink: true,
        cbor: true,
        msg: false,
    };

    fn wink(&mut self) {}

    fn cancel(&mut self) {
        println!("cancel");
        unimplemented!()
    }

    fn keepalive_needed(&mut self) -> KeepaliveResponse {
        println!("keepalive_needed");
        unimplemented!()
    }

    fn start_timer(&mut self) {
        println!("start_timer");
    }

    fn has_timed_out(&mut self) -> bool {
        println!("has_timed_out");
        false
    }
}

#[derive(Clone, Copy)]
pub struct HmacKeyCredential(pub(crate) [u8; 40]);

impl HmacKeyCredential {
    fn new(mac: &[u8; 32], nonce: [u8; 8]) -> Self {
        let mut new = [0; 40];
        new[..32].copy_from_slice(&mac[..]);
        new[32..].copy_from_slice(&nonce[..]);
        Self(new)
    }

    fn get_mac(&self) -> [u8; 32] {
        let mut mac = [0; 32];
        mac.copy_from_slice(&self.0[..32]);
        mac
    }
}

impl AsRef<[u8]> for HmacKeyCredential {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl Debug for HmacKeyCredential {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        write!(
            f,
            "HmacKeyCredential {{ mac: {:?}, nonce: {:?} }}",
            &self.0[..32],
            &self.0[32..]
        )
    }
}

impl PartialEq for HmacKeyCredential {
    fn eq(&self, other: &Self) -> bool {
        self.0[..] == other.0[..]
    }
}

impl Eq for HmacKeyCredential {}

const ITERATOR_MAX_LENGTH: usize = 8;

pub struct HmacCredentialQueue {
    index: usize,
    length: usize,
    list: [Option<(HmacKeyCredential, u32)>; ITERATOR_MAX_LENGTH],
}

impl HmacCredentialQueue {
    pub(crate) fn new() -> Self {
        Self {
            index: 0,
            length: 0,
            list: [None; ITERATOR_MAX_LENGTH],
        }
    }

    pub(crate) fn push(&mut self, credential: HmacKeyCredential, counter: u32) -> Result<(), ()> {
        if self.length < ITERATOR_MAX_LENGTH {
            self.list[self.length] = Some((credential, counter));
            self.length += 1;
            Ok(())
        } else {
            Err(())
        }
    }

    pub(crate) fn len(&self) -> usize {
        self.length
    }
}

impl Iterator for HmacCredentialQueue {
    type Item = (HmacKeyCredential, u32);

    fn next(&mut self) -> Option<Self::Item> {
        if self.list[self.index].is_some() {
            self.index += 1;
            self.list[self.index - 1]
        } else {
            None
        }
    }
}

pub(crate) struct CtapPlatform {
    hmac: HmacDriverFactory,
}

impl CtapPlatform {
    fn new(hmac: HmacDriverFactory) -> Self {
        Self { hmac }
    }
}

impl fmt::Debug for CtapPlatform {
    fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
        Ok(())
    }
}

impl CtapPlatform {
    fn generate_key_seed(
        &mut self,
        rp_id: &str,
        nonce: &[u8; 8],
    ) -> [u8; libtock::hmac::DEST_BUFFER_SIZE] {
        let hmac_driver = self.hmac.init_driver().unwrap();

        let mut key_buffer = HmacKeyBuffer::default();
        for (i, d) in rp_id.as_bytes().iter().enumerate() {
            key_buffer[i] = *d;
        }
        let _key_buffer = hmac_driver.init_key_buffer(&mut key_buffer).unwrap();

        let mut data_buffer = HmacDataBuffer::default();
        for (i, d) in nonce.iter().enumerate() {
            data_buffer[i] = *d;
        }
        let _data_buffer = hmac_driver.init_data_buffer(&mut data_buffer).unwrap();

        let mut seed_buffer = HmacDestBuffer::default();
        let seed_buffer = hmac_driver.init_dest_buffer(&mut seed_buffer).unwrap();

        let mut callback = |_result, _digest| {};

        let _subscription = hmac_driver.subscribe(&mut callback).unwrap();

        hmac_driver.run().unwrap();

        // Yield waiting for the HMAC callback
        unsafe {
            syscalls::raw::yieldk();
        }

        let mut temp_buffer = [0; libtock::hmac::DEST_BUFFER_SIZE];
        seed_buffer.read_bytes(&mut temp_buffer[..]);
        temp_buffer
    }
}

impl AuthenticatorPlatform for CtapPlatform {
    const AAGUID: [u8; 16] = [
        0xf8, 0xa0, 0x11, 0xf3, 0x8c, 0x0a, 0x4d, 0x15, 0x80, 0x06, 0x17, 0x11, 0x1f, 0x9e, 0xdc,
        0x7d,
    ];
    const CERTIFICATE: Option<&'static [u8]> = None;
    const MAX_MSG_LENGTH: u16 = 7609;
    const SUPPORTED_OPTIONS: CtapOptions = CtapOptions {
        plat: None,
        rk: Some(true),
        client_pin: None,
        up: Some(true),
        uv: Some(false),
    };
    type CredentialId = HmacKeyCredential;
    type PublicKeyBuffer = [u8; 32];
    type SignatureBuffer = [u8; 32];
    type CredentialIterator = HmacCredentialQueue;

    fn reset(&mut self) -> Result<(), ()> {
        unimplemented!()
    }

    fn check_exclude_list(&mut self, _rp_id: &str, _list: CredentialDescriptorList) -> bool {
        unimplemented!()
    }

    fn locate_credentials(
        &mut self,
        _rp_id: &str,
        list: Option<CredentialDescriptorList>,
    ) -> (u16, Self::CredentialIterator) {
        match list {
            None => (0, HmacCredentialQueue::new()),
            Some(_list) => {
                let result = HmacCredentialQueue::new();

                // TODO Add location

                (result.len() as u16, result)
            }
        }
    }

    fn create_credential(
        &mut self,
        rp: RpCredential<&[u8], &str>,
        _user: UserCredential<&[u8], &str>,
    ) -> (Self::CredentialId, PublicKey<Self::PublicKeyBuffer>, u32) {
        println!("create_credential");
        let rp_id = core::str::from_utf8(rp.rp_id()).unwrap();
        // This nonce is static!!!!
        // This is really bad cryptowise, let's print a warning
        // TODO: Convert this to generate a nonce from Tock's TRNG
        println!("WARNING!!! The nonce is static, this key is insecure");
        println!("WARNING!!! Do not use this anywhere important");
        let nonce = [0; 8];

        let seed_buffer = self.generate_key_seed(rp_id, &nonce);

        let credential = HmacKeyCredential::new(&seed_buffer, nonce.try_into().unwrap());

        let secret_key = SecretKey::from_bytes(&seed_buffer[0..32]).unwrap();

        let pub_key = p256::EncodedPoint::from_secret_key(&secret_key, false);

        let x: [u8; 32] = pub_key.as_bytes()[1..33].try_into().unwrap();
        let y: [u8; 32] = pub_key.as_bytes()[33..].try_into().unwrap();

        let ret_key = PublicKey::nistp256(x, y);

        (credential, ret_key, 0)
    }

    fn attest(
        &mut self,
        id: &Self::CredentialId,
        data: &[u8],
    ) -> Option<Signature<Self::SignatureBuffer>> {
        // TODO: Should attest be different then sign?
        let attest = {
            let secret_key = SecretKey::from_bytes(id.get_mac()).unwrap();

            let signer = SigningKey::from(&secret_key);

            let sig = signer.sign(data);

            let mut r: [u8; 32] = [0; 32];
            r.clone_from_slice(&sig.r().as_ref().to_bytes());
            let mut s: [u8; 32] = [0; 32];
            s.clone_from_slice(&sig.s().as_ref().to_bytes());

            Signature::nistp256(r, s)
        };

        Some(attest)
    }

    fn sign(
        &mut self,
        id: &Self::CredentialId,
        data: &[u8],
    ) -> Option<Signature<Self::SignatureBuffer>> {
        let attest = {
            let secret_key = SecretKey::from_bytes(id.get_mac()).unwrap();
            let signer = SigningKey::from(&secret_key);
            let sig = signer.sign(data);

            let mut r: [u8; 32] = [0; 32];
            r.clone_from_slice(&sig.r().as_ref().to_bytes());
            let mut s: [u8; 32] = [0; 32];
            s.clone_from_slice(&sig.s().as_ref().to_bytes());

            Signature::nistp256(r, s)
        };

        Some(attest)
    }

    fn start_timeout(&mut self) {
        unimplemented!()
    }

    fn has_timed_out(&mut self, _time: Duration) -> bool {
        unimplemented!()
    }
}

#[alloc_error_handler]
unsafe fn alloc_error_handler(_: Layout) -> ! {
    println!("alloc_error_handler called");
    loop {
        syscalls::raw::yieldk();
    }
}

#[libtock::main]
async fn main() -> TockResult<()> {
    let mut tp: Cell<TransactionProcessor<&'static mut [u8], UsbKeyHidPlatform>>;
    let mut temp_buffer = [0; libtock::ctap::RECV_BUFFER_SIZE];
    let mut drivers = libtock::retrieve_drivers()?;
    drivers.console.create_console();

    println!("Starting CTAP feature example");

    let ctap_driver = drivers.ctap.init_driver()?;

    let mut recv_buffer = CtapRecvBuffer::default();
    let recv_buffer = ctap_driver.init_recv_buffer(&mut recv_buffer)?;

    let mut send_buffer = CtapSendBuffer::default();
    let mut send_buffer = ctap_driver.init_send_buffer(&mut send_buffer)?;

    static mut BUFFER: &'static mut [u8] = &mut [0; 2048];
    unsafe {
        tp = Cell::new(TransactionProcessor::new(BUFFER, UsbKeyHidPlatform::new()));
    }

    let mut authenticator = Authenticator::create(CtapPlatform::new(drivers.hmac));

    let mut callback = |sent, _| {
        let mut looping = true;
        recv_buffer.read_bytes(&mut temp_buffer[..]);

        let temp_tp = tp.get_mut();
        let mut poll_msg = if sent == 0 { Some(&temp_buffer) } else { None };

        while looping {
            looping = false;
            let (maybe_request, maybe_data) = temp_tp.poll(poll_msg);

            if let Some(request) = maybe_request {
                let response = authenticator.process(request);
                temp_tp.response(response);
                looping = true;
            }

            if let Some(mut data) = maybe_data {
                send_buffer.write_bytes(&mut data[0..]);
                let _ = ctap_driver.send_data();
            }

            poll_msg = None;
        }
        let _ = ctap_driver.allow_receive();
    };

    let _subscription = ctap_driver.subscribe(&mut callback)?;
    ctap_driver.allow_receive()?;

    loop {
        unsafe { syscalls::raw::yieldk() };
    }
}
<commit_msg>examples-features/ctap: Add support for locate_credentials<commit_after>#![no_std]
#![feature(alloc_error_handler)]

extern crate alloc;

use core::alloc::Layout;
use core::cell::Cell;
use core::convert::TryInto;
use core::fmt::{self, Debug, Error, Formatter};
use core::time::Duration;
use ctap2_authenticator::credentials::{RpCredential, UserCredential};
use ctap2_authenticator::usbhid::{
    CtapHidCapabilities, CtapHidPlatform, KeepaliveResponse, TransactionProcessor,
};
use ctap2_authenticator::Authenticator;
use ctap2_authenticator::{
    AuthenticatorPlatform, CredentialDescriptorList, CtapOptions, PublicKey, Signature,
};
use generic_array::GenericArray;
use libtock::ctap::{CtapRecvBuffer, CtapSendBuffer};
use libtock::hmac::{HmacDataBuffer, HmacDestBuffer, HmacDriverFactory, HmacKeyBuffer};
use libtock::println;
use libtock::result::TockResult;
use libtock::syscalls;
use p256::ecdsa::{signature::Signer, SigningKey};
use p256::elliptic_curve::ff::PrimeField;
use p256::{Scalar, SecretKey};
use subtle::{Choice, ConditionallySelectable};

#[derive(Debug, Clone, PartialEq, Eq, Copy)]
pub struct PrivateKey(Scalar);

impl ConditionallySelectable for PrivateKey {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self(ConditionallySelectable::conditional_select(
            &a.0, &b.0, choice,
        ))
    }
}

impl Default for PrivateKey {
    fn default() -> Self {
        Self(Scalar::default())
    }
}

impl PrivateKey {
    pub fn from_bytes(bytes: &[u8; 32]) -> Option<Self> {
        Scalar::from_repr(GenericArray::clone_from_slice(bytes)).map(|s| Self(s))
    }
}

/// This is the provided implementation of `CtapHidPlatform` to be used in the `UsbContext`
pub(crate) struct UsbKeyHidPlatform {}

impl UsbKeyHidPlatform {
    pub fn new() -> Self {
        Self {}
    }
}

impl CtapHidPlatform for UsbKeyHidPlatform {
    // Should be kept in sync with the Crates version if possible
    const MAJOR_VERSION: u8 = 0;
    const MINOR_VERSION: u8 = 0;
    const BUILD_VERSION: u8 = 0;
    const CAPABILITIES: CtapHidCapabilities = CtapHidCapabilities {
        wink: true,
        cbor: true,
        msg: false,
    };

    fn wink(&mut self) {}

    fn cancel(&mut self) {
        println!("cancel");
        unimplemented!()
    }

    fn keepalive_needed(&mut self) -> KeepaliveResponse {
        println!("keepalive_needed");
        unimplemented!()
    }

    fn start_timer(&mut self) {
        println!("start_timer");
    }

    fn has_timed_out(&mut self) -> bool {
        println!("has_timed_out");
        false
    }
}

#[derive(Clone, Copy)]
pub struct HmacKeyCredential(pub(crate) [u8; 40]);

impl HmacKeyCredential {
    fn new(mac: &[u8; 32], nonce: [u8; 8]) -> Self {
        let mut new = [0; 40];
        new[..32].copy_from_slice(&mac[..]);
        new[32..].copy_from_slice(&nonce[..]);
        Self(new)
    }

    fn get_mac(&self) -> [u8; 32] {
        let mut mac = [0; 32];
        mac.copy_from_slice(&self.0[..32]);
        mac
    }
}

impl AsRef<[u8]> for HmacKeyCredential {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl Debug for HmacKeyCredential {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        write!(
            f,
            "HmacKeyCredential {{ mac: {:?}, nonce: {:?} }}",
            &self.0[..32],
            &self.0[32..]
        )
    }
}

impl PartialEq for HmacKeyCredential {
    fn eq(&self, other: &Self) -> bool {
        self.0[..] == other.0[..]
    }
}

impl Eq for HmacKeyCredential {}

const ITERATOR_MAX_LENGTH: usize = 8;

pub struct HmacCredentialQueue {
    index: usize,
    length: usize,
    list: [Option<(HmacKeyCredential, u32)>; ITERATOR_MAX_LENGTH],
}

impl HmacCredentialQueue {
    pub(crate) fn new() -> Self {
        Self {
            index: 0,
            length: 0,
            list: [None; ITERATOR_MAX_LENGTH],
        }
    }

    pub(crate) fn push(&mut self, credential: HmacKeyCredential, counter: u32) -> Result<(), ()> {
        if self.length < ITERATOR_MAX_LENGTH {
            self.list[self.length] = Some((credential, counter));
            self.length += 1;
            Ok(())
        } else {
            Err(())
        }
    }

    pub(crate) fn len(&self) -> usize {
        self.length
    }
}

impl Iterator for HmacCredentialQueue {
    type Item = (HmacKeyCredential, u32);

    fn next(&mut self) -> Option<Self::Item> {
        if self.list[self.index].is_some() {
            self.index += 1;
            self.list[self.index - 1]
        } else {
            None
        }
    }
}

pub(crate) struct CtapPlatform {
    hmac: HmacDriverFactory,
}

impl CtapPlatform {
    fn new(hmac: HmacDriverFactory) -> Self {
        Self { hmac }
    }
}

impl fmt::Debug for CtapPlatform {
    fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
        Ok(())
    }
}

impl CtapPlatform {
    fn generate_key_seed(
        &mut self,
        rp_id: &str,
        nonce: &[u8; 8],
    ) -> [u8; libtock::hmac::DEST_BUFFER_SIZE] {
        let hmac_driver = self.hmac.init_driver().unwrap();

        let mut key_buffer = HmacKeyBuffer::default();
        for (i, d) in rp_id.as_bytes().iter().enumerate() {
            key_buffer[i] = *d;
        }
        let _key_buffer = hmac_driver.init_key_buffer(&mut key_buffer).unwrap();

        let mut data_buffer = HmacDataBuffer::default();
        for (i, d) in nonce.iter().enumerate() {
            data_buffer[i] = *d;
        }
        let _data_buffer = hmac_driver.init_data_buffer(&mut data_buffer).unwrap();

        let mut seed_buffer = HmacDestBuffer::default();
        let seed_buffer = hmac_driver.init_dest_buffer(&mut seed_buffer).unwrap();

        let mut callback = |_result, _digest| {};

        let _subscription = hmac_driver.subscribe(&mut callback).unwrap();

        hmac_driver.run().unwrap();

        // Yield waiting for the HMAC callback
        unsafe {
            syscalls::raw::yieldk();
        }

        let mut temp_buffer = [0; libtock::hmac::DEST_BUFFER_SIZE];
        seed_buffer.read_bytes(&mut temp_buffer[..]);
        temp_buffer
    }

    /// Generates a Credential from a CredentialId and checks that the Id is valid
    fn checked_generate_hmac_cred(
        &mut self,
        rp_id: &str,
        credential_id: &[u8],
    ) -> Option<HmacKeyCredential> {
        if credential_id.len() != 40 {
            return None;
        }

        let received_mac = &credential_id[..32];

        let seed_buffer = self.generate_key_seed(rp_id, &credential_id[32..40].try_into().unwrap());

        let credential =
            HmacKeyCredential::new(&seed_buffer, credential_id[32..].try_into().unwrap());

        let generated_mac = &credential.0[..32];

        if generated_mac == received_mac {
            Some(credential)
        } else {
            None
        }
    }
}

impl AuthenticatorPlatform for CtapPlatform {
    const AAGUID: [u8; 16] = [
        0xf8, 0xa0, 0x11, 0xf3, 0x8c, 0x0a, 0x4d, 0x15, 0x80, 0x06, 0x17, 0x11, 0x1f, 0x9e, 0xdc,
        0x7d,
    ];
    const CERTIFICATE: Option<&'static [u8]> = None;
    const MAX_MSG_LENGTH: u16 = 7609;
    const SUPPORTED_OPTIONS: CtapOptions = CtapOptions {
        plat: None,
        rk: Some(true),
        client_pin: None,
        up: Some(true),
        uv: Some(false),
    };
    type CredentialId = HmacKeyCredential;
    type PublicKeyBuffer = [u8; 32];
    type SignatureBuffer = [u8; 32];
    type CredentialIterator = HmacCredentialQueue;

    fn reset(&mut self) -> Result<(), ()> {
        unimplemented!()
    }

    fn check_exclude_list(&mut self, _rp_id: &str, _list: CredentialDescriptorList) -> bool {
        unimplemented!()
    }

    fn locate_credentials(
        &mut self,
        rp_id: &str,
        list: Option<CredentialDescriptorList>,
    ) -> (u16, Self::CredentialIterator) {
        match list {
            None => (0, HmacCredentialQueue::new()),
            Some(list) => {
                let mut result = HmacCredentialQueue::new();

                for descriptor in list {
                    match self.checked_generate_hmac_cred(rp_id, descriptor.get_id()) {
                        None => (),
                        Some(credential) => {
                            match result.push(credential, 0) {
                                Ok(()) => (),
                                // NOTE: We just truncate if we don't have enough space to store all
                                // fitting credentials
                                Err(()) => break,
                            }
                        }
                    }
                }

                (result.len() as u16, result)
            }
        }
    }

    fn create_credential(
        &mut self,
        rp: RpCredential<&[u8], &str>,
        _user: UserCredential<&[u8], &str>,
    ) -> (Self::CredentialId, PublicKey<Self::PublicKeyBuffer>, u32) {
        println!("create_credential");
        let rp_id = core::str::from_utf8(rp.rp_id()).unwrap();
        // This nonce is static!!!!
        // This is really bad cryptowise, let's print a warning
        // TODO: Convert this to generate a nonce from Tock's TRNG
        println!("WARNING!!! The nonce is static, this key is insecure");
        println!("WARNING!!! Do not use this anywhere important");
        let nonce = [0; 8];

        let seed_buffer = self.generate_key_seed(rp_id, &nonce);

        let credential = HmacKeyCredential::new(&seed_buffer, nonce.try_into().unwrap());

        let secret_key = SecretKey::from_bytes(&seed_buffer[0..32]).unwrap();

        let pub_key = p256::EncodedPoint::from_secret_key(&secret_key, false);

        let x: [u8; 32] = pub_key.as_bytes()[1..33].try_into().unwrap();
        let y: [u8; 32] = pub_key.as_bytes()[33..].try_into().unwrap();

        let ret_key = PublicKey::nistp256(x, y);

        (credential, ret_key, 0)
    }

    fn attest(
        &mut self,
        id: &Self::CredentialId,
        data: &[u8],
    ) -> Option<Signature<Self::SignatureBuffer>> {
        // TODO: Should attest be different then sign?
        let attest = {
            let secret_key = SecretKey::from_bytes(id.get_mac()).unwrap();

            let signer = SigningKey::from(&secret_key);

            let sig = signer.sign(data);

            let mut r: [u8; 32] = [0; 32];
            r.clone_from_slice(&sig.r().as_ref().to_bytes());
            let mut s: [u8; 32] = [0; 32];
            s.clone_from_slice(&sig.s().as_ref().to_bytes());

            Signature::nistp256(r, s)
        };

        Some(attest)
    }

    fn sign(
        &mut self,
        id: &Self::CredentialId,
        data: &[u8],
    ) -> Option<Signature<Self::SignatureBuffer>> {
        let attest = {
            let secret_key = SecretKey::from_bytes(id.get_mac()).unwrap();
            let signer = SigningKey::from(&secret_key);
            let sig = signer.sign(data);

            let mut r: [u8; 32] = [0; 32];
            r.clone_from_slice(&sig.r().as_ref().to_bytes());
            let mut s: [u8; 32] = [0; 32];
            s.clone_from_slice(&sig.s().as_ref().to_bytes());

            Signature::nistp256(r, s)
        };

        Some(attest)
    }

    fn start_timeout(&mut self) {
        unimplemented!()
    }

    fn has_timed_out(&mut self, _time: Duration) -> bool {
        unimplemented!()
    }
}

#[alloc_error_handler]
unsafe fn alloc_error_handler(_: Layout) -> ! {
    println!("alloc_error_handler called");
    loop {
        syscalls::raw::yieldk();
    }
}

#[libtock::main]
async fn main() -> TockResult<()> {
    let mut tp: Cell<TransactionProcessor<&'static mut [u8], UsbKeyHidPlatform>>;
    let mut temp_buffer = [0; libtock::ctap::RECV_BUFFER_SIZE];
    let mut drivers = libtock::retrieve_drivers()?;
    drivers.console.create_console();

    println!("Starting CTAP feature example");

    let ctap_driver = drivers.ctap.init_driver()?;

    let mut recv_buffer = CtapRecvBuffer::default();
    let recv_buffer = ctap_driver.init_recv_buffer(&mut recv_buffer)?;

    let mut send_buffer = CtapSendBuffer::default();
    let mut send_buffer = ctap_driver.init_send_buffer(&mut send_buffer)?;

    static mut BUFFER: &'static mut [u8] = &mut [0; 2048];
    unsafe {
        tp = Cell::new(TransactionProcessor::new(BUFFER, UsbKeyHidPlatform::new()));
    }

    let mut authenticator = Authenticator::create(CtapPlatform::new(drivers.hmac));

    let mut callback = |sent, _| {
        let mut looping = true;
        recv_buffer.read_bytes(&mut temp_buffer[..]);

        let temp_tp = tp.get_mut();
        let mut poll_msg = if sent == 0 { Some(&temp_buffer) } else { None };

        while looping {
            looping = false;
            let (maybe_request, maybe_data) = temp_tp.poll(poll_msg);

            if let Some(request) = maybe_request {
                let response = authenticator.process(request);
                temp_tp.response(response);
                looping = true;
            }

            if let Some(mut data) = maybe_data {
                send_buffer.write_bytes(&mut data[0..]);
                let _ = ctap_driver.send_data();
            }

            poll_msg = None;
        }
        let _ = ctap_driver.allow_receive();
    };

    let _subscription = ctap_driver.subscribe(&mut callback)?;
    ctap_driver.allow_receive()?;

    loop {
        unsafe { syscalls::raw::yieldk() };
    }
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::io::ErrorKind::NotFound;
use std::fs;

use hyper::method::Method::{Get, Head};
use hyper::header::Connection;

use status::StatusCode;
use request::Request;
use response::Response;
use middleware::{Middleware, MiddlewareResult};

// this should be much simpler after unboxed closures land in Rust.

#[derive(Clone)]
pub struct StaticFilesHandler {
    root_path: PathBuf
}

impl<D> Middleware<D> for StaticFilesHandler {
    fn invoke<'a>(&self, req: &mut Request<D>, res: Response<'a, D>)
            -> MiddlewareResult<'a, D> {
        match req.origin.method {
            Get | Head => self.with_file(self.extract_path(req), res),
            _ => res.next_middleware()
        }
    }
}

impl StaticFilesHandler {
    /// Create a new middleware to serve files from within a given root directory.
    /// The file to serve will be determined by combining the requested Url with
    /// the provided root directory.
    ///
    ///
    /// # Examples
    /// ```{rust}
    /// use nickel::{Nickel, StaticFilesHandler};
    /// let mut server = Nickel::new();
    ///
    /// server.utilize(StaticFilesHandler::new("/path/to/serve/"));
    /// ```
    pub fn new<P: AsRef<Path>>(root_path: P) -> StaticFilesHandler {
        StaticFilesHandler {
            root_path: root_path.as_ref().to_path_buf()
        }
    }

    fn extract_path<'a, D>(&self, req: &'a mut Request<D>) -> Option<&'a str> {
        req.path_without_query().map(|path| {
            debug!("{:?} {:?}{:?}", req.origin.method, self.root_path.display(), path);

            match path {
                "/" => "index.html",
                path => &path[1..],
            }
        })
    }

    fn with_file<'a, 'b, D, P>(&self,
                            relative_path: Option<P>,
                            mut res: Response<'a, D>)
            -> MiddlewareResult<'a, D> where P: AsRef<Path> {
        if let Some(path) = relative_path {
            let path = path.as_ref();
            if !safe_path(path) {
                let log_msg = format!("The path '{:?}' was denied access.", path);
                return res.error(StatusCode::BadRequest, log_msg);
            }

            let path = self.root_path.join(path);
            match fs::metadata(&path) {
                Ok(ref attr) if attr.is_file() => {
                    res.set(Connection::close());
                    return res.send_file(&path);
                },
                Err(ref e) if e.kind() != NotFound => debug!("Error getting metadata \
                                                              for file '{:?}': {:?}",
                                                              path, e),
                _ => {}
            }
        };

        res.next_middleware()
    }
}

/// Block paths from accessing the parent directory
fn safe_path<P: AsRef<Path>>(path: P) -> bool {
    use std::path::Component;

    path.as_ref().components().all(|c| match c {
        // whitelist non-suspicious in case new things get added in future
        Component::CurDir | Component::Normal(_) => true,
        _ => false
    })
}

#[test]
fn bad_paths() {
    let bad_paths = &[
        "foo/bar/../baz/index.html",
        "foo/bar/../baz",
        "../bar/",
        "..",
        "/" // Root path should be handled already
    ];

    for &path in bad_paths {
        assert!(!safe_path(path), "expected {:?} to be suspicious", path);
    }
}

#[test]
fn valid_paths() {
    let good_paths = &[
        "foo/bar/./baz/index.html",
        "foo/bar/./baz",
        "./bar/",
        ".",
        "index.html"
    ];

    for &path in good_paths {
        assert!(safe_path(path), "expected {:?} to not be suspicious", path);
    }
}
<commit_msg>[Improved] nickel.rs StaticFilesHandler honor file mime types based on file extension.<commit_after>use std::path::{Path, PathBuf};
use std::io::ErrorKind::NotFound;
use std::fs;

use hyper::method::Method::{Get, Head};
use hyper::header::Connection;

use status::StatusCode;
use request::Request;
use response::Response;
use middleware::{Middleware, MiddlewareResult};

// this should be much simpler after unboxed closures land in Rust.

#[derive(Clone)]
pub struct StaticFilesHandler {
    root_path: PathBuf
}

impl<D> Middleware<D> for StaticFilesHandler {
    fn invoke<'a>(&self, req: &mut Request<D>, res: Response<'a, D>)
            -> MiddlewareResult<'a, D> {
        match req.origin.method {
            Get | Head => self.with_file(self.extract_path(req), res),
            _ => res.next_middleware()
        }
    }
}

impl StaticFilesHandler {
    /// Create a new middleware to serve files from within a given root directory.
    /// The file to serve will be determined by combining the requested Url with
    /// the provided root directory.
    ///
    ///
    /// # Examples
    /// ```{rust}
    /// use nickel::{Nickel, StaticFilesHandler};
    /// let mut server = Nickel::new();
    ///
    /// server.utilize(StaticFilesHandler::new("/path/to/serve/"));
    /// ```
    pub fn new<P: AsRef<Path>>(root_path: P) -> StaticFilesHandler {
        StaticFilesHandler {
            root_path: root_path.as_ref().to_path_buf()
        }
    }

    fn extract_path<'a, D>(&self, req: &'a mut Request<D>) -> Option<&'a str> {
        req.path_without_query().map(|path| {
            debug!("{:?} {:?}{:?}", req.origin.method, self.root_path.display(), path);

            match path {
                "/" => "index.html",
                path => &path[1..],
            }
        })
    }

    fn get_mime_by_filename<P: AsRef<Path>>(&self, path : P) -> &'static str {
        let _path = path.as_ref().to_str().unwrap();
        if _path.ends_with(".jpg") || _path.ends_with(".jpeg"){
            "image/jpeg"
        }else if _path.ends_with(".png"){
            "image/png"
        }else if _path.ends_with(".gif"){
            "image/gif"
        }else if _path.ends_with(".pdf"){
            "application/pdf"
        }else if _path.ends_with(".rtf"){
            "application/rtf"
        }else if _path.ends_with(".json"){
            "application/json"
        }else if _path.ends_with(".zip"){
            "application/x-zip-compressed"
        }else if _path.ends_with(".rar"){
            "application/x-compressed"
        }else if _path.ends_with(".gz") || _path.ends_with(".gzip"){
            "application/x-gzip"
        }else if _path.ends_with(".doc") || _path.ends_with(".docx"){
            "application/msword"
        }else if _path.ends_with(".xls") || _path.ends_with(".xlsx") || _path.ends_with(".xlt"){
            "application/excel"
        }else if _path.ends_with(".xml"){
            "text/xml"
        }else if _path.ends_with(".mov"){
            "video/quicktime"
        }else if _path.ends_with(".mp3"){
            "audio/mp3"
        }else{
            "application/octet-stream"
        }
    }

    fn with_file<'a, 'b, D, P>(&self,
                            relative_path: Option<P>,
                            mut res: Response<'a, D>)
            -> MiddlewareResult<'a, D> where P: AsRef<Path> {

        if let Some(path) = relative_path {
            let path = path.as_ref();
            if !safe_path(path) {
                let log_msg = format!("The path '{:?}' was denied access.", path);
                return res.error(StatusCode::BadRequest, log_msg);
            }

            let path = self.root_path.join(path);
            match fs::metadata(&path) {
                Ok(ref attr) if attr.is_file() => {

                    // res.set(MediaType::new("image".to_string(), "jpeg".to_string()));

                    let mime = self.get_mime_by_filename(&path);

                    res.headers_mut().set_raw("Content-Type", vec![mime.as_bytes().to_vec()]);

                    res.set(Connection::close());
                    return res.send_file(&path);
                },
                Err(ref e) if e.kind() != NotFound => debug!("Error getting metadata \
                                                              for file '{:?}': {:?}",
                                                              path, e),
                _ => {}
            }
        };

        res.next_middleware()
    }
}

/// Block paths from accessing the parent directory
fn safe_path<P: AsRef<Path>>(path: P) -> bool {
    use std::path::Component;

    path.as_ref().components().all(|c| match c {
        // whitelist non-suspicious in case new things get added in future
        Component::CurDir | Component::Normal(_) => true,
        _ => false
    })
}

#[test]
fn bad_paths() {
    let bad_paths = &[
        "foo/bar/../baz/index.html",
        "foo/bar/../baz",
        "../bar/",
        "..",
        "/" // Root path should be handled already
    ];

    for &path in bad_paths {
        assert!(!safe_path(path), "expected {:?} to be suspicious", path);
    }
}

#[test]
fn valid_paths() {
    let good_paths = &[
        "foo/bar/./baz/index.html",
        "foo/bar/./baz",
        "./bar/",
        ".",
        "index.html"
    ];

    for &path in good_paths {
        assert!(safe_path(path), "expected {:?} to not be suspicious", path);
    }
}
<|endoftext|>


<commit_before>//! Look-up table (LUT) implementation and transformation functions.
//!
//! This module contains the [`Lut`] data type,
//! designed to turn pixel data native sample values
//! (as encoded in the _Pixel Data_ attribute)
//! into displayable or otherwise more meaningful values.
//!
//! The type also provides easy-to-use constructor functions
//! for common DICOM sample value transformations.

use num_traits::{NumCast, ToPrimitive};
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use snafu::{OptionExt, Snafu};

use crate::{Rescale, WindowLevelTransform};

/// The LUT could not be created:
/// entry #{index} was mapped to {y_value},
/// which could not be cast to the target type.
#[derive(Debug, PartialEq, Snafu)]
pub struct CreateLutError {
    index: usize,
    y_value: f64,
}

impl CreateLutError {
    /// Get the original index in the LUT.
    pub fn index(&self) -> usize {
        self.index
    }
    /// Get the value which could not be converted to the target type.
    pub fn y_value(&self) -> f64 {
        self.y_value
    }
}

/// A look up table for pixel data sample value transformations.
///
/// # Example
///
/// The LUT can be populated with common transformations
/// via the functions [`new_rescale`](Lut::new_rescale),
/// [`new_window`](Lut::new_window),
/// and [`new_rescale_and_window`](Lut::new_rescale_and_window).
/// The function [`new_with_fn`](Lut::new_with_fn)
/// can be used to create a LUT with a custom function.
///
/// ```
/// # use dicom_pixeldata::{
/// #     CreateLutError, Lut, Rescale, VoiLutFunction,
/// #     WindowLevel, WindowLevelTransform,
/// # };
/// let bits_stored = 8;
/// let lut = Lut::new_rescale_and_window(
///     bits_stored,
///     false,
///     Rescale::new(1., -1024.),
///     WindowLevelTransform::new(
///         VoiLutFunction::Linear,
///         WindowLevel {
///             width: 300.,
///             center: 50.
///         }
///     ),
/// )?;
///
/// let val: u8 = lut.get(100_u8);
/// # Result::<(), CreateLutError>::Ok(())
/// ```
#[derive(Debug)]
pub struct Lut<T> {
    /// the table which maps an index to a transformed value,
    /// of size 2 to the power of `bits_stored`
    table: Vec<T>,
    /// whether the input sample values are signed (Pixel Representation = 1)
    signed: bool,
}

impl<T: 'static> Lut<T>
where
    T: NumCast,
    T: Copy,
    T: Send + Sync,
{
    /// Create a new LUT with the given characteristics
    /// and populates it with the outputs of the provided function.
    /// The function may be called concurrently.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `f`: the mapping function
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_with_fn(
        bits_stored: u16,
        signed: bool,
        f: impl Fn(f64) -> f64 + Sync,
    ) -> Result<Self, CreateLutError> {
        assert!(bits_stored != 0 && bits_stored <= 32);
        let size = (1 << bits_stored as u32) as usize;
        debug_assert!(size.is_power_of_two());

        let table: Result<Vec<_>, _> = (0..size)
            .into_par_iter()
            .map(|i| {
                // account for signedness to determine input pixel value
                let x = if signed && i >= size / 2 {
                    i as f64 - size as f64
                } else {
                    i as f64
                };
                let value = f(x);
                T::from(value).context(CreateLutSnafu {
                    index: i,
                    y_value: value,
                })
            })
            .collect();
        Ok(Self {
            table: table?,
            signed,
        })
    }

    /// Create a new LUT containing only the modality rescale transformation.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `rescale`: the rescale parameters
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_rescale(
        bits_stored: u16,
        signed: bool,
        rescale: Rescale,
    ) -> Result<Self, CreateLutError> {
        Self::new_with_fn(bits_stored, signed, |v| rescale.apply(v))
    }

    /// Create a new LUT containing the given modality rescale transformation
    /// and a min-max normalization
    /// which satisfies the raw samples given.
    /// The sample type `I` is expected to be either `u8` or `u16`,
    /// even if the sample is meant to be interpreted as signed.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `rescale`: the rescale parameters
    /// - `samples`: the raw pixel data samples expected to be fed to the LUT
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub(crate) fn new_rescale_and_normalize<I>(
        bits_stored: u16,
        signed: bool,
        rescale: Rescale,
        samples: I,
    ) -> Result<Self, CreateLutError>
    where
        I: IntoIterator,
        I: Clone,
        I::IntoIter: Clone,
        I::Item: Clone,
        I::Item: NumCast,
        I::Item: ToPrimitive,
    {
        let samples_f64 = samples.into_iter().filter_map(|v| v.to_f64());
        let min: f64 = samples_f64.clone().fold(f64::MAX, |a, b| a.min(b));
        let max: f64 = samples_f64.fold(f64::MIN, |a, b| a.max(b));

        // convert to f64 and swap if negative slope
        let (min, max) = if rescale.slope < 0. {
            (max, min)
        } else {
            (min, max)
        };

        // create a linear window level transform
        let voi = WindowLevelTransform::linear(crate::WindowLevel {
            width: max - min,
            center: min + max / 2.,
        });

        Self::new_rescale_and_window(bits_stored, signed, rescale, voi)
    }

    /// Create a new LUT containing the modality rescale transformation
    /// and the VOI transformation defined by a window level.
    ///
    /// The amplitude of the output values
    /// goes from 0 to `2^n - 1`, where `n` is the power of two
    /// which follows `bits_stored` (or itself if it is a power of two).
    /// For instance, if `bits_stored` is 12, the output values
    /// will go from 0 to 65535.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `rescale`: the rescale parameters
    /// - `voi`: the value of interest (VOI) function and parameters
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_rescale_and_window(
        bits_stored: u16,
        signed: bool,
        rescale: Rescale,
        voi: WindowLevelTransform,
    ) -> Result<Self, CreateLutError> {
        let bits_allocated = (bits_stored as usize).next_power_of_two();
        let y_max = ((1 << bits_allocated) - 1) as f64;
        Self::new_with_fn(bits_stored, signed, |v| {
            let x = v as f64;
            let v = rescale.apply(x);
            voi.apply(v, y_max)
        })
    }

    /// Create a new LUT containing
    /// a VOI transformation defined by a window level.
    ///
    /// The amplitude of the output values
    /// goes from 0 to `2^n - 1`, where `n` is the power of two
    /// which follows `bits_stored` (or itself if it is a power of two).
    /// For instance, if `bits_stored` is 12, the output values
    /// will go from 0 to 65535.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `voi`: the value of interest (VOI) function and parameters
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_window(
        bits_stored: u16,
        signed: bool,
        voi: WindowLevelTransform,
    ) -> Result<Self, CreateLutError> {
        let bits_allocated = (bits_stored as usize).next_power_of_two();
        let y_max = ((1 << bits_allocated) - 1) as f64;
        Self::new_with_fn(bits_stored, signed, |v| voi.apply(v as f64, y_max))
    }

    /// Apply the transformation to a single pixel sample value.
    ///
    /// Although the input is expected to be one of `u8`, `u16`, or `u32`,
    /// this method works for signed sample values as well,
    /// with the bits reinterpreted as their unsigned counterpart.
    ///
    /// # Panics
    ///
    /// Panics if `sample_value` is larger or equal to `2^bits_stored`.
    pub fn get<I: 'static>(&self, sample_value: I) -> T
    where
        I: Copy,
        I: Into<u32>,
    {
        let val = sample_value.into();
        let index = if self.signed {
            // adjust for signedness by masking out the extra sign bits
            let mask = self.table.len() - 1;
            val as usize & mask
        } else {
            val as usize
        };
        assert!((index as usize) < self.table.len());

        self.table[index as usize]
    }

    /// Adapts an iterator of pixel data sample values
    /// to an iterator of transformed values.
    pub fn map_iter<'a, I: 'static>(
        &'a self,
        iter: impl IntoIterator<Item = I> + 'a,
    ) -> impl Iterator<Item = T> + 'a
    where
        I: Copy,
        I: Into<u32>,
    {
        iter.into_iter().map(move |i| self.get(i))
    }

    /// Adapts a parallel iterator of pixel data sample values
    /// to a parallel iterator of transformed values.
    pub fn map_par_iter<'a, I: 'static>(
        &'a self,
        iter: impl ParallelIterator<Item = I> + 'a,
    ) -> impl ParallelIterator<Item = T> + 'a
    where
        I: Copy,
        I: Into<u32>,
    {
        iter.map(move |i| self.get(i))
    }
}

#[cfg(test)]
mod tests {
    use crate::{VoiLutFunction, WindowLevel};

    use super::*;

    /// Applying a common rescale function to a value
    /// gives the expected output.
    #[test]
    fn modality_lut_baseline_2() {
        let rescale = Rescale::new(2., -1024.);

        assert_eq!(rescale.apply(0.), -1024.);
        assert_eq!(rescale.apply(1.), -1022.);
        assert_eq!(rescale.apply(2.), -1020.);
        assert_eq!(rescale.apply(1024.), 1024.);
    }

    #[test]
    fn lut_signed_numbers() {
        // 10-bit precision input, signed output
        let lut: Lut<i16> = Lut::new_rescale(10, true, Rescale::new(2., -1024.)).unwrap();

        assert_eq!(lut.get(0 as u16), -1024);
        assert_eq!(lut.get(1 as u16), -1022);
        assert_eq!(lut.get(-1_i16 as u16), -1026);
        assert_eq!(lut.get(-2_i16 as u16), -1028);
        assert_eq!(lut.get(500 as u16), -24);
    }

    #[test]
    fn lut_rescale_and_window_16bit() {
        let bits_stored = 16;
        let lut = Lut::new_rescale_and_window(
            bits_stored,
            false,
            Rescale::new(1., -1024.),
            WindowLevelTransform::new(
                VoiLutFunction::Linear,
                WindowLevel {
                    width: 300.,
                    center: 50.,
                },
            ),
        )
        .unwrap();

        // < 0
        let val: u16 = lut.get(824_u16);
        assert_eq!(val, 0);

        // > 200
        let val: u16 = lut.get(1224_u16);
        assert_eq!(val, 65535);

        // around the middle

        let val: u16 = lut.get(1024_u16 + 50);
        let expected_range = 32_500..=33_000;
        assert!(
            expected_range.contains(&val),
            "outcome was {}, expected to be in {:?}",
            val,
            expected_range,
        );
    }
}
<commit_msg>[pixeldata] fix Lut::new_rescale_and_normalize<commit_after>//! Look-up table (LUT) implementation and transformation functions.
//!
//! This module contains the [`Lut`] data type,
//! designed to turn pixel data native sample values
//! (as encoded in the _Pixel Data_ attribute)
//! into displayable or otherwise more meaningful values.
//!
//! The type also provides easy-to-use constructor functions
//! for common DICOM sample value transformations.

use num_traits::{NumCast, ToPrimitive};
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use snafu::{OptionExt, Snafu};

use crate::{Rescale, WindowLevelTransform};

/// The LUT could not be created:
/// entry #{index} was mapped to {y_value},
/// which could not be cast to the target type.
#[derive(Debug, PartialEq, Snafu)]
pub struct CreateLutError {
    index: usize,
    y_value: f64,
}

impl CreateLutError {
    /// Get the original index in the LUT.
    pub fn index(&self) -> usize {
        self.index
    }
    /// Get the value which could not be converted to the target type.
    pub fn y_value(&self) -> f64 {
        self.y_value
    }
}

/// A look up table for pixel data sample value transformations.
///
/// # Example
///
/// The LUT can be populated with common transformations
/// via the functions [`new_rescale`](Lut::new_rescale),
/// [`new_window`](Lut::new_window),
/// and [`new_rescale_and_window`](Lut::new_rescale_and_window).
/// The function [`new_with_fn`](Lut::new_with_fn)
/// can be used to create a LUT with a custom function.
///
/// ```
/// # use dicom_pixeldata::{
/// #     CreateLutError, Lut, Rescale, VoiLutFunction,
/// #     WindowLevel, WindowLevelTransform,
/// # };
/// let bits_stored = 8;
/// let lut = Lut::new_rescale_and_window(
///     bits_stored,
///     false,
///     Rescale::new(1., -1024.),
///     WindowLevelTransform::new(
///         VoiLutFunction::Linear,
///         WindowLevel {
///             width: 300.,
///             center: 50.
///         }
///     ),
/// )?;
///
/// let val: u8 = lut.get(100_u8);
/// # Result::<(), CreateLutError>::Ok(())
/// ```
#[derive(Debug)]
pub struct Lut<T> {
    /// the table which maps an index to a transformed value,
    /// of size 2 to the power of `bits_stored`
    table: Vec<T>,
    /// whether the input sample values are signed (Pixel Representation = 1)
    signed: bool,
}

impl<T: 'static> Lut<T>
where
    T: NumCast,
    T: Copy,
    T: Send + Sync,
{
    /// Create a new LUT with the given characteristics
    /// and populates it with the outputs of the provided function.
    /// The function may be called concurrently.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `f`: the mapping function
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_with_fn(
        bits_stored: u16,
        signed: bool,
        f: impl Fn(f64) -> f64 + Sync,
    ) -> Result<Self, CreateLutError> {
        assert!(bits_stored != 0 && bits_stored <= 32);
        let size = (1 << bits_stored as u32) as usize;
        debug_assert!(size.is_power_of_two());

        let table: Result<Vec<_>, _> = (0..size)
            .into_par_iter()
            .map(|i| {
                // account for signedness to determine input pixel value
                let x = if signed && i >= size / 2 {
                    i as f64 - size as f64
                } else {
                    i as f64
                };
                let value = f(x);
                T::from(value).context(CreateLutSnafu {
                    index: i,
                    y_value: value,
                })
            })
            .collect();
        Ok(Self {
            table: table?,
            signed,
        })
    }

    /// Create a new LUT containing only the modality rescale transformation.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `rescale`: the rescale parameters
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_rescale(
        bits_stored: u16,
        signed: bool,
        rescale: Rescale,
    ) -> Result<Self, CreateLutError> {
        Self::new_with_fn(bits_stored, signed, |v| rescale.apply(v))
    }

    /// Create a new LUT containing the given modality rescale transformation
    /// and a min-max normalization
    /// which satisfies the raw samples given.
    /// The sample type `I` is expected to be either `u8` or `u16`,
    /// even if the sample is meant to be interpreted as signed.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `rescale`: the rescale parameters
    /// - `samples`: the raw pixel data samples expected to be fed to the LUT
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub(crate) fn new_rescale_and_normalize<I>(
        bits_stored: u16,
        signed: bool,
        rescale: Rescale,
        samples: I,
    ) -> Result<Self, CreateLutError>
    where
        I: IntoIterator,
        I: Clone,
        I::IntoIter: Clone,
        I::Item: Clone,
        I::Item: NumCast,
        I::Item: ToPrimitive,
    {
        let samples_f64 = samples.into_iter().filter_map(|v| v.to_f64());
        let min: f64 = samples_f64.clone().fold(f64::MAX, |a, b| a.min(b));
        let max: f64 = samples_f64.fold(f64::MIN, |a, b| a.max(b));

        // convert to f64 and swap if negative slope
        let (min, max) = if rescale.slope < 0. {
            (max, min)
        } else {
            (min, max)
        };

        let max = max * rescale.slope + rescale.intercept;
        let min = min * rescale.slope + rescale.intercept;

        // create a linear window level transform
        let voi = WindowLevelTransform::linear(crate::WindowLevel {
            width: max - min + 1.,
            center: (min + max) / 2.,
        });

        Self::new_rescale_and_window(bits_stored, signed, rescale, voi)
    }

    /// Create a new LUT containing the modality rescale transformation
    /// and the VOI transformation defined by a window level.
    ///
    /// The amplitude of the output values
    /// goes from 0 to `2^n - 1`, where `n` is the power of two
    /// which follows `bits_stored` (or itself if it is a power of two).
    /// For instance, if `bits_stored` is 12, the output values
    /// will go from 0 to 65535.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `rescale`: the rescale parameters
    /// - `voi`: the value of interest (VOI) function and parameters
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_rescale_and_window(
        bits_stored: u16,
        signed: bool,
        rescale: Rescale,
        voi: WindowLevelTransform,
    ) -> Result<Self, CreateLutError> {
        let bits_allocated = (bits_stored as usize).next_power_of_two();
        let y_max = ((1 << bits_allocated) - 1) as f64;
        Self::new_with_fn(bits_stored, signed, |v| {
            let x = v as f64;
            let v = rescale.apply(x);
            voi.apply(v, y_max)
        })
    }

    /// Create a new LUT containing
    /// a VOI transformation defined by a window level.
    ///
    /// The amplitude of the output values
    /// goes from 0 to `2^n - 1`, where `n` is the power of two
    /// which follows `bits_stored` (or itself if it is a power of two).
    /// For instance, if `bits_stored` is 12, the output values
    /// will go from 0 to 65535.
    ///
    /// - `bits_stored`:
    ///   the number of bits effectively used to represent the sample values
    ///   (the _Bits Stored_ DICOM attribute)
    /// - `signed`:
    ///   whether the input sample values are expected to be signed
    ///   (_Pixel Representation_ = 1)
    /// - `voi`: the value of interest (VOI) function and parameters
    ///
    /// # Panics
    ///
    /// Panics if `bits_stored` is 0 or too large.
    pub fn new_window(
        bits_stored: u16,
        signed: bool,
        voi: WindowLevelTransform,
    ) -> Result<Self, CreateLutError> {
        let bits_allocated = (bits_stored as usize).next_power_of_two();
        let y_max = ((1 << bits_allocated) - 1) as f64;
        Self::new_with_fn(bits_stored, signed, |v| voi.apply(v as f64, y_max))
    }

    /// Apply the transformation to a single pixel sample value.
    ///
    /// Although the input is expected to be one of `u8`, `u16`, or `u32`,
    /// this method works for signed sample values as well,
    /// with the bits reinterpreted as their unsigned counterpart.
    ///
    /// # Panics
    ///
    /// Panics if `sample_value` is larger or equal to `2^bits_stored`.
    pub fn get<I: 'static>(&self, sample_value: I) -> T
    where
        I: Copy,
        I: Into<u32>,
    {
        let val = sample_value.into();
        let index = if self.signed {
            // adjust for signedness by masking out the extra sign bits
            let mask = self.table.len() - 1;
            val as usize & mask
        } else {
            val as usize
        };
        assert!((index as usize) < self.table.len());

        self.table[index as usize]
    }

    /// Adapts an iterator of pixel data sample values
    /// to an iterator of transformed values.
    pub fn map_iter<'a, I: 'static>(
        &'a self,
        iter: impl IntoIterator<Item = I> + 'a,
    ) -> impl Iterator<Item = T> + 'a
    where
        I: Copy,
        I: Into<u32>,
    {
        iter.into_iter().map(move |i| self.get(i))
    }

    /// Adapts a parallel iterator of pixel data sample values
    /// to a parallel iterator of transformed values.
    pub fn map_par_iter<'a, I: 'static>(
        &'a self,
        iter: impl ParallelIterator<Item = I> + 'a,
    ) -> impl ParallelIterator<Item = T> + 'a
    where
        I: Copy,
        I: Into<u32>,
    {
        iter.map(move |i| self.get(i))
    }
}

#[cfg(test)]
mod tests {
    use crate::{VoiLutFunction, WindowLevel};

    use super::*;

    /// Applying a common rescale function to a value
    /// gives the expected output.
    #[test]
    fn modality_lut_baseline_2() {
        let rescale = Rescale::new(2., -1024.);

        assert_eq!(rescale.apply(0.), -1024.);
        assert_eq!(rescale.apply(1.), -1022.);
        assert_eq!(rescale.apply(2.), -1020.);
        assert_eq!(rescale.apply(1024.), 1024.);
    }

    #[test]
    fn lut_signed_numbers() {
        // 10-bit precision input, signed output
        let lut: Lut<i16> = Lut::new_rescale(10, true, Rescale::new(2., -1024.)).unwrap();

        assert_eq!(lut.get(0 as u16), -1024);
        assert_eq!(lut.get(1 as u16), -1022);
        assert_eq!(lut.get(-1_i16 as u16), -1026);
        assert_eq!(lut.get(-2_i16 as u16), -1028);
        assert_eq!(lut.get(500 as u16), -24);
    }

    #[test]
    fn lut_rescale_and_window_16bit() {
        let bits_stored = 16;
        let lut = Lut::new_rescale_and_window(
            bits_stored,
            false,
            Rescale::new(1., -1024.),
            WindowLevelTransform::new(
                VoiLutFunction::Linear,
                WindowLevel {
                    width: 300.,
                    center: 50.,
                },
            ),
        )
        .unwrap();

        // < 0
        let val: u16 = lut.get(824_u16);
        assert_eq!(val, 0);

        // > 200
        let val: u16 = lut.get(1224_u16);
        assert_eq!(val, 65535);

        // around the middle

        let val: u16 = lut.get(1024_u16 + 50);
        let expected_range = 32_500..=33_000;
        assert!(
            expected_range.contains(&val),
            "outcome was {}, expected to be in {:?}",
            val,
            expected_range,
        );
    }

    #[test]
    fn lut_rescale_and_normalize_16bit() {
        let bits_stored = 16;
        let lut: Lut<u16> = Lut::new_rescale_and_normalize(
            bits_stored,
            false,
            Rescale::new(1., -1024.),
            [0_u16, 1, 2, 500, 23].iter().copied(),
        ).unwrap();

        // samples are between 0 and 500
        assert_eq!(lut.get(0_u16), 0);
        assert_eq!(lut.get(500_u16), 0xFFFF);

        let y = lut.get(1_u16);
        assert!(y > 0 && y < 0xFFFF);

        let y = lut.get(498_u16);
        assert!(y > 0 && y < 0xFFFF);
    }
}
<|endoftext|>