<commit_before>// This test makes sure that the LLDB pretty printer does not throw an exception
// when trying to handle a Vec<> or anything else that contains zero-sized
// fields.

// min-lldb-version: 310
// ignore-gdb
// ignore-tidy-linelength

// compile-flags:-g

// === LLDB TESTS ==================================================================================
// lldb-command:run

// lldb-command:print v
// lldbg-check:[...]$0 = vec![1, 2, 3]
// lldbr-check:(alloc::vec::Vec<i32>) v = vec![1, 2, 3]
// lldb-command:print zs
// lldbg-check:[...]$1 = StructWithZeroSizedField { x: ZeroSizedStruct, y: 123, z: ZeroSizedStruct, w: 456 }
// lldbr-check:(issue_22656::StructWithZeroSizedField) zs = StructWithZeroSizedField { x: ZeroSizedStruct { }, y: 123, z: ZeroSizedStruct { }, w: 456 }
// lldbr-command:continue

#![allow(unused_variables)]
#![allow(dead_code)]
#![feature(omit_gdb_pretty_printer_section)]
#![omit_gdb_pretty_printer_section]

struct ZeroSizedStruct;

struct StructWithZeroSizedField {
    x: ZeroSizedStruct,
    y: u32,
    z: ZeroSizedStruct,
    w: u64
}

fn main() {
    let v = vec![1,2,3];

    let zs = StructWithZeroSizedField {
        x: ZeroSizedStruct,
        y: 123,
        z: ZeroSizedStruct,
        w: 456
    };

    zzz(); // #break
}

fn zzz() { () }
<commit_msg>Rollup merge of #64578 - max-sixty:22656-lldb-8, r=alexcrichton<commit_after>// This test makes sure that the LLDB pretty printer does not throw an exception
// when trying to handle a Vec<> or anything else that contains zero-sized
// fields.

// min-lldb-version: 310
// ignore-gdb
// ignore-tidy-linelength

// compile-flags:-g

// === LLDB TESTS ==================================================================================
// lldb-command:run

// lldb-command:print v
// lldbg-check:[...]$0 = vec![1, 2, 3]
// lldbr-check:(alloc::vec::Vec<i32>) v = vec![1, 2, 3]
// lldb-command:print zs
// lldbg-check:[...]$1 = StructWithZeroSizedField { x: ZeroSizedStruct[...], y: 123, z: ZeroSizedStruct[...], w: 456 }
// lldbr-check:(issue_22656::StructWithZeroSizedField) zs = StructWithZeroSizedField { x: ZeroSizedStruct { }, y: 123, z: ZeroSizedStruct { }, w: 456 }
// lldbr-command:continue

#![allow(unused_variables)]
#![allow(dead_code)]
#![feature(omit_gdb_pretty_printer_section)]
#![omit_gdb_pretty_printer_section]

struct ZeroSizedStruct;

struct StructWithZeroSizedField {
    x: ZeroSizedStruct,
    y: u32,
    z: ZeroSizedStruct,
    w: u64
}

fn main() {
    let v = vec![1,2,3];

    let zs = StructWithZeroSizedField {
        x: ZeroSizedStruct,
        y: 123,
        z: ZeroSizedStruct,
        w: 456
    };

    zzz(); // #break
}

fn zzz() { () }
<|endoftext|>


<commit_before>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// compile-pass
// only-x86_64
// no-prefer-dynamic
// compile-flags: -Ctarget-feature=+avx -Clto

fn main() {}<commit_msg>Add trailing newline<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// compile-pass
// only-x86_64
// no-prefer-dynamic
// compile-flags: -Ctarget-feature=+avx -Clto

fn main() {}
<|endoftext|>


<commit_before>use chrono::NaiveDateTime;
use conduit::{Request, Response};
use diesel::prelude::*;
use serde_json;

use db::RequestTransaction;
use schema::{crate_owner_invitations, crate_owners, crates, users};
use user::RequestUser;
use util::errors::{human, CargoResult};
use util::RequestUtils;
use owner::{CrateOwner, OwnerKind};

/// The model representing a row in the `crate_owner_invitations` database table.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Identifiable, Queryable)]
#[primary_key(invited_user_id, crate_id)]
pub struct CrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
    pub created_at: NaiveDateTime,
}

#[derive(Insertable, Clone, Copy, Debug)]
#[table_name = "crate_owner_invitations"]
pub struct NewCrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
}

impl CrateOwnerInvitation {
    pub fn invited_by_username(&self, conn: &PgConnection) -> String {
        users::table
            .find(self.invited_by_user_id)
            .select(users::gh_login)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown username)"))
    }

    pub fn crate_name(&self, conn: &PgConnection) -> String {
        crates::table
            .find(self.crate_id)
            .select(crates::name)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown crate name)"))
    }

    pub fn encodable(self, conn: &PgConnection) -> EncodableCrateOwnerInvitation {
        EncodableCrateOwnerInvitation {
            invited_by_username: self.invited_by_username(conn),
            crate_name: self.crate_name(conn),
            crate_id: self.crate_id,
            created_at: self.created_at,
        }
    }
}

/// The serialization format for the `CrateOwnerInvitation` model.
#[derive(Deserialize, Serialize, Debug)]
pub struct EncodableCrateOwnerInvitation {
    pub invited_by_username: String,
    pub crate_name: String,
    pub crate_id: i32,
    pub created_at: NaiveDateTime,
}

/// Handles the `GET /me/crate_owner_invitations` route.
pub fn list(req: &mut Request) -> CargoResult<Response> {
    let conn = &*req.db_conn()?;
    let user_id = req.user()?.id;

    let crate_owner_invitations = crate_owner_invitations::table
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .load::<CrateOwnerInvitation>(&*conn)?
        .into_iter()
        .map(|i| i.encodable(conn))
        .collect();

    #[derive(Serialize)]
    struct R {
        crate_owner_invitations: Vec<EncodableCrateOwnerInvitation>,
    }
    Ok(req.json(&R {
        crate_owner_invitations,
    }))
}

#[derive(Deserialize)]
struct OwnerInvitation {
    crate_owner_invite: InvitationResponse,
}

#[derive(Deserialize, Serialize, Debug, Copy, Clone)]
pub struct InvitationResponse {
    pub crate_id: i32,
    pub accepted: bool,
}

/// Handles the `PUT /me/crate_owner_invitations/:crate_id` route.
pub fn handle_invite(req: &mut Request) -> CargoResult<Response> {
    let conn = &*req.db_conn()?;

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;

    let crate_invite: OwnerInvitation =
        serde_json::from_str(&body).map_err(|_| human("invalid json request"))?;

    let crate_invite = crate_invite.crate_owner_invite;

    if crate_invite.accepted {
        accept_invite(req, conn, crate_invite)
    } else {
        decline_invite(req, conn, crate_invite)
    }
}

fn accept_invite(
    req: &mut Request,
    conn: &PgConnection,
    crate_invite: InvitationResponse,
) -> CargoResult<Response> {
    let user_id = req.user()?.id;
    use diesel::{insert, delete};
    use diesel::pg::upsert::do_update;
    use diesel::pg::upsert::OnConflictExtension;
    let pending_crate_owner = crate_owner_invitations::table
        .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .first::<CrateOwnerInvitation>(&*conn)?;

    let owner = CrateOwner {
        crate_id: crate_invite.crate_id,
        owner_id: user_id,
        created_by: pending_crate_owner.invited_by_user_id,
        owner_kind: OwnerKind::User as i32,
    };

    conn.transaction(|| {
        insert(&owner.on_conflict(
            crate_owners::table.primary_key(),
            do_update().set(crate_owners::deleted.eq(false)),
        )).into(crate_owners::table)
            .execute(conn)?;
        delete(
            crate_owner_invitations::table
                .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
                .filter(crate_owner_invitations::invited_user_id.eq(user_id)),
        ).execute(conn)?;

        #[derive(Serialize)]
        struct R {
            crate_owner_invitation: InvitationResponse,
        }
        Ok(req.json(&R {
            crate_owner_invitation: crate_invite,
        }))
    })
}

fn decline_invite(
    req: &mut Request,
    conn: &PgConnection,
    crate_invite: InvitationResponse,
) -> CargoResult<Response> {
    use diesel::delete;
    let user_id = req.user()?.id;

    delete(
        crate_owner_invitations::table
            .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
            .filter(crate_owner_invitations::invited_user_id.eq(user_id)),
    ).execute(conn)?;

    #[derive(Serialize)]
    struct R {
        crate_owner_invitation: InvitationResponse,
    }

    Ok(req.json(&R {
        crate_owner_invitation: crate_invite,
    }))
}
<commit_msg>reorder statements at top of function<commit_after>use chrono::NaiveDateTime;
use conduit::{Request, Response};
use diesel::prelude::*;
use serde_json;

use db::RequestTransaction;
use schema::{crate_owner_invitations, crate_owners, crates, users};
use user::RequestUser;
use util::errors::{human, CargoResult};
use util::RequestUtils;
use owner::{CrateOwner, OwnerKind};

/// The model representing a row in the `crate_owner_invitations` database table.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Identifiable, Queryable)]
#[primary_key(invited_user_id, crate_id)]
pub struct CrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
    pub created_at: NaiveDateTime,
}

#[derive(Insertable, Clone, Copy, Debug)]
#[table_name = "crate_owner_invitations"]
pub struct NewCrateOwnerInvitation {
    pub invited_user_id: i32,
    pub invited_by_user_id: i32,
    pub crate_id: i32,
}

impl CrateOwnerInvitation {
    pub fn invited_by_username(&self, conn: &PgConnection) -> String {
        users::table
            .find(self.invited_by_user_id)
            .select(users::gh_login)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown username)"))
    }

    pub fn crate_name(&self, conn: &PgConnection) -> String {
        crates::table
            .find(self.crate_id)
            .select(crates::name)
            .first(&*conn)
            .unwrap_or_else(|_| String::from("(unknown crate name)"))
    }

    pub fn encodable(self, conn: &PgConnection) -> EncodableCrateOwnerInvitation {
        EncodableCrateOwnerInvitation {
            invited_by_username: self.invited_by_username(conn),
            crate_name: self.crate_name(conn),
            crate_id: self.crate_id,
            created_at: self.created_at,
        }
    }
}

/// The serialization format for the `CrateOwnerInvitation` model.
#[derive(Deserialize, Serialize, Debug)]
pub struct EncodableCrateOwnerInvitation {
    pub invited_by_username: String,
    pub crate_name: String,
    pub crate_id: i32,
    pub created_at: NaiveDateTime,
}

/// Handles the `GET /me/crate_owner_invitations` route.
pub fn list(req: &mut Request) -> CargoResult<Response> {
    let conn = &*req.db_conn()?;
    let user_id = req.user()?.id;

    let crate_owner_invitations = crate_owner_invitations::table
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .load::<CrateOwnerInvitation>(&*conn)?
        .into_iter()
        .map(|i| i.encodable(conn))
        .collect();

    #[derive(Serialize)]
    struct R {
        crate_owner_invitations: Vec<EncodableCrateOwnerInvitation>,
    }
    Ok(req.json(&R {
        crate_owner_invitations,
    }))
}

#[derive(Deserialize)]
struct OwnerInvitation {
    crate_owner_invite: InvitationResponse,
}

#[derive(Deserialize, Serialize, Debug, Copy, Clone)]
pub struct InvitationResponse {
    pub crate_id: i32,
    pub accepted: bool,
}

/// Handles the `PUT /me/crate_owner_invitations/:crate_id` route.
pub fn handle_invite(req: &mut Request) -> CargoResult<Response> {
    let conn = &*req.db_conn()?;

    let mut body = String::new();
    req.body().read_to_string(&mut body)?;

    let crate_invite: OwnerInvitation =
        serde_json::from_str(&body).map_err(|_| human("invalid json request"))?;

    let crate_invite = crate_invite.crate_owner_invite;

    if crate_invite.accepted {
        accept_invite(req, conn, crate_invite)
    } else {
        decline_invite(req, conn, crate_invite)
    }
}

fn accept_invite(
    req: &mut Request,
    conn: &PgConnection,
    crate_invite: InvitationResponse,
) -> CargoResult<Response> {
    use diesel::{insert, delete};
    use diesel::pg::upsert::{OnConflictExtension, do_update};

    let user_id = req.user()?.id;
    let pending_crate_owner = crate_owner_invitations::table
        .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
        .filter(crate_owner_invitations::invited_user_id.eq(user_id))
        .first::<CrateOwnerInvitation>(&*conn)?;

    let owner = CrateOwner {
        crate_id: crate_invite.crate_id,
        owner_id: user_id,
        created_by: pending_crate_owner.invited_by_user_id,
        owner_kind: OwnerKind::User as i32,
    };

    conn.transaction(|| {
        insert(&owner.on_conflict(
            crate_owners::table.primary_key(),
            do_update().set(crate_owners::deleted.eq(false)),
        )).into(crate_owners::table)
            .execute(conn)?;
        delete(
            crate_owner_invitations::table
                .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
                .filter(crate_owner_invitations::invited_user_id.eq(user_id)),
        ).execute(conn)?;

        #[derive(Serialize)]
        struct R {
            crate_owner_invitation: InvitationResponse,
        }
        Ok(req.json(&R {
            crate_owner_invitation: crate_invite,
        }))
    })
}

fn decline_invite(
    req: &mut Request,
    conn: &PgConnection,
    crate_invite: InvitationResponse,
) -> CargoResult<Response> {
    use diesel::delete;
    let user_id = req.user()?.id;

    delete(
        crate_owner_invitations::table
            .filter(crate_owner_invitations::crate_id.eq(crate_invite.crate_id))
            .filter(crate_owner_invitations::invited_user_id.eq(user_id)),
    ).execute(conn)?;

    #[derive(Serialize)]
    struct R {
        crate_owner_invitation: InvitationResponse,
    }

    Ok(req.json(&R {
        crate_owner_invitation: crate_invite,
    }))
}
<|endoftext|>


<commit_before>// Copyright 2017 PingCAP, Inc.
//
// 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
//
//     http://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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fs;
use std::io::{Error, ErrorKind, Read, Result};
use std::sync::Mutex;

use libc::{self, pid_t};
use prometheus::core::{Collector, Desc};
use prometheus::{self, proto, CounterVec, IntGaugeVec, Opts};

/// monitor current process's threads.
pub fn monitor_threads<S: Into<String>>(namespace: S) -> Result<()> {
    let pid = unsafe { libc::getpid() };
    let tc = ThreadsCollector::new(pid, namespace);
    prometheus::register(Box::new(tc)).map_err(|e| to_io_err(format!("{:?}", e)))
}

/// A collector to collect threads metrics, including CPU usage
/// and threads state.
struct ThreadsCollector {
    pid: pid_t,
    descs: Vec<Desc>,
    // There are several step to collect CPU time, we need to synchronous them.
    cpu_totals: Mutex<CounterVec>,
    threads_state: IntGaugeVec,
}

impl ThreadsCollector {
    fn new<S: Into<String>>(pid: pid_t, namespace: S) -> ThreadsCollector {
        let mut descs: Vec<Desc> = vec![];
        let ns = namespace.into();
        let cpu_totals = CounterVec::new(
            Opts::new(
                "thread_cpu_seconds_total",
                "Total user and system CPU time spent in \
                 seconds by threads.",
            ).namespace(ns.clone()),
            &["name", "tid"],
        ).unwrap();
        descs.extend(cpu_totals.desc().into_iter().cloned());
        let threads_state = IntGaugeVec::new(
            Opts::new("threads_state", "Number of threads in each state.").namespace(ns),
            &["state"],
        ).unwrap();
        descs.extend(threads_state.desc().into_iter().cloned());

        ThreadsCollector {
            pid,
            descs,
            cpu_totals: Mutex::new(cpu_totals),
            threads_state,
        }
    }
}

impl Collector for ThreadsCollector {
    fn desc(&self) -> Vec<&Desc> {
        self.descs.iter().collect()
    }

    fn collect(&self) -> Vec<proto::MetricFamily> {
        // Synchronous collecting CPU time.
        let cpu_totals = self.cpu_totals.lock().unwrap();
        // Clean previous threads state.
        self.threads_state.reset();
        let tids = get_thread_ids(self.pid).unwrap();
        for tid in tids {
            if let Ok(Stat {
                name,
                state,
                utime,
                stime,
            }) = Stat::collect(self.pid, tid)
            {
                // Threads CPU time.
                let total = (utime + stime) / *CLK_TCK;
                let cpu_total = cpu_totals
                    .get_metric_with_label_values(&[&name, &format!("{}", tid)])
                    .unwrap();
                let past = cpu_total.get();
                let delta = total - past;
                if delta > 0.0 {
                    cpu_total.inc_by(delta);
                }

                // Threads states.
                let state = self
                    .threads_state
                    .get_metric_with_label_values(&[&state])
                    .unwrap();
                state.inc();
            }
        }

        let mut mfs = cpu_totals.collect();
        mfs.extend(self.threads_state.collect());
        mfs
    }
}

fn get_thread_ids(pid: pid_t) -> Result<Vec<pid_t>> {
    let mut tids = Vec::new();
    let dirs = fs::read_dir(format!("/proc/{}/task", pid))?;
    for task in dirs {
        let file_name = match task {
            Ok(t) => t.file_name(),
            Err(e) => {
                error!("fail to read task of {}, error {:?}", pid, e);
                continue;
            }
        };

        let tid = match file_name.to_str() {
            Some(tid) => match tid.parse() {
                Ok(tid) => tid,
                Err(e) => {
                    error!("fail to read task of {}, error {:?}", pid, e);
                    continue;
                }
            },
            None => {
                error!("fail to read task of {}", pid);
                continue;
            }
        };
        tids.push(tid);
    }

    Ok(tids)
}

// get thread name and the index of the last character(including ')').
fn get_thread_name(tid: pid_t, stat: &str) -> Result<(String, usize)> {
    let start = stat.find('(');
    let end = stat.rfind(')');

    if let (Some(start), Some(end)) = (start, end) {
        let raw = &stat[start + 1..end];
        let mut name = String::with_capacity(raw.len());

        // sanitize thread name.
        for c in raw.chars() {
            match c {
                // Prometheus label characters `[a-zA-Z0-9_:]`
                'a'...'z' | 'A'...'Z' | '0'...'9' | '_' | ':' => {
                    name.push(c);
                }
                '-' | ' ' => {
                    name.push('_');
                }
                _ => (),
            }
        }

        if name.is_empty() {
            name = format!("{}", tid)
        }

        return Ok((name, end));
    }

    Err(to_io_err(format!(
        "can not find thread name, stat: {}",
        stat
    )))
}

fn to_io_err(s: String) -> Error {
    Error::new(ErrorKind::Other, s)
}

struct Stat {
    name: String,
    state: String,
    utime: f64,
    stime: f64,
}

impl Stat {
    /// See more man proc.
    /// Index of utime and stime.
    const CPU_INDEX: [usize; 2] = [14 - 1, 15 - 1];
    /// Index of state.
    const PROCESS_STATE_INDEX: usize = 3 - 1;

    fn collect(pid: pid_t, tid: pid_t) -> Result<Stat> {
        let mut stat = String::new();
        fs::File::open(format!("/proc/{}/task/{}/stat", pid, tid))
            .and_then(|mut f| f.read_to_string(&mut stat))?;
        get_thread_stat_internal(tid, &stat)
    }
}

// Extracted from `Stat::collect`, for test purpose.
fn get_thread_stat_internal(tid: pid_t, stat: &str) -> Result<Stat> {
    let (name, end) = get_thread_name(tid, stat)?;
    let stats: Vec<_> = (&stat[end + 2..]).split_whitespace().collect(); // excluding ") ".
    let utime = stats
        .get(Stat::CPU_INDEX[0] - 2) // -2 because pid and comm is truncated.
        .unwrap_or(&"0")
        .parse()
        .map_err(|e| to_io_err(format!("{:?}: {}", e, stat)))?;
    let stime = stats
        .get(Stat::CPU_INDEX[1] - 2)
        .unwrap_or(&"0")
        .parse()
        .map_err(|e| to_io_err(format!("{:?}: {}", e, stat)))?;
    let state = stats
        .get(Stat::PROCESS_STATE_INDEX - 2)
        .unwrap_or(&"unknown")
        .to_string();
    Ok(Stat {
        name,
        state,
        utime,
        stime,
    })
}

lazy_static! {
    // getconf CLK_TCK
    static ref CLK_TCK: f64 = {
        unsafe {
            libc::sysconf(libc::_SC_CLK_TCK) as f64
        }
    };
}

#[cfg(test)]
mod tests {
    use std::sync;
    use std::thread;

    use libc;

    use super::*;

    #[test]
    fn test_thread_stat() {
        let name = "theadnametest66";
        let (tx, rx) = sync::mpsc::channel();
        let (tx1, rx1) = sync::mpsc::channel();
        let h = thread::Builder::new()
            .name(name.to_owned())
            .spawn(move || {
                tx1.send(()).unwrap();
                rx.recv().unwrap();
            })
            .unwrap();
        rx1.recv().unwrap();

        let pid = unsafe { libc::getpid() };
        let tids = get_thread_ids(pid).unwrap();
        assert!(tids.len() >= 2);

        tids.iter()
            .find(|t| {
                Stat::collect(pid, **t)
                    .map(|stat| stat.name == name)
                    .unwrap_or(false)
            })
            .unwrap();

        tx.send(()).unwrap();
        h.join().unwrap();
    }

    #[test]
    fn test_get_thread_name() {
        let cases = [
            ("(ok123)", "ok123", 6),
            ("(a-b)", "a_b", 4),
            ("(Az_1:)", "Az_1:", 6),
            ("(@123)", "123", 5),
            ("1 (ab) 1", "ab", 5),
            ("1 (a b) 1", "a_b", 6),
            ("1 ((a b )) 1", "a_b_", 9),
        ];
        for &(i, o, idx) in &cases {
            let (name, end) = get_thread_name(1, i).unwrap();
            assert_eq!(name, o);
            assert_eq!(end, idx);
        }

        assert_eq!(get_thread_name(1, "(@#)").unwrap().0, "1");
        assert!(get_thread_name(1, "invalid_stat").is_err());
    }

    #[test]
    fn test_get_thread_stat() {
        let sample = "2810 (test thd) S 2550 2621 2621 0 -1 4210688 2632 0 52 0 839 138 0 \
                      0 20 0 4 0 13862 709652480 3647 18446744073709551615 4194304 4319028 \
                      140732554845776 140732554845392 140439688777693 0 0 4096 16384 0 0 0 17 3 \
                      0 0 245 0 0 6417696 6421000 8478720 140732554851684 140732554851747 \
                      140732554851747 140732554854339 0";

        let Stat {
            name,
            state,
            utime,
            stime,
        } = get_thread_stat_internal(2810, sample).unwrap();
        assert_eq!(name, "test_thd");
        assert_eq!(state, "S");
        assert_eq!(utime as i64, 839);
        assert_eq!(stime as i64, 138);
    }

    #[test]
    fn test_smoke() {
        monitor_threads("smoke").unwrap();
    }
}
<commit_msg>util/metrics: add threads IO metrics (#3338)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// 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
//
//     http://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,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fs;
use std::io::{Error, ErrorKind, Read, Result};
use std::sync::Mutex;

use libc::{self, pid_t};
use prometheus::core::{Collector, Desc};
use prometheus::{self, proto, CounterVec, IntGaugeVec, Opts};

/// monitor current process's threads.
pub fn monitor_threads<S: Into<String>>(namespace: S) -> Result<()> {
    let pid = unsafe { libc::getpid() };
    let tc = ThreadsCollector::new(pid, namespace);
    prometheus::register(Box::new(tc)).map_err(|e| to_io_err(format!("{:?}", e)))
}

struct Metrics {
    cpu_totals: CounterVec,
    io_totals: CounterVec,
    threads_state: IntGaugeVec,
}

/// A collector to collect threads metrics, including CPU usage
/// and threads state.
struct ThreadsCollector {
    pid: pid_t,
    descs: Vec<Desc>,
    metrics: Mutex<Metrics>,
}

impl ThreadsCollector {
    fn new<S: Into<String>>(pid: pid_t, namespace: S) -> ThreadsCollector {
        let mut descs: Vec<Desc> = vec![];
        let ns = namespace.into();
        let cpu_totals = CounterVec::new(
            Opts::new(
                "thread_cpu_seconds_total",
                "Total user and system CPU time spent in \
                 seconds by threads.",
            ).namespace(ns.clone()),
            &["name", "tid"],
        ).unwrap();
        descs.extend(cpu_totals.desc().into_iter().cloned());
        let threads_state = IntGaugeVec::new(
            Opts::new("threads_state", "Number of threads in each state.").namespace(ns.clone()),
            &["state"],
        ).unwrap();
        descs.extend(threads_state.desc().into_iter().cloned());
        let io_totals = CounterVec::new(
            Opts::new(
                "threads_io_bytes_total",
                "Total number of bytes which threads cause to be fetched from or sent to the storage layer.",
            ).namespace(ns),
            &["name", "tid", "io"],
        ).unwrap();
        descs.extend(io_totals.desc().into_iter().cloned());

        ThreadsCollector {
            pid,
            descs,
            metrics: Mutex::new(Metrics {
                cpu_totals,
                io_totals,
                threads_state,
            }),
        }
    }
}

impl Collector for ThreadsCollector {
    fn desc(&self) -> Vec<&Desc> {
        self.descs.iter().collect()
    }

    fn collect(&self) -> Vec<proto::MetricFamily> {
        // Synchronous collecting metrics.
        let metrics = self.metrics.lock().unwrap();
        // Clean previous threads state.
        metrics.threads_state.reset();

        let tids = get_thread_ids(self.pid).unwrap();
        for tid in tids {
            if let Ok(Stat {
                name,
                state,
                utime,
                stime,
            }) = Stat::collect(self.pid, tid)
            {
                // Threads CPU time.
                let total = (utime + stime) / *CLK_TCK;
                let cpu_total = metrics
                    .cpu_totals
                    .get_metric_with_label_values(&[&name, &format!("{}", tid)])
                    .unwrap();
                let past = cpu_total.get();
                let delta = total - past;
                if delta > 0.0 {
                    cpu_total.inc_by(delta);
                }

                // Threads states.
                let state = metrics
                    .threads_state
                    .get_metric_with_label_values(&[&state])
                    .unwrap();
                state.inc();

                if let Ok(Io {
                    read_bytes,
                    write_bytes,
                }) = Io::collect(self.pid, tid)
                {
                    // Threads IO.
                    let read_total = metrics
                        .io_totals
                        .get_metric_with_label_values(&[&name, &format!("{}", tid), "read"])
                        .unwrap();
                    let read_past = read_total.get();
                    let read_delta = read_bytes as f64 - read_past;
                    if read_delta > 0.0 {
                        read_total.inc_by(read_delta);
                    }

                    let write_total = metrics
                        .io_totals
                        .get_metric_with_label_values(&[&name, &format!("{}", tid), "write"])
                        .unwrap();
                    let write_past = write_total.get();
                    let write_delta = write_bytes as f64 - write_past;
                    if write_delta > 0.0 {
                        write_total.inc_by(write_delta);
                    }
                }
            }
        }
        let mut mfs = metrics.cpu_totals.collect();
        mfs.extend(metrics.threads_state.collect());
        mfs.extend(metrics.io_totals.collect());
        mfs
    }
}

fn get_thread_ids(pid: pid_t) -> Result<Vec<pid_t>> {
    let mut tids = Vec::new();
    let dirs = fs::read_dir(format!("/proc/{}/task", pid))?;
    for task in dirs {
        let file_name = match task {
            Ok(t) => t.file_name(),
            Err(e) => {
                error!("fail to read task of {}, error {:?}", pid, e);
                continue;
            }
        };

        let tid = match file_name.to_str() {
            Some(tid) => match tid.parse() {
                Ok(tid) => tid,
                Err(e) => {
                    error!("fail to read task of {}, error {:?}", pid, e);
                    continue;
                }
            },
            None => {
                error!("fail to read task of {}", pid);
                continue;
            }
        };
        tids.push(tid);
    }

    Ok(tids)
}

// get thread name and the index of the last character(including ')').
fn get_thread_name(tid: pid_t, stat: &str) -> Result<(String, usize)> {
    let start = stat.find('(');
    let end = stat.rfind(')');

    if let (Some(start), Some(end)) = (start, end) {
        let raw = &stat[start + 1..end];
        let mut name = String::with_capacity(raw.len());

        // sanitize thread name.
        for c in raw.chars() {
            match c {
                // Prometheus label characters `[a-zA-Z0-9_:]`
                'a'...'z' | 'A'...'Z' | '0'...'9' | '_' | ':' => {
                    name.push(c);
                }
                '-' | ' ' => {
                    name.push('_');
                }
                _ => (),
            }
        }

        if name.is_empty() {
            name = format!("{}", tid)
        }

        return Ok((name, end));
    }

    Err(to_io_err(format!(
        "can not find thread name, stat: {}",
        stat
    )))
}

fn to_io_err(s: String) -> Error {
    Error::new(ErrorKind::Other, s)
}

struct Stat {
    name: String,
    state: String,
    utime: f64,
    stime: f64,
}

impl Stat {
    /// See more man proc.
    /// Index of utime and stime.
    const CPU_INDEX: [usize; 2] = [14 - 1, 15 - 1];
    /// Index of state.
    const PROCESS_STATE_INDEX: usize = 3 - 1;

    fn collect(pid: pid_t, tid: pid_t) -> Result<Stat> {
        let mut stat = String::new();
        fs::File::open(format!("/proc/{}/task/{}/stat", pid, tid))
            .and_then(|mut f| f.read_to_string(&mut stat))?;
        get_thread_stat_internal(tid, &stat)
    }
}

// Extracted from `Stat::collect`, for test purpose.
fn get_thread_stat_internal(tid: pid_t, stat: &str) -> Result<Stat> {
    let (name, end) = get_thread_name(tid, stat)?;
    let stats: Vec<_> = (&stat[end + 2..]).split_whitespace().collect(); // excluding ") ".
    let utime = stats
        .get(Stat::CPU_INDEX[0] - 2) // -2 because pid and comm is truncated.
        .unwrap_or(&"0")
        .parse()
        .map_err(|e| to_io_err(format!("{:?}: {}", e, stat)))?;
    let stime = stats
        .get(Stat::CPU_INDEX[1] - 2)
        .unwrap_or(&"0")
        .parse()
        .map_err(|e| to_io_err(format!("{:?}: {}", e, stat)))?;
    let state = stats
        .get(Stat::PROCESS_STATE_INDEX - 2)
        .unwrap_or(&"unknown")
        .to_string();
    Ok(Stat {
        name,
        state,
        utime,
        stime,
    })
}

// I/O statistics for threads.
struct Io {
    // Attempt to count the number of bytes which this process really did cause
    // to be fetched from the storage layer.  This is accurate for block-backed
    // filesystems.
    read_bytes: u64,
    // Attempt to count the number of bytes which this process caused to be
    // sent to the storage layer.
    write_bytes: u64,
}

impl Io {
    // # cat /proc/3828/io
    // rchar: 323934931
    // wchar: 323929600
    // syscr: 632687
    // syscw: 632675
    // read_bytes: 0
    // write_bytes: 323932160
    // cancelled_write_bytes: 0
    const READ_BYTES_INDEX: usize = 4;
    const WRITE_BYTES_INDEX: usize = 5;

    fn collect(pid: pid_t, tid: pid_t) -> Result<Io> {
        let mut io = String::new();
        fs::File::open(format!("/proc/{}/task/{}/io", pid, tid))
            .and_then(|mut f| f.read_to_string(&mut io))?;
        get_thread_io_internal(&io)
    }
}

// Extracted from `Io::collect`, for test purpose.
fn get_thread_io_internal(io: &str) -> Result<Io> {
    let read_bytes = io
        .lines()
        .nth(Io::READ_BYTES_INDEX)
        .map_or_else(|| Err(to_io_err(io.to_owned())), Ok)?
        .split(':')
        .nth(1)
        .map_or_else(|| Err(to_io_err(io.to_owned())), Ok)?
        .trim()
        .parse()
        .map_err(|e| to_io_err(format!("{:?}: {}", e, io)))?;

    let write_bytes = io
        .lines()
        .nth(Io::WRITE_BYTES_INDEX)
        .map_or_else(|| Err(to_io_err(io.to_owned())), Ok)?
        .split(':')
        .nth(1)
        .map_or_else(|| Err(to_io_err(io.to_owned())), Ok)?
        .trim()
        .parse()
        .map_err(|e| to_io_err(format!("{:?}: {}", e, io)))?;

    Ok(Io {
        read_bytes,
        write_bytes,
    })
}

lazy_static! {
    // getconf CLK_TCK
    static ref CLK_TCK: f64 = {
        unsafe {
            libc::sysconf(libc::_SC_CLK_TCK) as f64
        }
    };
}

#[cfg(test)]
mod tests {
    use std::sync;
    use std::thread;

    use libc;

    use super::*;

    #[test]
    fn test_thread_stat() {
        let name = "theadnametest66";
        let (tx, rx) = sync::mpsc::channel();
        let (tx1, rx1) = sync::mpsc::channel();
        let h = thread::Builder::new()
            .name(name.to_owned())
            .spawn(move || {
                tx1.send(()).unwrap();
                rx.recv().unwrap();
            })
            .unwrap();
        rx1.recv().unwrap();

        let pid = unsafe { libc::getpid() };
        let tids = get_thread_ids(pid).unwrap();
        assert!(tids.len() >= 2);

        tids.iter()
            .find(|t| {
                Stat::collect(pid, **t)
                    .map(|stat| stat.name == name)
                    .unwrap_or(false)
            })
            .unwrap();

        tx.send(()).unwrap();
        h.join().unwrap();
    }

    #[test]
    fn test_get_thread_name() {
        let cases = [
            ("(ok123)", "ok123", 6),
            ("(a-b)", "a_b", 4),
            ("(Az_1:)", "Az_1:", 6),
            ("(@123)", "123", 5),
            ("1 (ab) 1", "ab", 5),
            ("1 (a b) 1", "a_b", 6),
            ("1 ((a b )) 1", "a_b_", 9),
        ];
        for &(i, o, idx) in &cases {
            let (name, end) = get_thread_name(1, i).unwrap();
            assert_eq!(name, o);
            assert_eq!(end, idx);
        }

        assert_eq!(get_thread_name(1, "(@#)").unwrap().0, "1");
        assert!(get_thread_name(1, "invalid_stat").is_err());
    }

    #[test]
    fn test_get_thread_stat() {
        let sample = "2810 (test thd) S 2550 2621 2621 0 -1 4210688 2632 0 52 0 839 138 0 \
                      0 20 0 4 0 13862 709652480 3647 18446744073709551615 4194304 4319028 \
                      140732554845776 140732554845392 140439688777693 0 0 4096 16384 0 0 0 17 3 \
                      0 0 245 0 0 6417696 6421000 8478720 140732554851684 140732554851747 \
                      140732554851747 140732554854339 0";

        let Stat {
            name,
            state,
            utime,
            stime,
        } = get_thread_stat_internal(2810, sample).unwrap();
        assert_eq!(name, "test_thd");
        assert_eq!(state, "S");
        assert_eq!(utime as i64, 839);
        assert_eq!(stime as i64, 138);
    }

    #[test]
    fn test_get_thread_io() {
        let sample = "rchar: 323934931
wchar: 323929600
syscr: 632687
syscw: 632675
read_bytes: 7878789
write_bytes: 323932170
cancelled_write_bytes: 0";

        let Io {
            read_bytes,
            write_bytes,
        } = get_thread_io_internal(sample).unwrap();
        assert_eq!(read_bytes as i64, 7878789);
        assert_eq!(write_bytes as i64, 323932170);
    }

    #[test]
    fn test_smoke() {
        let pid = unsafe { libc::getpid() };
        let tc = ThreadsCollector::new(pid, "smoke");
        tc.collect();
        tc.desc();
        monitor_threads("smoke").unwrap();
    }
}
<|endoftext|>


<commit_before>use business::db::Db;
use entities::*;
use rusted_cypher::GraphClient;
use business::error::RepoError;
use std::result;

type Result<T> = result::Result<T, RepoError>;

fn neo4j_edge_name<'a>(t: &'a Triple) -> &'a str {
    match t.predicate {
       Relation::IsTaggedWith    => "IS_TAGGED_WITH",
       Relation::IsRatedWith     => "IS_RATED_WITH",
       Relation::IsCommentedWith => "IS_COMMENTED_WITH",
       Relation::CreatedBy       => "CREATED_BY"
    }.into()
}

fn object_id_to_neo4j_label<'a>(id: &'a ObjectId) -> (&'a str,&'a str) {
    match *id {
        ObjectId::Entry(ref id) => ("Entry",id),
        ObjectId::Tag(ref id) => ("Tag",id),
        ObjectId::User(ref id) => ("User",id),
        ObjectId::Comment(ref id) => ("Comment",id),
        ObjectId::Rating(ref id) => ("Rating",id)
    }
}

fn neo4j_label_to_object_id(label: &str, id: String) -> Option<ObjectId> {
    match label {
        "Entry"     => Some(ObjectId::Entry(id)),
        "Tag"       => Some(ObjectId::Tag(id)),
        "User"      => Some(ObjectId::User(id)),
        "Comment"   => Some(ObjectId::Comment(id)),
        "Rating"    => Some(ObjectId::Rating(id)),
        _           => None,
    }
}

fn neo4j_relation_to_relation(rel: &str) -> Option<Relation> {
    match rel {
        "IS_TAGGED_WITH"    => Some(Relation::IsTaggedWith),
        "IS_RATED_WITH"     => Some(Relation::IsRatedWith),
        "IS_COMMENTED_WITH" => Some(Relation::IsCommentedWith),
        "CREATED_BY"        => Some(Relation::CreatedBy),
        _                   => None
    }
}

#[derive(Deserialize,Debug)]
struct Neo4jObj {
    labels: Vec<String>,
    id: String
}

#[derive(Deserialize,Debug)]
struct Neo4jTriple {
    subject: Neo4jObj,
    predicate: String,
    object: Neo4jObj
}

fn from_neo_triple(t: Neo4jTriple) -> Option<Triple> {
    if t.subject.labels.is_empty() || t.object.labels.is_empty() {
        warn!("No labels found in {:?}",t);
        return None;
    }
    let subject = neo4j_label_to_object_id(&t.subject.labels[0], t.subject.id);
    let predicate = neo4j_relation_to_relation(&t.predicate);
    let object = neo4j_label_to_object_id(&t.object.labels[0], t.object.id);

    if let (Some(s), Some(p), Some(o)) = (subject, predicate, object) {
        Some(Triple {
            subject: s,
            predicate: p,
            object: o
        })
    } else {
        None
    }
}

impl Db for GraphClient {

    fn get_entry(&self, id: &str) -> Result<Entry> {
        let result = self.exec(cypher_stmt!(
        "MATCH (e:Entry)<--(s:EntryState) WHERE e.id = {id}
         WITH max(s.version) as version
         MATCH (e:Entry)<--(s:EntryState)
         WHERE e.id = {id} AND s.version = version
         WITH e, s
         MATCH s<-[:BELONGS_TO]-(c:Category)
         WITH e, s, collect(DISTINCT c.id) as categories
         RETURN {
           id          : e.id,
           created     : s.created,
           version     : s.version,
           title       : s.title,
           description : s.description,
           lat         : s.lat,
           lng         : s.lng,
           street      : s.street,
           zip         : s.zip,
           city        : s.city,
           country     : s.country,
           email       : s.email,
           telephone   : s.telephone,
           homepage    : s.homepage,
           categories  : categories,
           license     : s.license
         } AS e
         ORDER BY e.created DESC", {"id" => &id})?)?;
        let r = result.rows().next().ok_or(RepoError::NotFound)?;
        let e = r.get::<Entry>("e")?;
        Ok(e)
    }

    fn get_user(&self, username: &str) -> Result<User> {
        let result = self.exec(cypher_stmt!(
        "MATCH u:User
         WHERE u.username = {username}",
        { "username" => username })?)?;
        let r = result.rows().next().ok_or(RepoError::NotFound)?;
        let u = r.get::<User>("u")?;
        Ok(u)
    }

    fn all_entries(&self) -> Result<Vec<Entry>> {
        let result = self.exec(
        "MATCH (e:Entry)<--(x:EntryState)
          WITH distinct e, max(x.created) as max
          MATCH e<--(s:EntryState)
          WHERE s.created = max
          WITH e,s
          MATCH e<-[:BELONGS_TO]-s
          OPTIONAL MATCH (s)<-[:BELONGS_TO]-(c:Category)
          WITH e, s, collect(DISTINCT c.id) as categories
          RETURN {
            id          : e.id,
            created     : s.created,
            version     : s.version,
            title       : s.title,
            description : s.description,
            lat         : s.lat,
            lng         : s.lng,
            street      : s.street,
            zip         : s.zip,
            city        : s.city,
            country     : s.country,
            email       : s.email,
            telephone   : s.telephone,
            homepage    : s.homepage,
            categories  : categories,
            license     : s.license
          } AS e
          ORDER BY e.created DESC")?;
        Ok(result.rows()
            .filter_map(|r| r.get::<Entry>("e").ok())
            .collect::<Vec<Entry>>())
    }

    fn create_entry(&mut self, e: &Entry) -> Result<()> {
        self.exec(cypher_stmt!(
        "MATCH (c:Category)
         WHERE c.id in {categories}
         WITH
            collect(DISTINCT c)    AS cats,
            collect(DISTINCT c.id) AS cat_ids,
            count(c)               AS cnt
         WHERE cnt > 0
         CREATE (e:Entry {id:{id}})
         MERGE e<-[:BELONGS_TO]-(s:EntryState {
           created : timestamp(),
           version : 1
         })
         SET s.title       = {title},
             s.description = {description},
             s.lat         = {lat},
             s.lng         = {lng},
             s.street      = {street},
             s.zip         = {zip},
             s.city        = {city},
             s.country     = {country},
             s.email       = {email},
             s.telephone   = {telephone},
             s.homepage    = {homepage},
             s.license     = {license}
         FOREACH (c IN cats |
           MERGE c-[:BELONGS_TO]->s
         )",
        {
            "id"          => &e.id,
            "title"       => &e.title,
            "description" => &e.description,
            "lat"         => &e.lat,
            "lng"         => &e.lng,
            "street"      => &e.street,
            "zip"         => &e.zip,
            "city"        => &e.city,
            "country"     => &e.country,
            "email"       => &e.email,
            "telephone"   => &e.telephone,
            "homepage"    => &e.homepage,
            "license"     => &e.license,
            "categories"  => &e.categories
        })?)?;
        Ok(())
    }

    fn create_tag(&mut self, t: &Tag) -> Result<()> {
        self.exec(cypher_stmt!("MERGE (t:Tag {id:{id}})",
        { "id" => &t.id })?)?;
        Ok(())
    }

    fn create_user(&mut self, u: &User) -> Result<()> {
        self.exec(cypher_stmt!(
        "MERGE (
           u:User {
             username:{username},
             password:{password},
             email:{email},
           }
        )",
        {
            "username" => &u.username,
            "password" => &u.password,
            "email"    => &u.email
        })?)?;
        Ok(())
    }

    fn create_rating(&mut self, r: &Rating) -> Result<()> {
        self.exec(cypher_stmt!(
        "MERGE (r:Rating {
             id         : {id},
             title      : {title},
             created    : {created},
             value      : {value},
             context    : {context},
             references : {references}
        })",
        {
            "id"         => &r.id,
            "title"      => &r.title,
            "created"    => &r.created,
            "value"      => &r.value,
            "context"    => &r.context,
            "references" => &r.references
        })?)?;
        Ok(())
    }

    fn create_comment(&mut self, c: &Comment) -> Result<()> {
        self.exec(cypher_stmt!(
        "MERGE (
           c:Comment {
             id      : {id},
             created : {created},
             text    : {text}
           }
        )",
        {
            "id"        => &c.id,
            "created"   => &c.created,
            "text"      => &c.text
        })?)?;
        Ok(())
    }

    fn create_triple(&mut self, t: &Triple) -> Result<()> {

        let predicate = neo4j_edge_name(t);
        let (subject_type, subject_id) = object_id_to_neo4j_label(&t.subject);
        let (object_type, object_id) = object_id_to_neo4j_label(&t.object);

        let stmt = format!(
           "MATCH (s:{s_type})
            WHERE s.id = \"{s_id}\"
            WITH s
            MATCH (o:{o_type})
            WHERE o.id = \"{o_id}\"
            WITH s,o
            MERGE (s)-[:{predicate}]->(o)",
                s_type = subject_type,
                s_id = subject_id,
                o_type = object_type,
                o_id = object_id,
                predicate = predicate
            );
        self.exec(stmt)?;
        Ok(())
    }

    fn update_entry(&mut self, e: &Entry) -> Result<()> {
        self.exec(cypher_stmt!(
        "MATCH (e:Entry)<--(s:EntryState) WHERE e.id = {id}
         WITH max(s.version) as v
         MATCH (e:Entry)<--(old:EntryState)
         WHERE e.id = {id} AND old.version = v AND old.version + 1 = {version}
         WITH e
         MATCH (c:Category)
         WHERE c.id in {categories}
         WITH e, collect(DISTINCT c) AS cats, count(c) AS cnt
         WHERE cnt > 0
         MERGE e<-[:BELONGS_TO]-(s:EntryState {
           created : timestamp(),
           version : {version}
         })
         SET s.title       = {title},
             s.description = {description},
             s.lat         = {lat},
             s.lng         = {lng},
             s.street      = {street},
             s.zip         = {zip},
             s.city        = {city},
             s.country     = {country},
             s.email       = {email},
             s.telephone   = {telephone},
             s.homepage    = {homepage},
             s.license     = {license}
         FOREACH (c IN cats |
           MERGE (c)-[:BELONGS_TO]->s
         )",
        {
            "id"          => &e.id,
            "version"     => &e.version,
            "title"       => &e.title,
            "description" => &e.description,
            "lat"         => &e.lat,
            "lng"         => &e.lng,
            "street"      => &e.street,
            "zip"         => &e.zip,
            "city"        => &e.city,
            "country"     => &e.country,
            "email"       => &e.email,
            "telephone"   => &e.telephone,
            "homepage"    => &e.homepage,
            "license"     => &e.license,
            "categories"  => &e.categories
        })?)?;
        Ok(())
    }

    fn all_categories(&self) -> Result<Vec<Category>> {
        let result = self.exec(
        "MATCH (c:Category)<--(s:CategoryState)
         RETURN {
           id      : c.id,
           version : s.version,
           created : s.created,
           name    : s.name
         } AS c")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Category>("c").ok())
            .collect::<Vec<Category>>())
    }

    fn all_triples(&self) -> Result<Vec<Triple>> {
        let result = self.exec(
        "MATCH (s)-[p]->(o)
         WHERE (s.id IS NOT NULL) AND (o.id IS NOT NULL)
         RETURN {
            subject: { id: s.id, labels: labels(s) },
            predicate: type(p),
            object: { id: o.id, labels: labels(o) }
         } AS t")?;
        Ok(result
            .rows()
            .filter_map(|r|{
                match r.get::<Neo4jTriple>("t") {
                    Err(err) => {
                         warn!("{}",err);
                         None
                    }
                    Ok(t) => from_neo_triple(t)
                }
            })
            .collect::<Vec<Triple>>())
    }

    fn all_tags(&self) -> Result<Vec<Tag>> {
        let result = self.exec(
        "MATCH (t:Tag) RETURN t")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Tag>("t").ok())
            .collect::<Vec<Tag>>())
    }

    fn all_ratings(&self) -> Result<Vec<Rating>> {
        let result = self.exec(
        "MATCH (r:Rating) RETURN r")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Rating>("r").ok())
            .collect::<Vec<Rating>>())
    }

    fn all_comments(&self) -> Result<Vec<Comment>> {
        let result = self.exec(
        "MATCH (c:Comment) RETURN c")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Comment>("c").ok())
            .collect::<Vec<Comment>>())
    }

    fn delete_triple(&mut self, t: &Triple) -> Result<()> {
        let predicate = neo4j_edge_name(t);
        let (subject_type, subject_id) = object_id_to_neo4j_label(&t.subject);
        let (object_type, object_id) = object_id_to_neo4j_label(&t.object);
        let stmt = format!(
           "MATCH (s:{s_type})-[p:{predicate}]->(o:{o_type})
            WHERE s.id = \"{s_id}\" AND o.id = \"{o_id}\"
            DELETE p",
                s_type = subject_type,
                s_id = subject_id,
                o_type = object_type,
                o_id = object_id,
                predicate = predicate
            );
        self.exec(stmt)?;
        Ok(())
    }
}
<commit_msg>fix(neo4j): fix user queries<commit_after>use business::db::Db;
use entities::*;
use rusted_cypher::GraphClient;
use business::error::RepoError;
use std::result;

type Result<T> = result::Result<T, RepoError>;

fn neo4j_edge_name<'a>(t: &'a Triple) -> &'a str {
    match t.predicate {
       Relation::IsTaggedWith    => "IS_TAGGED_WITH",
       Relation::IsRatedWith     => "IS_RATED_WITH",
       Relation::IsCommentedWith => "IS_COMMENTED_WITH",
       Relation::CreatedBy       => "CREATED_BY"
    }.into()
}

fn object_id_to_neo4j_label<'a>(id: &'a ObjectId) -> (&'a str,&'a str) {
    match *id {
        ObjectId::Entry(ref id) => ("Entry",id),
        ObjectId::Tag(ref id) => ("Tag",id),
        ObjectId::User(ref id) => ("User",id),
        ObjectId::Comment(ref id) => ("Comment",id),
        ObjectId::Rating(ref id) => ("Rating",id)
    }
}

fn neo4j_label_to_object_id(label: &str, id: String) -> Option<ObjectId> {
    match label {
        "Entry"     => Some(ObjectId::Entry(id)),
        "Tag"       => Some(ObjectId::Tag(id)),
        "User"      => Some(ObjectId::User(id)),
        "Comment"   => Some(ObjectId::Comment(id)),
        "Rating"    => Some(ObjectId::Rating(id)),
        _           => None,
    }
}

fn neo4j_relation_to_relation(rel: &str) -> Option<Relation> {
    match rel {
        "IS_TAGGED_WITH"    => Some(Relation::IsTaggedWith),
        "IS_RATED_WITH"     => Some(Relation::IsRatedWith),
        "IS_COMMENTED_WITH" => Some(Relation::IsCommentedWith),
        "CREATED_BY"        => Some(Relation::CreatedBy),
        _                   => None
    }
}

#[derive(Deserialize,Debug)]
struct Neo4jObj {
    labels: Vec<String>,
    id: String
}

#[derive(Deserialize,Debug)]
struct Neo4jTriple {
    subject: Neo4jObj,
    predicate: String,
    object: Neo4jObj
}

fn from_neo_triple(t: Neo4jTriple) -> Option<Triple> {
    if t.subject.labels.is_empty() || t.object.labels.is_empty() {
        warn!("No labels found in {:?}",t);
        return None;
    }
    let subject = neo4j_label_to_object_id(&t.subject.labels[0], t.subject.id);
    let predicate = neo4j_relation_to_relation(&t.predicate);
    let object = neo4j_label_to_object_id(&t.object.labels[0], t.object.id);

    if let (Some(s), Some(p), Some(o)) = (subject, predicate, object) {
        Some(Triple {
            subject: s,
            predicate: p,
            object: o
        })
    } else {
        None
    }
}

impl Db for GraphClient {

    fn get_entry(&self, id: &str) -> Result<Entry> {
        let result = self.exec(cypher_stmt!(
        "MATCH (e:Entry)<--(s:EntryState) WHERE e.id = {id}
         WITH max(s.version) as version
         MATCH (e:Entry)<--(s:EntryState)
         WHERE e.id = {id} AND s.version = version
         WITH e, s
         MATCH s<-[:BELONGS_TO]-(c:Category)
         WITH e, s, collect(DISTINCT c.id) as categories
         RETURN {
           id          : e.id,
           created     : s.created,
           version     : s.version,
           title       : s.title,
           description : s.description,
           lat         : s.lat,
           lng         : s.lng,
           street      : s.street,
           zip         : s.zip,
           city        : s.city,
           country     : s.country,
           email       : s.email,
           telephone   : s.telephone,
           homepage    : s.homepage,
           categories  : categories,
           license     : s.license
         } AS e
         ORDER BY e.created DESC", {"id" => &id})?)?;
        let r = result.rows().next().ok_or(RepoError::NotFound)?;
        let e = r.get::<Entry>("e")?;
        Ok(e)
    }

    fn get_user(&self, username: &str) -> Result<User> {
        let result = self.exec(cypher_stmt!(
        "MATCH (u:User)
         WHERE u.username = {username}
         RETURN u",
        { "username" => username })?)?;
        let r = result.rows().next().ok_or(RepoError::NotFound)?;
        let u = r.get::<User>("u")?;
        Ok(u)
    }

    fn all_entries(&self) -> Result<Vec<Entry>> {
        let result = self.exec(
        "MATCH (e:Entry)<--(x:EntryState)
          WITH distinct e, max(x.created) as max
          MATCH e<--(s:EntryState)
          WHERE s.created = max
          WITH e,s
          MATCH e<-[:BELONGS_TO]-s
          OPTIONAL MATCH (s)<-[:BELONGS_TO]-(c:Category)
          WITH e, s, collect(DISTINCT c.id) as categories
          RETURN {
            id          : e.id,
            created     : s.created,
            version     : s.version,
            title       : s.title,
            description : s.description,
            lat         : s.lat,
            lng         : s.lng,
            street      : s.street,
            zip         : s.zip,
            city        : s.city,
            country     : s.country,
            email       : s.email,
            telephone   : s.telephone,
            homepage    : s.homepage,
            categories  : categories,
            license     : s.license
          } AS e
          ORDER BY e.created DESC")?;
        Ok(result.rows()
            .filter_map(|r| r.get::<Entry>("e").ok())
            .collect::<Vec<Entry>>())
    }

    fn create_entry(&mut self, e: &Entry) -> Result<()> {
        self.exec(cypher_stmt!(
        "MATCH (c:Category)
         WHERE c.id in {categories}
         WITH
            collect(DISTINCT c)    AS cats,
            collect(DISTINCT c.id) AS cat_ids,
            count(c)               AS cnt
         WHERE cnt > 0
         CREATE (e:Entry {id:{id}})
         MERGE e<-[:BELONGS_TO]-(s:EntryState {
           created : timestamp(),
           version : 1
         })
         SET s.title       = {title},
             s.description = {description},
             s.lat         = {lat},
             s.lng         = {lng},
             s.street      = {street},
             s.zip         = {zip},
             s.city        = {city},
             s.country     = {country},
             s.email       = {email},
             s.telephone   = {telephone},
             s.homepage    = {homepage},
             s.license     = {license}
         FOREACH (c IN cats |
           MERGE c-[:BELONGS_TO]->s
         )",
        {
            "id"          => &e.id,
            "title"       => &e.title,
            "description" => &e.description,
            "lat"         => &e.lat,
            "lng"         => &e.lng,
            "street"      => &e.street,
            "zip"         => &e.zip,
            "city"        => &e.city,
            "country"     => &e.country,
            "email"       => &e.email,
            "telephone"   => &e.telephone,
            "homepage"    => &e.homepage,
            "license"     => &e.license,
            "categories"  => &e.categories
        })?)?;
        Ok(())
    }

    fn create_tag(&mut self, t: &Tag) -> Result<()> {
        self.exec(cypher_stmt!("MERGE (t:Tag {id:{id}})",
        { "id" => &t.id })?)?;
        Ok(())
    }

    fn create_user(&mut self, u: &User) -> Result<()> {
        self.exec(cypher_stmt!(
        "MERGE (
           u:User {
             username:{username},
             password:{password},
             email:{email}
           }
        )",
        {
            "username" => &u.username,
            "password" => &u.password,
            "email"    => &u.email
        })?)?;
        Ok(())
    }

    fn create_rating(&mut self, r: &Rating) -> Result<()> {
        self.exec(cypher_stmt!(
        "MERGE (r:Rating {
             id         : {id},
             title      : {title},
             created    : {created},
             value      : {value},
             context    : {context},
             references : {references}
        })",
        {
            "id"         => &r.id,
            "title"      => &r.title,
            "created"    => &r.created,
            "value"      => &r.value,
            "context"    => &r.context,
            "references" => &r.references
        })?)?;
        Ok(())
    }

    fn create_comment(&mut self, c: &Comment) -> Result<()> {
        self.exec(cypher_stmt!(
        "MERGE (
           c:Comment {
             id      : {id},
             created : {created},
             text    : {text}
           }
        )",
        {
            "id"        => &c.id,
            "created"   => &c.created,
            "text"      => &c.text
        })?)?;
        Ok(())
    }

    fn create_triple(&mut self, t: &Triple) -> Result<()> {

        let predicate = neo4j_edge_name(t);
        let (subject_type, subject_id) = object_id_to_neo4j_label(&t.subject);
        let (object_type, object_id) = object_id_to_neo4j_label(&t.object);

        let stmt = format!(
           "MATCH (s:{s_type})
            WHERE s.id = \"{s_id}\"
            WITH s
            MATCH (o:{o_type})
            WHERE o.id = \"{o_id}\"
            WITH s,o
            MERGE (s)-[:{predicate}]->(o)",
                s_type = subject_type,
                s_id = subject_id,
                o_type = object_type,
                o_id = object_id,
                predicate = predicate
            );
        self.exec(stmt)?;
        Ok(())
    }

    fn update_entry(&mut self, e: &Entry) -> Result<()> {
        self.exec(cypher_stmt!(
        "MATCH (e:Entry)<--(s:EntryState) WHERE e.id = {id}
         WITH max(s.version) as v
         MATCH (e:Entry)<--(old:EntryState)
         WHERE e.id = {id} AND old.version = v AND old.version + 1 = {version}
         WITH e
         MATCH (c:Category)
         WHERE c.id in {categories}
         WITH e, collect(DISTINCT c) AS cats, count(c) AS cnt
         WHERE cnt > 0
         MERGE e<-[:BELONGS_TO]-(s:EntryState {
           created : timestamp(),
           version : {version}
         })
         SET s.title       = {title},
             s.description = {description},
             s.lat         = {lat},
             s.lng         = {lng},
             s.street      = {street},
             s.zip         = {zip},
             s.city        = {city},
             s.country     = {country},
             s.email       = {email},
             s.telephone   = {telephone},
             s.homepage    = {homepage},
             s.license     = {license}
         FOREACH (c IN cats |
           MERGE (c)-[:BELONGS_TO]->s
         )",
        {
            "id"          => &e.id,
            "version"     => &e.version,
            "title"       => &e.title,
            "description" => &e.description,
            "lat"         => &e.lat,
            "lng"         => &e.lng,
            "street"      => &e.street,
            "zip"         => &e.zip,
            "city"        => &e.city,
            "country"     => &e.country,
            "email"       => &e.email,
            "telephone"   => &e.telephone,
            "homepage"    => &e.homepage,
            "license"     => &e.license,
            "categories"  => &e.categories
        })?)?;
        Ok(())
    }

    fn all_categories(&self) -> Result<Vec<Category>> {
        let result = self.exec(
        "MATCH (c:Category)<--(s:CategoryState)
         RETURN {
           id      : c.id,
           version : s.version,
           created : s.created,
           name    : s.name
         } AS c")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Category>("c").ok())
            .collect::<Vec<Category>>())
    }

    fn all_triples(&self) -> Result<Vec<Triple>> {
        let result = self.exec(
        "MATCH (s)-[p]->(o)
         WHERE (s.id IS NOT NULL) AND (o.id IS NOT NULL)
         RETURN {
            subject: { id: s.id, labels: labels(s) },
            predicate: type(p),
            object: { id: o.id, labels: labels(o) }
         } AS t")?;
        Ok(result
            .rows()
            .filter_map(|r|{
                match r.get::<Neo4jTriple>("t") {
                    Err(err) => {
                         warn!("{}",err);
                         None
                    }
                    Ok(t) => from_neo_triple(t)
                }
            })
            .collect::<Vec<Triple>>())
    }

    fn all_tags(&self) -> Result<Vec<Tag>> {
        let result = self.exec(
        "MATCH (t:Tag) RETURN t")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Tag>("t").ok())
            .collect::<Vec<Tag>>())
    }

    fn all_ratings(&self) -> Result<Vec<Rating>> {
        let result = self.exec(
        "MATCH (r:Rating) RETURN r")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Rating>("r").ok())
            .collect::<Vec<Rating>>())
    }

    fn all_comments(&self) -> Result<Vec<Comment>> {
        let result = self.exec(
        "MATCH (c:Comment) RETURN c")?;
        Ok(result
            .rows()
            .filter_map(|r| r.get::<Comment>("c").ok())
            .collect::<Vec<Comment>>())
    }

    fn delete_triple(&mut self, t: &Triple) -> Result<()> {
        let predicate = neo4j_edge_name(t);
        let (subject_type, subject_id) = object_id_to_neo4j_label(&t.subject);
        let (object_type, object_id) = object_id_to_neo4j_label(&t.object);
        let stmt = format!(
           "MATCH (s:{s_type})-[p:{predicate}]->(o:{o_type})
            WHERE s.id = \"{s_id}\" AND o.id = \"{o_id}\"
            DELETE p",
                s_type = subject_type,
                s_id = subject_id,
                o_type = object_type,
                o_id = object_id,
                predicate = predicate
            );
        self.exec(stmt)?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use chrono::offset::Utc;
use chrono::prelude::*;
use serde_derive::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashSet};

use crate::crypto;
use crate::error::Error;
use crate::metadata::{self, Metadata};
use crate::Result;

const SPEC_VERSION: &str = "1.0";

fn parse_datetime(ts: &str) -> Result<DateTime<Utc>> {
    Utc.datetime_from_str(ts, "%FT%TZ")
        .map_err(|e| Error::Encoding(format!("Can't parse DateTime: {:?}", e)))
}

fn format_datetime(ts: &DateTime<Utc>) -> String {
    format!(
        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
        ts.year(),
        ts.month(),
        ts.day(),
        ts.hour(),
        ts.minute(),
        ts.second()
    )
}

#[derive(Debug, Serialize, Deserialize)]
pub struct RootMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    consistent_snapshot: bool,
    expires: String,
    #[serde(deserialize_with = "deserialize_reject_duplicates::deserialize")]
    keys: BTreeMap<crypto::KeyId, crypto::PublicKey>,
    roles: RoleDefinitions,
}

impl RootMetadata {
    pub fn from(meta: &metadata::RootMetadata) -> Result<Self> {
        Ok(RootMetadata {
            typ: metadata::Role::Root,
            spec_version: SPEC_VERSION.to_string(),
            version: meta.version(),
            expires: format_datetime(&meta.expires()),
            consistent_snapshot: meta.consistent_snapshot(),
            keys: meta
                .keys()
                .iter()
                .map(|(id, key)| (id.clone(), key.clone()))
                .collect(),
            roles: RoleDefinitions {
                root: meta.root().clone(),
                snapshot: meta.snapshot().clone(),
                targets: meta.targets().clone(),
                timestamp: meta.timestamp().clone(),
            },
        })
    }

    pub fn try_into(self) -> Result<metadata::RootMetadata> {
        if self.typ != metadata::Role::Root {
            return Err(Error::Encoding(format!(
                "Attempted to decode root metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        // Ignore all keys with incorrect key IDs. We should give an error if the key ID is not
        // correct according to TUF spec. However, due to backward compatibility, we may receive
        // metadata with key IDs generated by TUF 0.9. We simply ignore those old keys.
        let keys_with_correct_key_id = self
            .keys
            .into_iter()
            .filter(|(key_id, pkey)| key_id == pkey.key_id())
            .collect();

        metadata::RootMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.consistent_snapshot,
            keys_with_correct_key_id,
            self.roles.root,
            self.roles.snapshot,
            self.roles.targets,
            self.roles.timestamp,
        )
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct RoleDefinitions {
    root: metadata::RoleDefinition,
    snapshot: metadata::RoleDefinition,
    targets: metadata::RoleDefinition,
    timestamp: metadata::RoleDefinition,
}

#[derive(Serialize, Deserialize)]
pub struct RoleDefinition {
    threshold: u32,
    #[serde(rename = "keyids")]
    key_ids: Vec<crypto::KeyId>,
}

impl RoleDefinition {
    pub fn from(role: &metadata::RoleDefinition) -> Result<Self> {
        let key_ids = role.key_ids().to_vec();

        Ok(RoleDefinition {
            threshold: role.threshold(),
            key_ids,
        })
    }

    pub fn try_into(self) -> Result<metadata::RoleDefinition> {
        let vec_len = self.key_ids.len();
        if vec_len < 1 {
            return Err(Error::Encoding(
                "Role defined with no assoiciated key IDs.".into(),
            ));
        }

        let mut seen = HashSet::new();
        let mut dupes = 0;
        for key_id in self.key_ids.iter() {
            if !seen.insert(key_id) {
                dupes += 1;
            }
        }

        if dupes != 0 {
            return Err(Error::Encoding(format!(
                "Found {} duplicate key IDs.",
                dupes
            )));
        }

        Ok(metadata::RoleDefinition::new(self.threshold, self.key_ids)?)
    }
}

#[derive(Serialize, Deserialize)]
pub struct TimestampMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    expires: String,
    meta: TimestampMeta,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct TimestampMeta {
    #[serde(rename = "snapshot.json")]
    snapshot: metadata::MetadataDescription,
}

impl TimestampMetadata {
    pub fn from(metadata: &metadata::TimestampMetadata) -> Result<Self> {
        Ok(TimestampMetadata {
            typ: metadata::Role::Timestamp,
            spec_version: SPEC_VERSION.to_string(),
            version: metadata.version(),
            expires: format_datetime(metadata.expires()),
            meta: TimestampMeta {
                snapshot: metadata.snapshot().clone(),
            },
        })
    }

    pub fn try_into(self) -> Result<metadata::TimestampMetadata> {
        if self.typ != metadata::Role::Timestamp {
            return Err(Error::Encoding(format!(
                "Attempted to decode timestamp metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        metadata::TimestampMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.meta.snapshot,
        )
    }
}

#[derive(Serialize, Deserialize)]
pub struct SnapshotMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    expires: String,
    #[serde(deserialize_with = "deserialize_reject_duplicates::deserialize")]
    meta: BTreeMap<String, metadata::MetadataDescription>,
}

impl SnapshotMetadata {
    pub fn from(metadata: &metadata::SnapshotMetadata) -> Result<Self> {
        Ok(SnapshotMetadata {
            typ: metadata::Role::Snapshot,
            spec_version: SPEC_VERSION.to_string(),
            version: metadata.version(),
            expires: format_datetime(&metadata.expires()),
            meta: metadata
                .meta()
                .iter()
                .map(|(p, d)| (format!("{}.json", p), d.clone()))
                .collect(),
        })
    }

    pub fn try_into(self) -> Result<metadata::SnapshotMetadata> {
        if self.typ != metadata::Role::Snapshot {
            return Err(Error::Encoding(format!(
                "Attempted to decode snapshot metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        metadata::SnapshotMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.meta
                .into_iter()
                .map(|(p, d)| {
                    if !p.ends_with(".json") {
                        return Err(Error::Encoding(format!(
                            "Metadata does not end with .json: {}",
                            p
                        )));
                    }

                    let s = p.split_at(p.len() - ".json".len()).0;
                    let p = metadata::MetadataPath::new(s)?;

                    Ok((p, d))
                })
                .collect::<Result<_>>()?,
        )
    }
}

#[derive(Serialize, Deserialize)]
pub struct TargetsMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    expires: String,
    targets: BTreeMap<metadata::VirtualTargetPath, metadata::TargetDescription>,
    #[serde(skip_serializing_if = "Option::is_none")]
    delegations: Option<metadata::Delegations>,
}

impl TargetsMetadata {
    pub fn from(metadata: &metadata::TargetsMetadata) -> Result<Self> {
        Ok(TargetsMetadata {
            typ: metadata::Role::Targets,
            spec_version: SPEC_VERSION.to_string(),
            version: metadata.version(),
            expires: format_datetime(&metadata.expires()),
            targets: metadata
                .targets()
                .iter()
                .map(|(p, d)| (p.clone(), d.clone()))
                .collect(),
            delegations: metadata.delegations().cloned(),
        })
    }

    pub fn try_into(self) -> Result<metadata::TargetsMetadata> {
        if self.typ != metadata::Role::Targets {
            return Err(Error::Encoding(format!(
                "Attempted to decode targets metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        metadata::TargetsMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.targets.into_iter().collect(),
            self.delegations,
        )
    }
}

#[derive(Serialize, Deserialize)]
pub struct PublicKey {
    keytype: crypto::KeyType,
    scheme: crypto::SignatureScheme,
    #[serde(skip_serializing_if = "Option::is_none")]
    keyid_hash_algorithms: Option<Vec<String>>,
    keyval: PublicKeyValue,
}

impl PublicKey {
    pub fn new(
        keytype: crypto::KeyType,
        scheme: crypto::SignatureScheme,
        keyid_hash_algorithms: Option<Vec<String>>,
        public_key: String,
    ) -> Self {
        PublicKey {
            keytype,
            scheme,
            keyid_hash_algorithms,
            keyval: PublicKeyValue { public: public_key },
        }
    }

    pub fn public_key(&self) -> &str {
        &self.keyval.public
    }

    pub fn scheme(&self) -> &crypto::SignatureScheme {
        &self.scheme
    }

    pub fn keytype(&self) -> &crypto::KeyType {
        &self.keytype
    }

    pub fn keyid_hash_algorithms(&self) -> &Option<Vec<String>> {
        &self.keyid_hash_algorithms
    }
}

#[derive(Serialize, Deserialize)]
pub struct PublicKeyValue {
    public: String,
}

#[derive(Serialize, Deserialize)]
pub struct Delegation {
    role: metadata::MetadataPath,
    terminating: bool,
    threshold: u32,
    #[serde(rename = "keyids")]
    key_ids: Vec<crypto::KeyId>,
    paths: Vec<metadata::VirtualTargetPath>,
}

impl Delegation {
    pub fn from(meta: &metadata::Delegation) -> Self {
        let mut paths = meta
            .paths()
            .iter()
            .cloned()
            .collect::<Vec<metadata::VirtualTargetPath>>();
        paths.sort();
        let mut key_ids = meta
            .key_ids()
            .iter()
            .cloned()
            .collect::<Vec<crypto::KeyId>>();
        key_ids.sort();

        Delegation {
            role: meta.role().clone(),
            terminating: meta.terminating(),
            threshold: meta.threshold(),
            key_ids,
            paths,
        }
    }

    pub fn try_into(self) -> Result<metadata::Delegation> {
        let paths = self
            .paths
            .iter()
            .cloned()
            .collect::<HashSet<metadata::VirtualTargetPath>>();
        if paths.len() != self.paths.len() {
            return Err(Error::Encoding("Non-unique delegation paths.".into()));
        }

        let key_ids = self
            .key_ids
            .iter()
            .cloned()
            .collect::<HashSet<crypto::KeyId>>();
        if key_ids.len() != self.key_ids.len() {
            return Err(Error::Encoding("Non-unique delegation key IDs.".into()));
        }

        metadata::Delegation::new(self.role, self.terminating, self.threshold, key_ids, paths)
    }
}

#[derive(Serialize, Deserialize)]
pub struct Delegations {
    #[serde(deserialize_with = "deserialize_reject_duplicates::deserialize")]
    keys: BTreeMap<crypto::KeyId, crypto::PublicKey>,
    roles: Vec<metadata::Delegation>,
}

impl Delegations {
    pub fn from(delegations: &metadata::Delegations) -> Delegations {
        Delegations {
            keys: delegations
                .keys()
                .iter()
                .map(|(id, key)| (id.clone(), key.clone()))
                .collect(),
            roles: delegations.roles().clone(),
        }
    }

    pub fn try_into(self) -> Result<metadata::Delegations> {
        metadata::Delegations::new(self.keys.into_iter().collect(), self.roles)
    }
}

#[derive(Serialize, Deserialize)]
pub struct TargetDescription {
    length: u64,
    hashes: BTreeMap<crypto::HashAlgorithm, crypto::HashValue>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    custom: Option<BTreeMap<String, serde_json::Value>>,
}

impl TargetDescription {
    pub fn from(description: &metadata::TargetDescription) -> TargetDescription {
        TargetDescription {
            length: description.length(),
            hashes: description
                .hashes()
                .iter()
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect(),
            custom: description
                .custom()
                .map(|custom| custom.iter().map(|(k, v)| (k.clone(), v.clone())).collect()),
        }
    }

    pub fn try_into(self) -> Result<metadata::TargetDescription> {
        metadata::TargetDescription::new(
            self.length,
            self.hashes.into_iter().collect(),
            self.custom.map(|custom| custom.into_iter().collect()),
        )
    }
}

#[derive(Deserialize)]
pub struct MetadataDescription {
    version: u32,
    length: usize,
    hashes: BTreeMap<crypto::HashAlgorithm, crypto::HashValue>,
}

impl MetadataDescription {
    pub fn try_into(self) -> Result<metadata::MetadataDescription> {
        metadata::MetadataDescription::new(
            self.version,
            self.length,
            self.hashes.into_iter().collect(),
        )
    }
}

/// Custom deserialize to reject duplicate keys.
mod deserialize_reject_duplicates {
    use serde::de::{Deserialize, Deserializer, Error, MapAccess, Visitor};
    use std::collections::BTreeMap;
    use std::fmt;
    use std::marker::PhantomData;
    use std::result::Result;

    pub fn deserialize<'de, K, V, D>(deserializer: D) -> Result<BTreeMap<K, V>, D::Error>
    where
        K: Deserialize<'de> + Ord,
        V: Deserialize<'de>,
        D: Deserializer<'de>,
    {
        struct BTreeVisitor<K, V> {
            marker: PhantomData<(K, V)>,
        };

        impl<'de, K, V> Visitor<'de> for BTreeVisitor<K, V>
        where
            K: Deserialize<'de> + Ord,
            V: Deserialize<'de>,
        {
            type Value = BTreeMap<K, V>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("map")
            }

            fn visit_map<M>(self, mut access: M) -> std::result::Result<Self::Value, M::Error>
            where
                M: MapAccess<'de>,
            {
                let mut map = BTreeMap::new();
                while let Some((key, value)) = access.next_entry()? {
                    if map.insert(key, value).is_some() {
                        return Err(M::Error::custom("Cannot have duplicate keys"));
                    }
                }
                Ok(map)
            }
        }

        deserializer.deserialize_map(BTreeVisitor {
            marker: PhantomData,
        })
    }
}
<commit_msg>Remove redundant semicolon<commit_after>use chrono::offset::Utc;
use chrono::prelude::*;
use serde_derive::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashSet};

use crate::crypto;
use crate::error::Error;
use crate::metadata::{self, Metadata};
use crate::Result;

const SPEC_VERSION: &str = "1.0";

fn parse_datetime(ts: &str) -> Result<DateTime<Utc>> {
    Utc.datetime_from_str(ts, "%FT%TZ")
        .map_err(|e| Error::Encoding(format!("Can't parse DateTime: {:?}", e)))
}

fn format_datetime(ts: &DateTime<Utc>) -> String {
    format!(
        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
        ts.year(),
        ts.month(),
        ts.day(),
        ts.hour(),
        ts.minute(),
        ts.second()
    )
}

#[derive(Debug, Serialize, Deserialize)]
pub struct RootMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    consistent_snapshot: bool,
    expires: String,
    #[serde(deserialize_with = "deserialize_reject_duplicates::deserialize")]
    keys: BTreeMap<crypto::KeyId, crypto::PublicKey>,
    roles: RoleDefinitions,
}

impl RootMetadata {
    pub fn from(meta: &metadata::RootMetadata) -> Result<Self> {
        Ok(RootMetadata {
            typ: metadata::Role::Root,
            spec_version: SPEC_VERSION.to_string(),
            version: meta.version(),
            expires: format_datetime(&meta.expires()),
            consistent_snapshot: meta.consistent_snapshot(),
            keys: meta
                .keys()
                .iter()
                .map(|(id, key)| (id.clone(), key.clone()))
                .collect(),
            roles: RoleDefinitions {
                root: meta.root().clone(),
                snapshot: meta.snapshot().clone(),
                targets: meta.targets().clone(),
                timestamp: meta.timestamp().clone(),
            },
        })
    }

    pub fn try_into(self) -> Result<metadata::RootMetadata> {
        if self.typ != metadata::Role::Root {
            return Err(Error::Encoding(format!(
                "Attempted to decode root metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        // Ignore all keys with incorrect key IDs. We should give an error if the key ID is not
        // correct according to TUF spec. However, due to backward compatibility, we may receive
        // metadata with key IDs generated by TUF 0.9. We simply ignore those old keys.
        let keys_with_correct_key_id = self
            .keys
            .into_iter()
            .filter(|(key_id, pkey)| key_id == pkey.key_id())
            .collect();

        metadata::RootMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.consistent_snapshot,
            keys_with_correct_key_id,
            self.roles.root,
            self.roles.snapshot,
            self.roles.targets,
            self.roles.timestamp,
        )
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct RoleDefinitions {
    root: metadata::RoleDefinition,
    snapshot: metadata::RoleDefinition,
    targets: metadata::RoleDefinition,
    timestamp: metadata::RoleDefinition,
}

#[derive(Serialize, Deserialize)]
pub struct RoleDefinition {
    threshold: u32,
    #[serde(rename = "keyids")]
    key_ids: Vec<crypto::KeyId>,
}

impl RoleDefinition {
    pub fn from(role: &metadata::RoleDefinition) -> Result<Self> {
        let key_ids = role.key_ids().to_vec();

        Ok(RoleDefinition {
            threshold: role.threshold(),
            key_ids,
        })
    }

    pub fn try_into(self) -> Result<metadata::RoleDefinition> {
        let vec_len = self.key_ids.len();
        if vec_len < 1 {
            return Err(Error::Encoding(
                "Role defined with no assoiciated key IDs.".into(),
            ));
        }

        let mut seen = HashSet::new();
        let mut dupes = 0;
        for key_id in self.key_ids.iter() {
            if !seen.insert(key_id) {
                dupes += 1;
            }
        }

        if dupes != 0 {
            return Err(Error::Encoding(format!(
                "Found {} duplicate key IDs.",
                dupes
            )));
        }

        Ok(metadata::RoleDefinition::new(self.threshold, self.key_ids)?)
    }
}

#[derive(Serialize, Deserialize)]
pub struct TimestampMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    expires: String,
    meta: TimestampMeta,
}

#[derive(Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct TimestampMeta {
    #[serde(rename = "snapshot.json")]
    snapshot: metadata::MetadataDescription,
}

impl TimestampMetadata {
    pub fn from(metadata: &metadata::TimestampMetadata) -> Result<Self> {
        Ok(TimestampMetadata {
            typ: metadata::Role::Timestamp,
            spec_version: SPEC_VERSION.to_string(),
            version: metadata.version(),
            expires: format_datetime(metadata.expires()),
            meta: TimestampMeta {
                snapshot: metadata.snapshot().clone(),
            },
        })
    }

    pub fn try_into(self) -> Result<metadata::TimestampMetadata> {
        if self.typ != metadata::Role::Timestamp {
            return Err(Error::Encoding(format!(
                "Attempted to decode timestamp metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        metadata::TimestampMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.meta.snapshot,
        )
    }
}

#[derive(Serialize, Deserialize)]
pub struct SnapshotMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    expires: String,
    #[serde(deserialize_with = "deserialize_reject_duplicates::deserialize")]
    meta: BTreeMap<String, metadata::MetadataDescription>,
}

impl SnapshotMetadata {
    pub fn from(metadata: &metadata::SnapshotMetadata) -> Result<Self> {
        Ok(SnapshotMetadata {
            typ: metadata::Role::Snapshot,
            spec_version: SPEC_VERSION.to_string(),
            version: metadata.version(),
            expires: format_datetime(&metadata.expires()),
            meta: metadata
                .meta()
                .iter()
                .map(|(p, d)| (format!("{}.json", p), d.clone()))
                .collect(),
        })
    }

    pub fn try_into(self) -> Result<metadata::SnapshotMetadata> {
        if self.typ != metadata::Role::Snapshot {
            return Err(Error::Encoding(format!(
                "Attempted to decode snapshot metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        metadata::SnapshotMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.meta
                .into_iter()
                .map(|(p, d)| {
                    if !p.ends_with(".json") {
                        return Err(Error::Encoding(format!(
                            "Metadata does not end with .json: {}",
                            p
                        )));
                    }

                    let s = p.split_at(p.len() - ".json".len()).0;
                    let p = metadata::MetadataPath::new(s)?;

                    Ok((p, d))
                })
                .collect::<Result<_>>()?,
        )
    }
}

#[derive(Serialize, Deserialize)]
pub struct TargetsMetadata {
    #[serde(rename = "_type")]
    typ: metadata::Role,
    spec_version: String,
    version: u32,
    expires: String,
    targets: BTreeMap<metadata::VirtualTargetPath, metadata::TargetDescription>,
    #[serde(skip_serializing_if = "Option::is_none")]
    delegations: Option<metadata::Delegations>,
}

impl TargetsMetadata {
    pub fn from(metadata: &metadata::TargetsMetadata) -> Result<Self> {
        Ok(TargetsMetadata {
            typ: metadata::Role::Targets,
            spec_version: SPEC_VERSION.to_string(),
            version: metadata.version(),
            expires: format_datetime(&metadata.expires()),
            targets: metadata
                .targets()
                .iter()
                .map(|(p, d)| (p.clone(), d.clone()))
                .collect(),
            delegations: metadata.delegations().cloned(),
        })
    }

    pub fn try_into(self) -> Result<metadata::TargetsMetadata> {
        if self.typ != metadata::Role::Targets {
            return Err(Error::Encoding(format!(
                "Attempted to decode targets metdata labeled as {:?}",
                self.typ
            )));
        }

        if self.spec_version != SPEC_VERSION {
            return Err(Error::Encoding(format!(
                "Unknown spec version {}",
                self.spec_version
            )));
        }

        metadata::TargetsMetadata::new(
            self.version,
            parse_datetime(&self.expires)?,
            self.targets.into_iter().collect(),
            self.delegations,
        )
    }
}

#[derive(Serialize, Deserialize)]
pub struct PublicKey {
    keytype: crypto::KeyType,
    scheme: crypto::SignatureScheme,
    #[serde(skip_serializing_if = "Option::is_none")]
    keyid_hash_algorithms: Option<Vec<String>>,
    keyval: PublicKeyValue,
}

impl PublicKey {
    pub fn new(
        keytype: crypto::KeyType,
        scheme: crypto::SignatureScheme,
        keyid_hash_algorithms: Option<Vec<String>>,
        public_key: String,
    ) -> Self {
        PublicKey {
            keytype,
            scheme,
            keyid_hash_algorithms,
            keyval: PublicKeyValue { public: public_key },
        }
    }

    pub fn public_key(&self) -> &str {
        &self.keyval.public
    }

    pub fn scheme(&self) -> &crypto::SignatureScheme {
        &self.scheme
    }

    pub fn keytype(&self) -> &crypto::KeyType {
        &self.keytype
    }

    pub fn keyid_hash_algorithms(&self) -> &Option<Vec<String>> {
        &self.keyid_hash_algorithms
    }
}

#[derive(Serialize, Deserialize)]
pub struct PublicKeyValue {
    public: String,
}

#[derive(Serialize, Deserialize)]
pub struct Delegation {
    role: metadata::MetadataPath,
    terminating: bool,
    threshold: u32,
    #[serde(rename = "keyids")]
    key_ids: Vec<crypto::KeyId>,
    paths: Vec<metadata::VirtualTargetPath>,
}

impl Delegation {
    pub fn from(meta: &metadata::Delegation) -> Self {
        let mut paths = meta
            .paths()
            .iter()
            .cloned()
            .collect::<Vec<metadata::VirtualTargetPath>>();
        paths.sort();
        let mut key_ids = meta
            .key_ids()
            .iter()
            .cloned()
            .collect::<Vec<crypto::KeyId>>();
        key_ids.sort();

        Delegation {
            role: meta.role().clone(),
            terminating: meta.terminating(),
            threshold: meta.threshold(),
            key_ids,
            paths,
        }
    }

    pub fn try_into(self) -> Result<metadata::Delegation> {
        let paths = self
            .paths
            .iter()
            .cloned()
            .collect::<HashSet<metadata::VirtualTargetPath>>();
        if paths.len() != self.paths.len() {
            return Err(Error::Encoding("Non-unique delegation paths.".into()));
        }

        let key_ids = self
            .key_ids
            .iter()
            .cloned()
            .collect::<HashSet<crypto::KeyId>>();
        if key_ids.len() != self.key_ids.len() {
            return Err(Error::Encoding("Non-unique delegation key IDs.".into()));
        }

        metadata::Delegation::new(self.role, self.terminating, self.threshold, key_ids, paths)
    }
}

#[derive(Serialize, Deserialize)]
pub struct Delegations {
    #[serde(deserialize_with = "deserialize_reject_duplicates::deserialize")]
    keys: BTreeMap<crypto::KeyId, crypto::PublicKey>,
    roles: Vec<metadata::Delegation>,
}

impl Delegations {
    pub fn from(delegations: &metadata::Delegations) -> Delegations {
        Delegations {
            keys: delegations
                .keys()
                .iter()
                .map(|(id, key)| (id.clone(), key.clone()))
                .collect(),
            roles: delegations.roles().clone(),
        }
    }

    pub fn try_into(self) -> Result<metadata::Delegations> {
        metadata::Delegations::new(self.keys.into_iter().collect(), self.roles)
    }
}

#[derive(Serialize, Deserialize)]
pub struct TargetDescription {
    length: u64,
    hashes: BTreeMap<crypto::HashAlgorithm, crypto::HashValue>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    custom: Option<BTreeMap<String, serde_json::Value>>,
}

impl TargetDescription {
    pub fn from(description: &metadata::TargetDescription) -> TargetDescription {
        TargetDescription {
            length: description.length(),
            hashes: description
                .hashes()
                .iter()
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect(),
            custom: description
                .custom()
                .map(|custom| custom.iter().map(|(k, v)| (k.clone(), v.clone())).collect()),
        }
    }

    pub fn try_into(self) -> Result<metadata::TargetDescription> {
        metadata::TargetDescription::new(
            self.length,
            self.hashes.into_iter().collect(),
            self.custom.map(|custom| custom.into_iter().collect()),
        )
    }
}

#[derive(Deserialize)]
pub struct MetadataDescription {
    version: u32,
    length: usize,
    hashes: BTreeMap<crypto::HashAlgorithm, crypto::HashValue>,
}

impl MetadataDescription {
    pub fn try_into(self) -> Result<metadata::MetadataDescription> {
        metadata::MetadataDescription::new(
            self.version,
            self.length,
            self.hashes.into_iter().collect(),
        )
    }
}

/// Custom deserialize to reject duplicate keys.
mod deserialize_reject_duplicates {
    use serde::de::{Deserialize, Deserializer, Error, MapAccess, Visitor};
    use std::collections::BTreeMap;
    use std::fmt;
    use std::marker::PhantomData;
    use std::result::Result;

    pub fn deserialize<'de, K, V, D>(deserializer: D) -> Result<BTreeMap<K, V>, D::Error>
    where
        K: Deserialize<'de> + Ord,
        V: Deserialize<'de>,
        D: Deserializer<'de>,
    {
        struct BTreeVisitor<K, V> {
            marker: PhantomData<(K, V)>,
        }

        impl<'de, K, V> Visitor<'de> for BTreeVisitor<K, V>
        where
            K: Deserialize<'de> + Ord,
            V: Deserialize<'de>,
        {
            type Value = BTreeMap<K, V>;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("map")
            }

            fn visit_map<M>(self, mut access: M) -> std::result::Result<Self::Value, M::Error>
            where
                M: MapAccess<'de>,
            {
                let mut map = BTreeMap::new();
                while let Some((key, value)) = access.next_entry()? {
                    if map.insert(key, value).is_some() {
                        return Err(M::Error::custom("Cannot have duplicate keys"));
                    }
                }
                Ok(map)
            }
        }

        deserializer.deserialize_map(BTreeVisitor {
            marker: PhantomData,
        })
    }
}
<|endoftext|>