<commit_before>use std::collections::BTreeMap;
use rustc_serialize::json::{ToJson, Json};
use postgres;
use uuid::Uuid;
use std::fmt;
use chrono::{NaiveDateTime, UTC, DateTime};

use apple_music;
use spotify;
use error::Error;
use super::{conn, Model};
use model::provider::Provider;
use model::state::State;
use model::enclosure::Enclosure;
use model::artist::Artist;

static PROPS: [&'static str; 14]  = ["id",
                                     "provider",
                                     "identifier",
                                     "owner_id",
                                     "owner_name",
                                     "url",
                                     "title",
                                     "description",
                                     "thumbnail_url",
                                     "artwork_url",
                                     "published_at",
                                     "created_at",
                                     "updated_at",
                                     "state"];

#[derive(Debug, Clone)]
pub struct Album {
    pub id:            Uuid,
    pub provider:      Provider,
    pub identifier:    String,
    pub owner_id:      Option<String>,
    pub owner_name:    Option<String>,
    pub url:           String,
    pub title:         String,
    pub description:   Option<String>,
    pub thumbnail_url: Option<String>,
    pub artwork_url:   Option<String>,
    pub published_at:  NaiveDateTime,
    pub created_at:    NaiveDateTime,
    pub updated_at:    NaiveDateTime,
    pub state:         State,
    pub artists:       Option<Vec<Artist>>,
}

impl PartialEq for Album {
    fn eq(&self, p: &Album) -> bool {
        return self.identifier == p.identifier && self.provider == p.provider
    }
}

impl ToJson for Album {
    fn to_json(&self) -> Json {
        let published_at = DateTime::<UTC>::from_utc(self.published_at, UTC);
        let created_at   = DateTime::<UTC>::from_utc(self.created_at  , UTC);
        let updated_at   = DateTime::<UTC>::from_utc(self.updated_at  , UTC);
        let mut d = BTreeMap::new();
        d.insert("id".to_string()           , self.id.to_string().to_json());
        d.insert("provider".to_string()     , self.provider.to_json());
        d.insert("identifier".to_string()   , self.identifier.to_json());
        d.insert("owner_id".to_string()     , self.owner_id.to_json());
        d.insert("owner_name".to_string()   , self.owner_name.to_json());
        d.insert("url".to_string()          , self.url.to_json());
        d.insert("title".to_string()        , self.title.to_json());
        d.insert("description".to_string()  , self.description.to_json());
        d.insert("thumbnail_url".to_string(), self.thumbnail_url.to_json());
        d.insert("artwork_url".to_string()  , self.artwork_url.to_json());
        d.insert("published_at".to_string() , published_at.to_rfc3339().to_json());
        d.insert("created_at".to_string()   , created_at.to_rfc3339().to_json());
        d.insert("updated_at".to_string()   , updated_at.to_rfc3339().to_json());
        d.insert("state".to_string()        , self.state.to_json());
        d.insert("artists".to_string()      , self.artists.to_json());
        Json::Object(d)
    }
}

impl fmt::Display for Album {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.provider, self.identifier)
    }
}

impl Model for Album {
    fn table_name() -> String { "albums".to_string() }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }

    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Album> {
        let mut albums = Vec::new();
        for row in rows.iter() {
            let album = Album {
                id:            row.get(0),
                provider:      Provider::new(row.get(1)),
                identifier:    row.get(2),
                owner_id:      row.get(3),
                owner_name:    row.get(4),
                url:           row.get(5),
                title:         row.get(6),
                description:   row.get(7),
                thumbnail_url: row.get(8),
                artwork_url:   row.get(9),
                published_at:  row.get(10),
                created_at:    row.get(11),
                updated_at:    row.get(12),
                state:         State::new(row.get(13)),
                artists:       None,
            };
            albums.push(album)
        }
        albums
    }

    fn create(&self) -> Result<Album, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO albums (provider, identifier, url, title)
                                      VALUES ($1, $2, $3, $4) RETURNING id"));
        let rows = try!(stmt.query(&[&self.provider.to_string(), &self.identifier, &self.url, &self.title]));
        let mut album = self.clone();
        for row in rows.iter() {
            album.id = row.get(0);
        }
        Ok(album)
    }

    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = UTC::now().naive_utc();
        let conn = try!(conn());
        let stmt = try!(conn.prepare("UPDATE albums SET
                                      provider      = $2,
                                      identifier    = $3,
                                      owner_id      = $4,
                                      owner_name    = $5,
                                      url           = $6,
                                      title         = $7,
                                      description   = $8,
                                      thumbnail_url = $9,
                                      artwork_url   = $10,
                                      published_at  = $11,
                                      created_at    = $12,
                                      updated_at    = $13,
                                      state         = $14
                                      WHERE id = $1"));
        let result = stmt.query(&[&self.id,
                                  &self.provider.to_string(),
                                  &self.identifier,
                                  &self.owner_id,
                                  &self.owner_name,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.thumbnail_url,
                                  &self.artwork_url,
                                  &self.published_at,
                                  &self.created_at,
                                  &self.updated_at,
                                  &self.state.to_string(),
        ]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected),
        }
    }
}

impl Enclosure for Album {
    fn new(provider: Provider, identifier: String) -> Album {
        Album {
            id:            Uuid::new_v4(),
            provider:      provider,
            identifier:    identifier,
            owner_id:      None,
            owner_name:    None,
            url:           "".to_string(),
            title:         "".to_string(),
            description:   None,
            thumbnail_url: None,
            artwork_url:   None,
            published_at:  UTC::now().naive_utc(),
            created_at:    UTC::now().naive_utc(),
            updated_at:    UTC::now().naive_utc(),
            state:         State::Alive,
            artists:       None,
        }
    }
    fn find_by_entry_id(entry_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums p LEFT JOIN album_entries pe
                        ON p.id = pe.album_id
                        WHERE pe.entry_id = $1
                        ORDER BY p.published_at DESC",
                     Album::props_str("p."))).unwrap();
        let rows = stmt.query(&[&entry_id]).unwrap();
        Album::rows_to_items(rows)
    }
}

impl Album {
    fn add_artist(&mut self, artist: Artist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO track_artists (track_id, artist_id) VALUES ($1, $2)"));
        try!(stmt.query(&[&self.id, &artist.id]));
        match self.artists {
            Some(ref mut artists) => artists.push(artist),
            None                  => self.artists = Some(vec![artist]),

        }
        Ok(())
    }

    pub fn from_sp_album(album: &spotify::Album) -> Album {
        Album::new(Provider::Spotify, (*album).id.to_string())
            .update_with_sp_album(album)
            .clone()
    }

    pub fn from_am_album(album: &apple_music::Album) -> Album {
        Album::new(Provider::AppleMusic, (*album).id.to_string())
            .update_with_am_album(album)
            .clone()
    }

    pub fn update_with_sp_album(&mut self, album: &spotify::Album) -> &mut Album {
        self.provider       = Provider::Spotify;
        self.identifier     = album.id.to_string();
        if album.artists.len() > 0 {
            self.owner_id   = Some(album.artists[0].id.clone());
            self.owner_name = Some(album.artists[0].name.clone());
        }
        self.url            = album.uri.clone();
        self.title          = album.name.clone();
        self.description    = None;
        self.state          = State::Alive;
        self.published_at   = UTC::now().naive_utc();
        if album.images.len() > 0 {
            self.artwork_url   = Some(album.images[0].url.clone());
            self.thumbnail_url = Some(album.images[0].url.clone());
        }
        if album.images.len() > 1 {
            self.thumbnail_url = Some(album.images[1].url.clone());
        }
        let artists = album.artists
            .iter()
            .map(|artist| {
                Artist::find_or_create(self.provider, artist.id.to_string())
                    .map(|mut a| a.update_with_sp_artist(&artist).clone())
            })
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<Artist>>();
        for mut a in artists {
            let _ = a.save();
            let _ = self.add_artist(a);
        }
        self
    }

    pub fn update_with_am_album(&mut self, album: &apple_music::Album) -> &mut Album {
        self.provider      = Provider::AppleMusic;
        self.identifier    = album.id.to_string();
        self.owner_id      = Some(album.album_artist.to_string());
        self.owner_name    = Some(album.album_artist.to_string());
        self.url           = album.music_url.to_string();
        self.title         = album.title.to_string();
        self.description   = None;
        self.thumbnail_url = Some(album.artwork_url.to_string());
        self.artwork_url   = Some(album.artwork_url.to_string());
        self.state         = State::Alive;
        if let Ok(mut artist) = Artist::find_or_create(self.provider,
                                                       album.album_artist.to_string()) {
            artist.name  = album.album_artist.to_string();
            let _ = artist.save();
            let _ = self.add_artist(artist);
        }
        self
    }

    pub fn disable(&mut self) -> &mut Album {
        self.state = State::Dead;
        self
    }

    pub fn delete(&self) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("DELETE FROM albums WHERE id=$1"));
        let result = stmt.query(&[&self.id]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected)
        }
    }
}

#[cfg(test)]
mod test {
    use model::enclosure::Enclosure;
    use model::Model;
    use super::Album;
    use Provider;
    #[test]
    fn test_new() {
        let album = Album::new(Provider::YouTube,
                               "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string());
        assert_eq!(album.provider, Provider::YouTube);
        assert_eq!(&album.identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX")
    }
    #[test]
    fn test_find_of_create() {
        let identifier = "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string();
        let result     = Album::find_or_create(Provider::YouTube, identifier);
        assert!(result.is_ok());
    }
    #[test]
    fn test_delete() {
        let identifier = "test_delete".to_string();
        let album      = Album::find_or_create(Provider::YouTube, identifier).unwrap();
        let result     = album.delete();
        assert!(result.is_ok());
    }
    #[test]
    fn test_save() {
        let id = "test_save";
        let mut album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        album.title   = "title".to_string();
        let result    = album.save();
        assert!(result.is_ok());
        let album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        assert_eq!(&album.title, "title");
    }
}
<commit_msg>Add tracks to Album model<commit_after>use std::collections::BTreeMap;
use rustc_serialize::json::{ToJson, Json};
use postgres;
use uuid::Uuid;
use std::fmt;
use chrono::{NaiveDateTime, UTC, DateTime};

use apple_music;
use spotify;
use error::Error;
use super::{conn, Model};
use model::provider::Provider;
use model::state::State;
use model::enclosure::Enclosure;
use model::track::Track;
use model::artist::Artist;

static PROPS: [&'static str; 14]  = ["id",
                                     "provider",
                                     "identifier",
                                     "owner_id",
                                     "owner_name",
                                     "url",
                                     "title",
                                     "description",
                                     "thumbnail_url",
                                     "artwork_url",
                                     "published_at",
                                     "created_at",
                                     "updated_at",
                                     "state"];

#[derive(Debug, Clone)]
pub struct Album {
    pub id:            Uuid,
    pub provider:      Provider,
    pub identifier:    String,
    pub owner_id:      Option<String>,
    pub owner_name:    Option<String>,
    pub url:           String,
    pub title:         String,
    pub description:   Option<String>,
    pub thumbnail_url: Option<String>,
    pub artwork_url:   Option<String>,
    pub published_at:  NaiveDateTime,
    pub created_at:    NaiveDateTime,
    pub updated_at:    NaiveDateTime,
    pub state:         State,
    pub tracks:        Vec<Track>,
    pub artists:       Option<Vec<Artist>>,
}

impl PartialEq for Album {
    fn eq(&self, p: &Album) -> bool {
        return self.identifier == p.identifier && self.provider == p.provider
    }
}

impl ToJson for Album {
    fn to_json(&self) -> Json {
        let published_at = DateTime::<UTC>::from_utc(self.published_at, UTC);
        let created_at   = DateTime::<UTC>::from_utc(self.created_at  , UTC);
        let updated_at   = DateTime::<UTC>::from_utc(self.updated_at  , UTC);
        let mut d = BTreeMap::new();
        d.insert("id".to_string()           , self.id.to_string().to_json());
        d.insert("provider".to_string()     , self.provider.to_json());
        d.insert("identifier".to_string()   , self.identifier.to_json());
        d.insert("owner_id".to_string()     , self.owner_id.to_json());
        d.insert("owner_name".to_string()   , self.owner_name.to_json());
        d.insert("url".to_string()          , self.url.to_json());
        d.insert("title".to_string()        , self.title.to_json());
        d.insert("description".to_string()  , self.description.to_json());
        d.insert("thumbnail_url".to_string(), self.thumbnail_url.to_json());
        d.insert("artwork_url".to_string()  , self.artwork_url.to_json());
        d.insert("published_at".to_string() , published_at.to_rfc3339().to_json());
        d.insert("created_at".to_string()   , created_at.to_rfc3339().to_json());
        d.insert("updated_at".to_string()   , updated_at.to_rfc3339().to_json());
        d.insert("state".to_string()        , self.state.to_json());
        d.insert("artists".to_string()      , self.artists.to_json());
        Json::Object(d)
    }
}

impl fmt::Display for Album {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.provider, self.identifier)
    }
}

impl Model for Album {
    fn table_name() -> String { "albums".to_string() }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }

    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Album> {
        let mut albums = Vec::new();
        for row in rows.iter() {
            let album = Album {
                id:            row.get(0),
                provider:      Provider::new(row.get(1)),
                identifier:    row.get(2),
                owner_id:      row.get(3),
                owner_name:    row.get(4),
                url:           row.get(5),
                title:         row.get(6),
                description:   row.get(7),
                thumbnail_url: row.get(8),
                artwork_url:   row.get(9),
                published_at:  row.get(10),
                created_at:    row.get(11),
                updated_at:    row.get(12),
                state:         State::new(row.get(13)),
                tracks:        vec![],
                artists:       None,
            };
            albums.push(album)
        }
        albums
    }

    fn create(&self) -> Result<Album, Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO albums (provider, identifier, url, title)
                                      VALUES ($1, $2, $3, $4) RETURNING id"));
        let rows = try!(stmt.query(&[&self.provider.to_string(), &self.identifier, &self.url, &self.title]));
        let mut album = self.clone();
        for row in rows.iter() {
            album.id = row.get(0);
        }
        Ok(album)
    }

    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = UTC::now().naive_utc();
        let conn = try!(conn());
        let stmt = try!(conn.prepare("UPDATE albums SET
                                      provider      = $2,
                                      identifier    = $3,
                                      owner_id      = $4,
                                      owner_name    = $5,
                                      url           = $6,
                                      title         = $7,
                                      description   = $8,
                                      thumbnail_url = $9,
                                      artwork_url   = $10,
                                      published_at  = $11,
                                      created_at    = $12,
                                      updated_at    = $13,
                                      state         = $14
                                      WHERE id = $1"));
        let result = stmt.query(&[&self.id,
                                  &self.provider.to_string(),
                                  &self.identifier,
                                  &self.owner_id,
                                  &self.owner_name,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.thumbnail_url,
                                  &self.artwork_url,
                                  &self.published_at,
                                  &self.created_at,
                                  &self.updated_at,
                                  &self.state.to_string(),
        ]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected),
        }
    }
}

impl Enclosure for Album {
    fn new(provider: Provider, identifier: String) -> Album {
        Album {
            id:            Uuid::new_v4(),
            provider:      provider,
            identifier:    identifier,
            owner_id:      None,
            owner_name:    None,
            url:           "".to_string(),
            title:         "".to_string(),
            description:   None,
            thumbnail_url: None,
            artwork_url:   None,
            published_at:  UTC::now().naive_utc(),
            created_at:    UTC::now().naive_utc(),
            updated_at:    UTC::now().naive_utc(),
            state:         State::Alive,
            tracks:        vec![],
            artists:       None,
        }
    }
    fn find_by_entry_id(entry_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums p LEFT JOIN album_entries pe
                        ON p.id = pe.album_id
                        WHERE pe.entry_id = $1
                        ORDER BY p.published_at DESC",
                     Album::props_str("p."))).unwrap();
        let rows = stmt.query(&[&entry_id]).unwrap();
        Album::rows_to_items(rows)
    }
}

impl Album {
    fn add_track(&mut self, track: &Track) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO album_tracks (track_id, album_id)
                                 VALUES ($1, $2)"));
        try!(stmt.query(&[&track.id, &self.id]));
        Ok(())
    }
    fn add_artist(&mut self, artist: Artist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO track_artists (track_id, artist_id) VALUES ($1, $2)"));
        try!(stmt.query(&[&self.id, &artist.id]));
        match self.artists {
            Some(ref mut artists) => artists.push(artist),
            None                  => self.artists = Some(vec![artist]),

        }
        Ok(())
    }

    pub fn from_sp_album(album: &spotify::Album) -> Album {
        Album::new(Provider::Spotify, (*album).id.to_string())
            .update_with_sp_album(album)
            .clone()
    }

    pub fn from_am_album(album: &apple_music::Album) -> Album {
        Album::new(Provider::AppleMusic, (*album).id.to_string())
            .update_with_am_album(album)
            .clone()
    }

    fn add_tracks(&mut self, tracks: Vec<Track>) {
        self.tracks = tracks.iter().map(|t| {
            let mut t = t.clone();
            if let Ok(new_track) = Track::find_or_create(t.provider,
                                                         t.identifier.to_string()) {
                t.id      = new_track.id;
                let _     = self.add_track(&t);
            };
            t
        }).collect::<Vec<_>>();
    }

    pub fn update_with_sp_album(&mut self, album: &spotify::Album) -> &mut Album {
        self.provider       = Provider::Spotify;
        self.identifier     = album.id.to_string();
        if album.artists.len() > 0 {
            self.owner_id   = Some(album.artists[0].id.clone());
            self.owner_name = Some(album.artists[0].name.clone());
        }
        self.url            = album.uri.clone();
        self.title          = album.name.clone();
        self.description    = None;
        self.state          = State::Alive;
        self.published_at   = UTC::now().naive_utc();
        if album.images.len() > 0 {
            self.artwork_url   = Some(album.images[0].url.clone());
            self.thumbnail_url = Some(album.images[0].url.clone());
        }
        if album.images.len() > 1 {
            self.thumbnail_url = Some(album.images[1].url.clone());
        }
        let artists = album.artists
            .iter()
            .map(|artist| {
                Artist::find_or_create(self.provider, artist.id.to_string())
                    .map(|mut a| a.update_with_sp_artist(&artist).clone())
            })
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<Artist>>();
        for mut a in artists {
            let _ = a.save();
            let _ = self.add_artist(a);
        }
        let tracks = album.tracks.clone()
            .map(|t| t.items).unwrap_or(vec![]).iter()
            .map(|ref t| Track::from_sp_track(t))
            .map(|ref mut t| t.update_with_sp_album(&album).clone())
            .collect::<Vec<_>>();
        self.add_tracks(tracks);
        self
    }

    pub fn update_with_am_album(&mut self, album: &apple_music::Album) -> &mut Album {
        self.provider      = Provider::AppleMusic;
        self.identifier    = album.id.to_string();
        self.owner_id      = Some(album.album_artist.to_string());
        self.owner_name    = Some(album.album_artist.to_string());
        self.url           = album.music_url.to_string();
        self.title         = album.title.to_string();
        self.description   = None;
        self.thumbnail_url = Some(album.artwork_url.to_string());
        self.artwork_url   = Some(album.artwork_url.to_string());
        self.state         = State::Alive;
        if let Ok(mut artist) = Artist::find_or_create(self.provider,
                                                       album.album_artist.to_string()) {
            artist.name  = album.album_artist.to_string();
            let _ = artist.save();
            let _ = self.add_artist(artist);
        }
        self
    }

    pub fn disable(&mut self) -> &mut Album {
        self.state = State::Dead;
        self
    }

    pub fn delete(&self) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("DELETE FROM albums WHERE id=$1"));
        let result = stmt.query(&[&self.id]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected)
        }
    }
}

#[cfg(test)]
mod test {
    use model::enclosure::Enclosure;
    use model::Model;
    use super::Album;
    use Provider;
    #[test]
    fn test_new() {
        let album = Album::new(Provider::YouTube,
                               "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string());
        assert_eq!(album.provider, Provider::YouTube);
        assert_eq!(&album.identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX")
    }
    #[test]
    fn test_find_of_create() {
        let identifier = "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string();
        let result     = Album::find_or_create(Provider::YouTube, identifier);
        assert!(result.is_ok());
    }
    #[test]
    fn test_delete() {
        let identifier = "test_delete".to_string();
        let album      = Album::find_or_create(Provider::YouTube, identifier).unwrap();
        let result     = album.delete();
        assert!(result.is_ok());
    }
    #[test]
    fn test_save() {
        let id = "test_save";
        let mut album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        album.title   = "title".to_string();
        let result    = album.save();
        assert!(result.is_ok());
        let album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        assert_eq!(&album.title, "title");
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use rustc_serialize::json::{ToJson, Json};
use postgres;
use uuid::Uuid;
use std::fmt;
use chrono::{NaiveDateTime, UTC, DateTime};

use apple_music;
use spotify;
use error::Error;
use super::{conn, Model};
use model::provider::Provider;
use model::state::State;
use model::enclosure::Enclosure;
use model::artist::Artist;

static PROPS: [&'static str; 14]  = ["id",
                                     "provider",
                                     "identifier",
                                     "owner_id",
                                     "owner_name",
                                     "url",
                                     "title",
                                     "description",
                                     "thumbnail_url",
                                     "artwork_url",
                                     "published_at",
                                     "created_at",
                                     "updated_at",
                                     "state"];

#[derive(Debug, Clone)]
pub struct Album {
    pub id:            Uuid,
    pub provider:      Provider,
    pub identifier:    String,
    pub owner_id:      Option<String>,
    pub owner_name:    Option<String>,
    pub url:           String,
    pub title:         String,
    pub description:   Option<String>,
    pub thumbnail_url: Option<String>,
    pub artwork_url:   Option<String>,
    pub published_at:  NaiveDateTime,
    pub created_at:    NaiveDateTime,
    pub updated_at:    NaiveDateTime,
    pub state:         State,
    pub artists:       Option<Vec<Artist>>,
}

impl PartialEq for Album {
    fn eq(&self, p: &Album) -> bool {
        return self.identifier == p.identifier && self.provider == p.provider
    }
}

impl ToJson for Album {
    fn to_json(&self) -> Json {
        let published_at = DateTime::<UTC>::from_utc(self.published_at, UTC);
        let created_at   = DateTime::<UTC>::from_utc(self.created_at  , UTC);
        let updated_at   = DateTime::<UTC>::from_utc(self.updated_at  , UTC);
        let mut d = BTreeMap::new();
        d.insert("id".to_string()           , self.id.to_string().to_json());
        d.insert("provider".to_string()     , self.provider.to_json());
        d.insert("identifier".to_string()   , self.identifier.to_json());
        d.insert("url".to_string()          , self.url.to_json());
        d.insert("title".to_string()        , self.title.to_json());
        d.insert("description".to_string()  , self.description.to_json());
        d.insert("thumbnail_url".to_string(), self.thumbnail_url.to_json());
        d.insert("artwork_url".to_string()  , self.artwork_url.to_json());
        d.insert("published_at".to_string() , published_at.to_rfc3339().to_json());
        d.insert("created_at".to_string()   , created_at.to_rfc3339().to_json());
        d.insert("updated_at".to_string()   , updated_at.to_rfc3339().to_json());
        d.insert("state".to_string()        , self.state.to_json());
        d.insert("artists".to_string()      , self.artists.to_json());
        Json::Object(d)
    }
}

impl fmt::Display for Album {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.provider, self.identifier)
    }
}

impl Model for Album {
    fn table_name() -> String { "albums".to_string() }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }

    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Album> {
        let mut albums = Vec::new();
        for row in rows.iter() {
            let album = Album {
                id:            row.get(0),
                provider:      Provider::new(row.get(1)),
                identifier:    row.get(2),
                owner_id:      row.get(3),
                owner_name:    row.get(4),
                url:           row.get(5),
                title:         row.get(6),
                description:   row.get(7),
                thumbnail_url: row.get(8),
                artwork_url:   row.get(9),
                published_at:  row.get(10),
                created_at:    row.get(11),
                updated_at:    row.get(12),
                state:         State::new(row.get(13)),
                artists:       None,
            };
            albums.push(album)
        }
        albums
    }

    fn create(&self) -> Result<Album, Error> {
        let conn = conn().unwrap();
        let stmt = conn.prepare("INSERT INTO albums (provider, identifier, url, title)
                                 VALUES ($1, $2, $3, $4) RETURNING id").unwrap();
        let rows = try!(stmt.query(&[&self.provider.to_string(), &self.identifier, &self.url, &self.title]));
        let mut album = self.clone();
        for row in rows.iter() {
            album.id = row.get(0);
        }
        Ok(album)
    }

    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = UTC::now().naive_utc();
        let conn = conn().unwrap();
        let stmt = conn.prepare("UPDATE albums SET
                                 provider      = $2,
                                 identifier    = $3,
                                 owner_id      = $4,
                                 owner_name    = $5,
                                 url           = $6,
                                 title         = $7,
                                 description   = $8,
                                 thumbnail_url = $9,
                                 artwork_url   = $10,
                                 published_at  = $11,
                                 created_at    = $12,
                                 updated_at    = $13,
                                 state         = $14
                                 WHERE id = $1").unwrap();
        let result = stmt.query(&[&self.id,
                                  &self.provider.to_string(),
                                  &self.identifier,
                                  &self.owner_id,
                                  &self.owner_name,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.thumbnail_url,
                                  &self.artwork_url,
                                  &self.published_at,
                                  &self.created_at,
                                  &self.updated_at,
                                  &self.state.to_string(),
        ]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected),
        }
    }
}

impl Enclosure for Album {
    fn new(provider: Provider, identifier: String) -> Album {
        Album {
            id:            Uuid::new_v4(),
            provider:      provider,
            identifier:    identifier,
            owner_id:      None,
            owner_name:    None,
            url:           "".to_string(),
            title:         "".to_string(),
            description:   None,
            thumbnail_url: None,
            artwork_url:   None,
            published_at:  UTC::now().naive_utc(),
            created_at:    UTC::now().naive_utc(),
            updated_at:    UTC::now().naive_utc(),
            state:         State::Alive,
            artists:       None,
        }
    }
    fn find_by_entry_id(entry_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums p LEFT JOIN album_entries pe
                        ON p.id = pe.album_id
                        WHERE pe.entry_id = $1
                        ORDER BY p.published_at DESC",
                     Album::props_str("p."))).unwrap();
        let rows = stmt.query(&[&entry_id]).unwrap();
        Album::rows_to_items(rows)
    }
}

impl Album {
    fn add_artist(&mut self, artist: Artist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO track_artists (track_id, artist_id) VALUES ($1, $2)"));
        try!(stmt.query(&[&self.id, &artist.id]));
        match self.artists {
            Some(ref mut artists) => artists.push(artist),
            None                  => self.artists = Some(vec![artist]),

        }
        Ok(())
    }

    pub fn from_sp_album(album: &spotify::Album) -> Album {
        Album::new(Provider::Spotify, (*album).id.to_string())
            .update_with_sp_album(album)
            .clone()
    }

    pub fn from_am_album(album: &apple_music::Album) -> Album {
        Album::new(Provider::AppleMusic, (*album).id.to_string())
            .update_with_am_album(album)
            .clone()
    }

    pub fn update_with_sp_album(&mut self, album: &spotify::Album) -> &mut Album {
        self.provider       = Provider::Spotify;
        self.identifier     = album.id.to_string();
        self.url            = album.uri.clone();
        self.title          = album.name.clone();
        self.description    = None;
        self.state          = State::Alive;
        self.published_at   = UTC::now().naive_utc();
        if album.images.len() > 0 {
            self.artwork_url   = Some(album.images[0].url.clone());
            self.thumbnail_url = Some(album.images[0].url.clone());
        }
        if album.images.len() > 1 {
            self.thumbnail_url = Some(album.images[1].url.clone());
        }
        let artists = album.artists
            .iter()
            .map(|artist| {
                Artist::find_or_create(self.provider, artist.id.to_string())
                    .map(|mut a| a.update_with_sp_artist(&artist).clone())
            })
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<Artist>>();
        for mut a in artists {
            let _ = a.save();
            let _ = self.add_artist(a);
        }
        self
    }

    pub fn update_with_am_album(&mut self, album: &apple_music::Album) -> &mut Album {
        self.provider      = Provider::AppleMusic;
        self.identifier    = album.id.to_string();
        self.url           = album.music_url.to_string();
        self.title         = album.title.to_string();
        self.description   = None;
        self.thumbnail_url = Some(album.artwork_url.to_string());
        self.artwork_url   = Some(album.artwork_url.to_string());
        self.state         = State::Alive;
        if let Ok(mut artist) = Artist::find_or_create(self.provider,
                                                       album.album_artist.to_string()) {
            artist.name  = album.album_artist.to_string();
            let _ = artist.save();
            let _ = self.add_artist(artist);
        }
        self
    }

    pub fn disable(&mut self) -> &mut Album {
        self.state = State::Dead;
        self
    }

    pub fn delete(&self) -> Result<(), Error> {
        let conn = conn().unwrap();
        let stmt = conn.prepare("DELETE FROM albums WHERE id=$1").unwrap();
        let result = stmt.query(&[&self.id]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected)
        }
    }
}

#[cfg(test)]
mod test {
    use model::enclosure::Enclosure;
    use model::Model;
    use super::Album;
    use Provider;
    #[test]
    fn test_new() {
        let album = Album::new(Provider::YouTube,
                               "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string());
        assert_eq!(album.provider, Provider::YouTube);
        assert_eq!(&album.identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX")
    }
    #[test]
    fn test_find_of_create() {
        let identifier = "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string();
        let result     = Album::find_or_create(Provider::YouTube, identifier);
        assert!(result.is_ok());
    }
    #[test]
    fn test_delete() {
        let identifier = "test_delete".to_string();
        let album      = Album::find_or_create(Provider::YouTube, identifier).unwrap();
        let result     = album.delete();
        assert!(result.is_ok());
    }
    #[test]
    fn test_save() {
        let id = "test_save";
        let mut album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        album.title   = "title".to_string();
        let result    = album.save();
        assert!(result.is_ok());
        let album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        assert_eq!(&album.title, "title");
    }
}
<commit_msg>Set owner_id and owner_name of Album correctly<commit_after>use std::collections::BTreeMap;
use rustc_serialize::json::{ToJson, Json};
use postgres;
use uuid::Uuid;
use std::fmt;
use chrono::{NaiveDateTime, UTC, DateTime};

use apple_music;
use spotify;
use error::Error;
use super::{conn, Model};
use model::provider::Provider;
use model::state::State;
use model::enclosure::Enclosure;
use model::artist::Artist;

static PROPS: [&'static str; 14]  = ["id",
                                     "provider",
                                     "identifier",
                                     "owner_id",
                                     "owner_name",
                                     "url",
                                     "title",
                                     "description",
                                     "thumbnail_url",
                                     "artwork_url",
                                     "published_at",
                                     "created_at",
                                     "updated_at",
                                     "state"];

#[derive(Debug, Clone)]
pub struct Album {
    pub id:            Uuid,
    pub provider:      Provider,
    pub identifier:    String,
    pub owner_id:      Option<String>,
    pub owner_name:    Option<String>,
    pub url:           String,
    pub title:         String,
    pub description:   Option<String>,
    pub thumbnail_url: Option<String>,
    pub artwork_url:   Option<String>,
    pub published_at:  NaiveDateTime,
    pub created_at:    NaiveDateTime,
    pub updated_at:    NaiveDateTime,
    pub state:         State,
    pub artists:       Option<Vec<Artist>>,
}

impl PartialEq for Album {
    fn eq(&self, p: &Album) -> bool {
        return self.identifier == p.identifier && self.provider == p.provider
    }
}

impl ToJson for Album {
    fn to_json(&self) -> Json {
        let published_at = DateTime::<UTC>::from_utc(self.published_at, UTC);
        let created_at   = DateTime::<UTC>::from_utc(self.created_at  , UTC);
        let updated_at   = DateTime::<UTC>::from_utc(self.updated_at  , UTC);
        let mut d = BTreeMap::new();
        d.insert("id".to_string()           , self.id.to_string().to_json());
        d.insert("provider".to_string()     , self.provider.to_json());
        d.insert("identifier".to_string()   , self.identifier.to_json());
        d.insert("owner_id".to_string()     , self.owner_id.to_json());
        d.insert("owner_name".to_string()   , self.owner_name.to_json());
        d.insert("url".to_string()          , self.url.to_json());
        d.insert("title".to_string()        , self.title.to_json());
        d.insert("description".to_string()  , self.description.to_json());
        d.insert("thumbnail_url".to_string(), self.thumbnail_url.to_json());
        d.insert("artwork_url".to_string()  , self.artwork_url.to_json());
        d.insert("published_at".to_string() , published_at.to_rfc3339().to_json());
        d.insert("created_at".to_string()   , created_at.to_rfc3339().to_json());
        d.insert("updated_at".to_string()   , updated_at.to_rfc3339().to_json());
        d.insert("state".to_string()        , self.state.to_json());
        d.insert("artists".to_string()      , self.artists.to_json());
        Json::Object(d)
    }
}

impl fmt::Display for Album {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}:{}", self.provider, self.identifier)
    }
}

impl Model for Album {
    fn table_name() -> String { "albums".to_string() }
    fn props_str(prefix: &str) -> String {
        PROPS
            .iter()
            .map(|&p| format!("{}{}", prefix, p))
            .collect::<Vec<String>>().join(",")
    }

    fn rows_to_items(rows: postgres::rows::Rows) -> Vec<Album> {
        let mut albums = Vec::new();
        for row in rows.iter() {
            let album = Album {
                id:            row.get(0),
                provider:      Provider::new(row.get(1)),
                identifier:    row.get(2),
                owner_id:      row.get(3),
                owner_name:    row.get(4),
                url:           row.get(5),
                title:         row.get(6),
                description:   row.get(7),
                thumbnail_url: row.get(8),
                artwork_url:   row.get(9),
                published_at:  row.get(10),
                created_at:    row.get(11),
                updated_at:    row.get(12),
                state:         State::new(row.get(13)),
                artists:       None,
            };
            albums.push(album)
        }
        albums
    }

    fn create(&self) -> Result<Album, Error> {
        let conn = conn().unwrap();
        let stmt = conn.prepare("INSERT INTO albums (provider, identifier, url, title)
                                 VALUES ($1, $2, $3, $4) RETURNING id").unwrap();
        let rows = try!(stmt.query(&[&self.provider.to_string(), &self.identifier, &self.url, &self.title]));
        let mut album = self.clone();
        for row in rows.iter() {
            album.id = row.get(0);
        }
        Ok(album)
    }

    fn save(&mut self) -> Result<(), Error> {
        self.updated_at = UTC::now().naive_utc();
        let conn = conn().unwrap();
        let stmt = conn.prepare("UPDATE albums SET
                                 provider      = $2,
                                 identifier    = $3,
                                 owner_id      = $4,
                                 owner_name    = $5,
                                 url           = $6,
                                 title         = $7,
                                 description   = $8,
                                 thumbnail_url = $9,
                                 artwork_url   = $10,
                                 published_at  = $11,
                                 created_at    = $12,
                                 updated_at    = $13,
                                 state         = $14
                                 WHERE id = $1").unwrap();
        let result = stmt.query(&[&self.id,
                                  &self.provider.to_string(),
                                  &self.identifier,
                                  &self.owner_id,
                                  &self.owner_name,
                                  &self.url,
                                  &self.title,
                                  &self.description,
                                  &self.thumbnail_url,
                                  &self.artwork_url,
                                  &self.published_at,
                                  &self.created_at,
                                  &self.updated_at,
                                  &self.state.to_string(),
        ]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected),
        }
    }
}

impl Enclosure for Album {
    fn new(provider: Provider, identifier: String) -> Album {
        Album {
            id:            Uuid::new_v4(),
            provider:      provider,
            identifier:    identifier,
            owner_id:      None,
            owner_name:    None,
            url:           "".to_string(),
            title:         "".to_string(),
            description:   None,
            thumbnail_url: None,
            artwork_url:   None,
            published_at:  UTC::now().naive_utc(),
            created_at:    UTC::now().naive_utc(),
            updated_at:    UTC::now().naive_utc(),
            state:         State::Alive,
            artists:       None,
        }
    }
    fn find_by_entry_id(entry_id: Uuid) -> Vec<Album> {
        let conn = conn().unwrap();
        let stmt = conn.prepare(
            &format!("SELECT {} FROM albums p LEFT JOIN album_entries pe
                        ON p.id = pe.album_id
                        WHERE pe.entry_id = $1
                        ORDER BY p.published_at DESC",
                     Album::props_str("p."))).unwrap();
        let rows = stmt.query(&[&entry_id]).unwrap();
        Album::rows_to_items(rows)
    }
}

impl Album {
    fn add_artist(&mut self, artist: Artist) -> Result<(), Error> {
        let conn = try!(conn());
        let stmt = try!(conn.prepare("INSERT INTO track_artists (track_id, artist_id) VALUES ($1, $2)"));
        try!(stmt.query(&[&self.id, &artist.id]));
        match self.artists {
            Some(ref mut artists) => artists.push(artist),
            None                  => self.artists = Some(vec![artist]),

        }
        Ok(())
    }

    pub fn from_sp_album(album: &spotify::Album) -> Album {
        Album::new(Provider::Spotify, (*album).id.to_string())
            .update_with_sp_album(album)
            .clone()
    }

    pub fn from_am_album(album: &apple_music::Album) -> Album {
        Album::new(Provider::AppleMusic, (*album).id.to_string())
            .update_with_am_album(album)
            .clone()
    }

    pub fn update_with_sp_album(&mut self, album: &spotify::Album) -> &mut Album {
        self.provider       = Provider::Spotify;
        self.identifier     = album.id.to_string();
        if album.artists.len() > 0 {
            self.owner_id   = Some(album.artists[0].id.clone());
            self.owner_name = Some(album.artists[0].name.clone());
        }
        self.url            = album.uri.clone();
        self.title          = album.name.clone();
        self.description    = None;
        self.state          = State::Alive;
        self.published_at   = UTC::now().naive_utc();
        if album.images.len() > 0 {
            self.artwork_url   = Some(album.images[0].url.clone());
            self.thumbnail_url = Some(album.images[0].url.clone());
        }
        if album.images.len() > 1 {
            self.thumbnail_url = Some(album.images[1].url.clone());
        }
        let artists = album.artists
            .iter()
            .map(|artist| {
                Artist::find_or_create(self.provider, artist.id.to_string())
                    .map(|mut a| a.update_with_sp_artist(&artist).clone())
            })
            .filter(|r| r.is_ok())
            .map(|r| r.unwrap())
            .collect::<Vec<Artist>>();
        for mut a in artists {
            let _ = a.save();
            let _ = self.add_artist(a);
        }
        self
    }

    pub fn update_with_am_album(&mut self, album: &apple_music::Album) -> &mut Album {
        self.provider      = Provider::AppleMusic;
        self.identifier    = album.id.to_string();
        self.owner_id      = Some(album.album_artist.to_string());
        self.owner_name    = Some(album.album_artist.to_string());
        self.url           = album.music_url.to_string();
        self.title         = album.title.to_string();
        self.description   = None;
        self.thumbnail_url = Some(album.artwork_url.to_string());
        self.artwork_url   = Some(album.artwork_url.to_string());
        self.state         = State::Alive;
        if let Ok(mut artist) = Artist::find_or_create(self.provider,
                                                       album.album_artist.to_string()) {
            artist.name  = album.album_artist.to_string();
            let _ = artist.save();
            let _ = self.add_artist(artist);
        }
        self
    }

    pub fn disable(&mut self) -> &mut Album {
        self.state = State::Dead;
        self
    }

    pub fn delete(&self) -> Result<(), Error> {
        let conn = conn().unwrap();
        let stmt = conn.prepare("DELETE FROM albums WHERE id=$1").unwrap();
        let result = stmt.query(&[&self.id]);
        match result {
            Ok(_)  => Ok(()),
            Err(_) => Err(Error::Unexpected)
        }
    }
}

#[cfg(test)]
mod test {
    use model::enclosure::Enclosure;
    use model::Model;
    use super::Album;
    use Provider;
    #[test]
    fn test_new() {
        let album = Album::new(Provider::YouTube,
                               "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string());
        assert_eq!(album.provider, Provider::YouTube);
        assert_eq!(&album.identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX")
    }
    #[test]
    fn test_find_of_create() {
        let identifier = "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string();
        let result     = Album::find_or_create(Provider::YouTube, identifier);
        assert!(result.is_ok());
    }
    #[test]
    fn test_delete() {
        let identifier = "test_delete".to_string();
        let album      = Album::find_or_create(Provider::YouTube, identifier).unwrap();
        let result     = album.delete();
        assert!(result.is_ok());
    }
    #[test]
    fn test_save() {
        let id = "test_save";
        let mut album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        album.title   = "title".to_string();
        let result    = album.save();
        assert!(result.is_ok());
        let album = Album::find_or_create(Provider::YouTube, id.to_string()).unwrap();
        assert_eq!(&album.title, "title");
    }
}
<|endoftext|>


<commit_before>use ast::AST;
use env::Env;
use runtime::Error;
use values::Value;
use primitives;
use special_forms;
use memory::Memory;

use std::rc::Rc;

pub fn eval(value: &AST, env: &mut Env, memory: &mut Memory) -> Result<Rc<Value>, Error> {
    match value {
        &AST::Atom(ref atom) =>
            eval_variable(atom, env),
        &AST::Bool(b) =>
            Ok(memory.boolean(b)),
        &AST::Integer(i) =>
            Ok(memory.integer(i)),
        &AST::List(ref list) =>
            eval_list(list, env, memory),
        _ =>
            Err(Error::IrreducibleValue(value.clone())),
    }
}

fn eval_list(list: &Vec<AST>, env: &mut Env, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if list.is_empty() {
        return Ok(mem.empty_list());
    }

    let head = list.head().unwrap();
    let tail = list.tail();

    if let &AST::List(ref list) = head {
        if list.head().unwrap() == &AST::Atom("quote".to_string()) {
            return Err(Error::UnappliableValue(Value::from_ast(head)));
        }
    }

    if let &AST::Atom(ref special_form) = head {
        let args = tail;

        match special_form.as_slice() {
            "and"    => return special_forms::and(args, env, mem),
            "define" => return special_forms::define(args, env, mem),
            "if"     => return special_forms::if_(args, env, mem),
            "lambda" => return special_forms::lambda(args, None, mem),
            "or"     => return special_forms::or(args, env, mem),
            "quote"  => return special_forms::quote(args, mem),
            _        => (),
        }
    }

    let fun  = try!(eval(head, env, mem));
    let args = try!(eval_args(tail, env, mem));

    match fun.deref() {
        &Value::Fn(ref name, ref args_names, ref body) =>
            apply(name, args_names, args, body, env, mem),
        &Value::PrimitiveFn(ref name) =>
            primitives::apply(name, args, mem),
        fun =>
            Err(Error::UnappliableValue(fun.clone()))
    }
}

fn eval_args(args: &[AST], env: &mut Env, mem: &mut Memory) -> Result<Vec<Rc<Value>>, Error> {
    let mut out = Vec::with_capacity(args.len());

    for arg in args.iter() {
        let evald_arg = try!(eval(arg, env, mem));
        out.push(evald_arg);
    }

    Ok(out)
}

fn eval_variable(name: &String, env: &mut Env) -> Result<Rc<Value>, Error> {
    match env.get(name) {
        Some(value) => Ok(value),
        None        => Err(Error::UnboundVariable(name.clone())),
    }
}

fn apply(name: &Option<String>, arg_names: &Vec<String>, arg_values: Vec<Rc<Value>>, body: &AST, env: &Env, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if arg_names.len() != arg_values.len() {
        return Err(Error::BadArity(name.clone()));
    }

    let mut inner_env = Env::wraps(env);
    for (name, value) in arg_names.iter().zip(arg_values.iter()) {
        inner_env.set(name.clone(), value.clone());
    }

    eval(body, &mut inner_env, mem)
}

fn atom_or_error(value: &AST) -> Result<String, Error> {
    if let &AST::Atom(ref atom) = value {
        Ok(atom.to_string())
    } else {
        Err(Error::WrongArgumentType(Value::from_ast(value)))
    }
}
<commit_msg>Improved flow clarity in eval.<commit_after>use ast::AST;
use env::Env;
use runtime::Error;
use values::Value;
use primitives;
use special_forms;
use memory::Memory;

use std::rc::Rc;

pub fn eval(value: &AST, env: &mut Env, memory: &mut Memory) -> Result<Rc<Value>, Error> {
    match value {
        &AST::Atom(ref atom) =>
            eval_variable(atom, env),
        &AST::Bool(b) =>
            Ok(memory.boolean(b)),
        &AST::Integer(i) =>
            Ok(memory.integer(i)),
        &AST::List(ref list) =>
            eval_list(list, env, memory),
        _ =>
            Err(Error::IrreducibleValue(value.clone())),
    }
}

fn eval_list(list: &Vec<AST>, env: &mut Env, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if list.is_empty() {
        return Ok(mem.empty_list());
    }

    let head = list.head().unwrap();
    let tail = list.tail();

    if is_quoted(head) {
        return Err(Error::UnappliableValue(Value::from_ast(head)));
    }

    if let &AST::Atom(ref special_form) = head {
        let args = tail;

        match special_form.as_slice() {
            "and"    => return special_forms::and(args, env, mem),
            "define" => return special_forms::define(args, env, mem),
            "if"     => return special_forms::if_(args, env, mem),
            "lambda" => return special_forms::lambda(args, None, mem),
            "or"     => return special_forms::or(args, env, mem),
            "quote"  => return special_forms::quote(args, mem),
            _        => (),
        }
    }

    let fun  = try!(eval(head, env, mem));
    let args = try!(eval_args(tail, env, mem));

    match fun.deref() {
        &Value::Fn(ref name, ref args_names, ref body) =>
            apply(name, args_names, args, body, env, mem),
        &Value::PrimitiveFn(ref name) =>
            primitives::apply(name, args, mem),
        fun =>
            Err(Error::UnappliableValue(fun.clone()))
    }
}

fn eval_args(args: &[AST], env: &mut Env, mem: &mut Memory) -> Result<Vec<Rc<Value>>, Error> {
    let mut out = Vec::with_capacity(args.len());

    for arg in args.iter() {
        let evald_arg = try!(eval(arg, env, mem));
        out.push(evald_arg);
    }

    Ok(out)
}

fn eval_variable(name: &String, env: &mut Env) -> Result<Rc<Value>, Error> {
    match env.get(name) {
        Some(value) => Ok(value),
        None        => Err(Error::UnboundVariable(name.clone())),
    }
}

fn apply(name: &Option<String>, arg_names: &Vec<String>, arg_values: Vec<Rc<Value>>, body: &AST, env: &Env, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if arg_names.len() != arg_values.len() {
        return Err(Error::BadArity(name.clone()));
    }

    let mut inner_env = wrap_env(arg_names, arg_values, env);
    eval(body, &mut inner_env, mem)
}

fn wrap_env<'a>(arg_names: &Vec<String>, arg_values: Vec<Rc<Value>>, outer_env: &'a Env) -> Env<'a> {
    let mut env = Env::wraps(outer_env);

    for (name, value) in arg_names.iter().zip(arg_values.iter()) {
        env.set(name.clone(), value.clone());
    }

    env
}

fn is_quoted(value: &AST) -> bool {
    match value {
        &AST::List(ref list) =>
            match list[0] {
                AST::Atom(ref a) if a.as_slice() == "quote" =>
                    true,
                _ =>
                    false
            },
        _ =>
            false
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Brendan Zabarauskas
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Macro instantiation tests

#[macro_use]
extern crate approx;

#[test]
fn test_abs_diff_eq() {
    let _: bool = abs_diff_eq!(1.0, 1.0);
    let _: bool = abs_diff_eq!(1.0, 1.0, epsilon = 1.0);
}

#[test]
fn test_abs_diff_ne() {
    let _: bool = abs_diff_ne!(1.0, 1.0);
    let _: bool = abs_diff_ne!(1.0, 1.0, epsilon = 1.0);
}

#[test]
fn test_relative_eq() {
    let _: bool = relative_eq!(1.0, 1.0);
    let _: bool = relative_eq!(1.0, 1.0, epsilon = 1.0);
    let _: bool = relative_eq!(1.0, 1.0, max_relative = 1.0);
    let _: bool = relative_eq!(1.0, 1.0, epsilon = 1.0, max_relative = 1.0);
}

#[test]
fn test_relative_ne() {
    let _: bool = relative_ne!(1.0, 1.0);
    let _: bool = relative_ne!(1.0, 1.0, epsilon = 1.0);
    let _: bool = relative_ne!(1.0, 1.0, max_relative = 1.0);
    let _: bool = relative_ne!(1.0, 1.0, epsilon = 1.0, max_relative = 1.0);
}

#[test]
fn test_ulps_eq() {
    let _: bool = ulps_eq!(1.0, 1.0);
    let _: bool = ulps_eq!(1.0, 1.0, epsilon = 1.0);
    let _: bool = ulps_eq!(1.0, 1.0, max_ulps = 1);
    let _: bool = ulps_eq!(1.0, 1.0, epsilon = 1.0, max_ulps = 1);
}

#[test]
fn test_ulps_ne() {
    let _: bool = ulps_ne!(1.0, 1.0);
    let _: bool = ulps_ne!(1.0, 1.0, epsilon = 1.0);
    let _: bool = ulps_ne!(1.0, 1.0, max_ulps = 1);
    let _: bool = ulps_ne!(1.0, 1.0, epsilon = 1.0, max_ulps = 1);
}
<commit_msg>Test for trailing commas on macros<commit_after>// Copyright 2015 Brendan Zabarauskas
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Macro instantiation tests

#[macro_use]
extern crate approx;

#[test]
fn test_abs_diff_eq() {
    let _: bool = abs_diff_eq!(1.0, 1.0);
    let _: bool = abs_diff_eq!(1.0, 1.0, epsilon = 1.0);
}

#[test]
fn test_abs_diff_eq_trailing_commas() {
    let _: bool = abs_diff_eq!(1.0, 1.0,);
    let _: bool = abs_diff_eq!(1.0, 1.0, epsilon = 1.0,);
}

#[test]
fn test_abs_diff_ne() {
    let _: bool = abs_diff_ne!(1.0, 1.0);
    let _: bool = abs_diff_ne!(1.0, 1.0, epsilon = 1.0);
}

#[test]
fn test_abs_diff_ne_trailing_commas() {
    let _: bool = abs_diff_ne!(1.0, 1.0,);
    let _: bool = abs_diff_ne!(1.0, 1.0, epsilon = 1.0,);
}

#[test]
fn test_relative_eq() {
    let _: bool = relative_eq!(1.0, 1.0);
    let _: bool = relative_eq!(1.0, 1.0, epsilon = 1.0);
    let _: bool = relative_eq!(1.0, 1.0, max_relative = 1.0);
    let _: bool = relative_eq!(1.0, 1.0, epsilon = 1.0, max_relative = 1.0);
}

#[test]
fn test_relative_eq_trailing_commas() {
    let _: bool = relative_eq!(1.0, 1.0,);
    let _: bool = relative_eq!(1.0, 1.0, epsilon = 1.0, max_relative = 1.0,);
}

#[test]
fn test_relative_ne() {
    let _: bool = relative_ne!(1.0, 1.0);
    let _: bool = relative_ne!(1.0, 1.0, epsilon = 1.0);
    let _: bool = relative_ne!(1.0, 1.0, max_relative = 1.0);
    let _: bool = relative_ne!(1.0, 1.0, epsilon = 1.0, max_relative = 1.0);
}

#[test]
fn test_relative_ne_trailing_commas() {
    let _: bool = relative_ne!(1.0, 1.0,);
    let _: bool = relative_ne!(1.0, 1.0, epsilon = 1.0, max_relative = 1.0,);
}

#[test]
fn test_ulps_eq() {
    let _: bool = ulps_eq!(1.0, 1.0);
    let _: bool = ulps_eq!(1.0, 1.0, epsilon = 1.0);
    let _: bool = ulps_eq!(1.0, 1.0, max_ulps = 1);
    let _: bool = ulps_eq!(1.0, 1.0, epsilon = 1.0, max_ulps = 1);
}

#[test]
fn test_ulps_eq_trailing_commas() {
    let _: bool = ulps_eq!(1.0, 1.0,);
    let _: bool = ulps_eq!(1.0, 1.0, epsilon = 1.0, max_ulps = 1,);
}

#[test]
fn test_ulps_ne() {
    let _: bool = ulps_ne!(1.0, 1.0);
    let _: bool = ulps_ne!(1.0, 1.0, epsilon = 1.0);
    let _: bool = ulps_ne!(1.0, 1.0, max_ulps = 1);
    let _: bool = ulps_ne!(1.0, 1.0, epsilon = 1.0, max_ulps = 1);
}

#[test]
fn test_ulps_ne_trailing_commas() {
    let _: bool = ulps_ne!(1.0, 1.0,);
    let _: bool = ulps_ne!(1.0, 1.0, epsilon = 1.0, max_ulps = 1,);
}

<|endoftext|>


<commit_before>use std;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use rand;
use rand::Rng;

use aes::{pad, padding_valid, Aes128, MODE};
use aes::BLOCK_SIZE;

use serialize::from_base64_file;

use unstable_features::all_bytes;
use unstable_features::MoveFrom;

use xor::XOR;

use helper::ceil_div;

use errors::*;

use prefix_suffix_oracles::{DeterministicOracle, Oracle};
use prefix_suffix_oracles::{Oracle11, Oracle12, Oracle13, Oracle14, Oracle16};

fn matasano2_9() -> Result<()> {
    compare(
        [
            89,
            69,
            76,
            76,
            79,
            87,
            32,
            83,
            85,
            66,
            77,
            65,
            82,
            73,
            78,
            69,
            4,
            4,
            4,
            4,
        ].as_ref(),
        &pad(b"YELLOW SUBMARINE".as_ref(), 20)?,
    )
}

#[test]
fn aes_128_cbc() {
    let iv = [0; BLOCK_SIZE];
    let key = b"YELLOW SUBMARINE";
    let input = b"ABCDEFGHIJKLMNOP";
    assert_eq!(
        input.as_ref(),
        &input
            .encrypt(key, Some(&iv), MODE::CBC)
            .unwrap()
            .decrypt(key, Some(&iv), MODE::CBC)
            .unwrap()[..]
    );
}

fn matasano2_10() -> Result<()> {
    let key = b"YELLOW SUBMARINE";
    let input = from_base64_file(Path::new("data/10.txt"))?;
    let cleartext = input.decrypt(key, Some(&[0; BLOCK_SIZE]), MODE::CBC)?;

    //Read reference cleartext
    let path = Path::new("data/10.ref.txt");
    let mut file = File::open(&path)?;
    let mut cleartext_ref = String::new();
    file.read_to_string(&mut cleartext_ref)?;

    compare(cleartext_ref.as_bytes(), &cleartext)
}

fn uses_ecb(oracle: &mut Oracle11) -> Result<bool> {
    // Assumes that oracle prepends at most one block of jibber
    // TODO: Can we relax this condition?
    let input = vec![0; 3 * BLOCK_SIZE];
    let ciphertext = oracle.encrypt(&input)?;
    let blocks: Vec<&[u8]> = ciphertext.chunks(BLOCK_SIZE).skip(1).take(2).collect();
    Ok(blocks[0] == blocks[1])
}

fn prefix_plus_suffix_length<T: Oracle>(oracle: &T) -> Result<usize> {
    let initial = oracle.encrypt(&[])?.len();
    let input = [0; BLOCK_SIZE];
    //Would profit from range_inclusive
    if let Some(index) = (1..BLOCK_SIZE + 1).find(|&i| {
        if let Ok(ciphertext) = oracle.encrypt(&input[BLOCK_SIZE - i..]) {
            initial != ciphertext.len()
        } else {
            false
        }
    }) {
        Ok(initial - index)
    } else {
        bail!(
            "length of oracle output did not change, something is wrong with the provided oracle"
        );
    }
}

/* For an oracle prepending prefix and appending suffix to its input, this function returns
 * prefix.len()/BLOCK_SIZE, that is the number of blocks fully occupied by the prefix.
 *
 * To determine this number, we pass two different cleartexts to the oracle and count the number
 * of identical blocks at the start of the corresponding ciphertexts. */
fn prefix_blocks_count<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    if let Some(result) = oracle
        .encrypt(&[0])?
        .chunks(BLOCK_SIZE)
        .zip(oracle.encrypt(&[1])?.chunks(BLOCK_SIZE))
        .position(|(x, y)| x != y)
    {
        Ok(result)
    } else {
        bail!("no differing blocks found, something is wrong with the provided oracle");
    }
}

// We look at the first block C not fully occupied by the prefix and fill it with a
// constant block B, say consisting of 0's.
// This part of the cleartext in oracle looks as follows:
//
//                               <------ B ------->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 || 0 0 0 suffix[0] suffix[1]
// <---------- C ------------------------>
//
// We then successively reduce the length of B until the ciphertext of C
// changes. This happens as soon as the cleartext in oracle looks as follows:
//
//                               <-- B -->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 suffix[0] || suffix[1]
// <---------- C ---------------------------------->
//
// This gives us the length of the prefix in C.
//
// We need to do this with two different constants because suffix[0] might accidentally
// coincide with the constant we have chosen.

pub fn prefix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    let n = prefix_blocks_count(oracle)?;
    let helper = |k: u8| -> Result<usize> {
        let constant_block = vec![k; BLOCK_SIZE];

        let mut prev = oracle.encrypt(&constant_block)?;

        for i in 0..BLOCK_SIZE {
            let cur = oracle.encrypt(&constant_block[i + 1..])?;
            if prev.chunks(BLOCK_SIZE).nth(n) != cur.chunks(BLOCK_SIZE).nth(n) {
                return Ok(i);
            }
            prev = cur;
        }
        Ok(BLOCK_SIZE)
    };

    Ok(n * BLOCK_SIZE + std::cmp::min(helper(0)?, helper(1)?))
}

fn suffix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    Ok(prefix_plus_suffix_length(oracle)? - prefix_length(oracle)?)
}

#[test]
fn test_length_functions() {
    let key = random_block();
    let mut prefix = Vec::new();
    let mut suffix = Vec::new();
    for _ in 0..64 {
        for _ in 0..64 {
            {
                let oracle = CommonPrefixSuffixOracle {
                    key: key.clone(),
                    prefix: prefix.clone(),
                    suffix: suffix.clone(),
                    mode: MODE::CTR,
                };
                println!("{}", prefix.len());
                println!("{}", prefix_length(&oracle).unwrap());
                assert!(prefix.len() == prefix_length(&oracle).unwrap());
                //assert!(suffix.len() == suffix_length(&oracle).unwrap());
            }
            suffix.push(1);
        }
        suffix.clear();
        prefix.push(0);
    }
}

fn matasano2_11() -> Result<()> {
    let mut oracle = Oracle11::new()?;
    let uses_ecb = uses_ecb(&mut oracle)?;
    oracle.verify_solution(uses_ecb)
}

fn decrypt_suffix<T: DeterministicOracle>(oracle: &T) -> Result<Vec<u8>> {
    // The following input is chosen in such a way that the cleartext in oracle looks as follows:
    //
    //            input start      input end
    //                ↓                ↓
    // <-- prefix --> 0 ... 0 || 0 ... 0 suffix[0] || suffix[1] ...
    //                ↑          ↑
    //            prefix_len  prefix_blocks*BLOCK_SIZE
    //
    // The resulting ciphertext is compared to oracle([input, u]). The u yielding a match is
    // equal to suffix[0].

    let prefix_len = prefix_length(oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let suffix_len = suffix_length(oracle)?;

    let mut suffix = Vec::with_capacity(suffix_len);

    let mut input = vec![0; prefix_padding + BLOCK_SIZE - 1];
    let reference_ciphertexts = (0..BLOCK_SIZE)
        .map(|left_shift| oracle.encrypt(&input[left_shift..]))
        .collect::<Result<Vec<Vec<u8>>>>()?;

    for i in 0..suffix_len {
        let block = prefix_blocks + i / BLOCK_SIZE;
        //let block_range = block*BLOCK_SIZE .. (block + 1)*BLOCK_SIZE;
        let left_shift = i % BLOCK_SIZE;
        for u in all_bytes() {
            input.push(u);
            if reference_ciphertexts[left_shift][block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
                == oracle.encrypt(&input[left_shift..])?
                    [block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
            {
                suffix.push(u);
                break;
            }
            input.pop();
        }
    }
    Ok(suffix)
}

fn matasano2_12() -> Result<()> {
    let oracle = Oracle12::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn decode_profile(u: &[u8], sep: u8) -> HashMap<&[u8], &[u8]> {
    let mut p = HashMap::new();
    for pair in u.split(|&x| x == sep) {
        let mut components = pair.split(|&x| x == b'=');
        p.insert(components.next().unwrap(), components.next().unwrap_or(&[]));
    }
    p
}

/* The following function works under the single assumption that the target value "user" (to be
   replaced by "admin") is stored at the very end of the profile. */
fn matasano2_13() -> Result<()> {
    let oracle = Oracle13::new()?;

    let prefix_len = prefix_length(&oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let target_cleartext = b"admin".pad();
    let mut input = vec![0; prefix_padding];
    input.extend_from_slice(&target_cleartext);
    let target_last_block = oracle
        .encrypt(&input)?
        .split_off(prefix_blocks * BLOCK_SIZE);

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding + "user".len()])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    ciphertext[blocks * BLOCK_SIZE..].move_from2(target_last_block, 0, BLOCK_SIZE);

    oracle.verify_solution(&ciphertext)
}

fn matasano2_14() -> Result<()> {
    let oracle = Oracle14::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn random_block() -> Vec<u8> {
    let mut rng = rand::thread_rng();
    rng.gen_iter().take(BLOCK_SIZE).collect()
}

fn matasano2_15() -> Result<()> {
    compare(true, b"ICE ICE BABY\x04\x04\x04\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x05\x05\x05\x05".padding_valid())?;
    compare(false, b"ICE ICE BABY\x01\x02\x03\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x03\x03\x03".padding_valid())?;
    compare(
        true,
        padding_valid(
            b"ICE ICE BABY\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C",
            12,
        ).unwrap(),
    )?;
    Ok(())
}

fn matasano2_16() -> Result<()> {
    let oracle = Oracle16::new()?;

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    let target_last_block = b";admin=true".pad();
    let current_last_block = vec![BLOCK_SIZE as u8; BLOCK_SIZE];
    let attack_bitflip = target_last_block.xor(&current_last_block);

    // Flip the next to last block
    ciphertext[(blocks - 1) * BLOCK_SIZE..blocks * BLOCK_SIZE].xor_inplace(&attack_bitflip);

    oracle.verify_solution(&ciphertext)
}

pub fn run() {
    println!("Set 2");
    run_exercise(matasano2_9, 9);
    run_exercise(matasano2_10, 10);
    run_exercise(matasano2_11, 11);
    run_exercise(matasano2_12, 12);
    run_exercise(matasano2_13, 13);
    run_exercise(matasano2_14, 14);
    run_exercise(matasano2_15, 15);
    run_exercise(matasano2_16, 16);
}
<commit_msg>Fix length computation for CTR mode<commit_after>use std;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use rand;
use rand::Rng;

use aes::{pad, padding_valid, Aes128, MODE};
use aes::BLOCK_SIZE;

use serialize::from_base64_file;

use unstable_features::all_bytes;
use unstable_features::MoveFrom;

use xor::XOR;

use helper::ceil_div;

use errors::*;

use prefix_suffix_oracles::{DeterministicOracle, Oracle};
use prefix_suffix_oracles::{Oracle11, Oracle12, Oracle13, Oracle14, Oracle16};

fn matasano2_9() -> Result<()> {
    compare(
        [
            89,
            69,
            76,
            76,
            79,
            87,
            32,
            83,
            85,
            66,
            77,
            65,
            82,
            73,
            78,
            69,
            4,
            4,
            4,
            4,
        ].as_ref(),
        &pad(b"YELLOW SUBMARINE".as_ref(), 20)?,
    )
}

#[test]
fn aes_128_cbc() {
    let iv = [0; BLOCK_SIZE];
    let key = b"YELLOW SUBMARINE";
    let input = b"ABCDEFGHIJKLMNOP";
    assert_eq!(
        input.as_ref(),
        &input
            .encrypt(key, Some(&iv), MODE::CBC)
            .unwrap()
            .decrypt(key, Some(&iv), MODE::CBC)
            .unwrap()[..]
    );
}

fn matasano2_10() -> Result<()> {
    let key = b"YELLOW SUBMARINE";
    let input = from_base64_file(Path::new("data/10.txt"))?;
    let cleartext = input.decrypt(key, Some(&[0; BLOCK_SIZE]), MODE::CBC)?;

    //Read reference cleartext
    let path = Path::new("data/10.ref.txt");
    let mut file = File::open(&path)?;
    let mut cleartext_ref = String::new();
    file.read_to_string(&mut cleartext_ref)?;

    compare(cleartext_ref.as_bytes(), &cleartext)
}

fn uses_ecb(oracle: &mut Oracle11) -> Result<bool> {
    // Assumes that oracle prepends at most one block of jibber
    // TODO: Can we relax this condition?
    let input = vec![0; 3 * BLOCK_SIZE];
    let ciphertext = oracle.encrypt(&input)?;
    let blocks: Vec<&[u8]> = ciphertext.chunks(BLOCK_SIZE).skip(1).take(2).collect();
    Ok(blocks[0] == blocks[1])
}

fn uses_padding<T: Oracle>(oracle: &T) -> Result<bool> {
    Ok((oracle.encrypt(&[0])?.len()-oracle.encrypt(&[])?.len()) % BLOCK_SIZE == 0)
}

fn prefix_plus_suffix_length<T: Oracle>(oracle: &T) -> Result<usize> {
    let initial = oracle.encrypt(&[])?.len();
    if !uses_padding(oracle)? {
        return Ok(initial);
    }

    let input = [0; BLOCK_SIZE];
    //Would profit from range_inclusive
    if let Some(index) = (1..BLOCK_SIZE + 1).find(|&i| {
        if let Ok(ciphertext) = oracle.encrypt(&input[BLOCK_SIZE - i..]) {
            initial != ciphertext.len()
        } else {
            false
        }
    }) {
        Ok(initial - index)
    } else {
        bail!(
            "length of oracle output did not change, something is wrong with the provided oracle"
        );
    }
}

/* For an oracle prepending prefix and appending suffix to its input, this function returns
 * prefix.len()/BLOCK_SIZE, that is the number of blocks fully occupied by the prefix.
 *
 * To determine this number, we pass two different cleartexts to the oracle and count the number
 * of identical blocks at the start of the corresponding ciphertexts. */
fn prefix_blocks_count<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    if let Some(result) = oracle
        .encrypt(&[0])?
        .chunks(BLOCK_SIZE)
        .zip(oracle.encrypt(&[1])?.chunks(BLOCK_SIZE))
        .position(|(x, y)| x != y)
    {
        Ok(result)
    } else {
        bail!("no differing blocks found, something is wrong with the provided oracle");
    }
}

// We look at the first block C not fully occupied by the prefix and fill it with a
// constant block B, say consisting of 0's.
// This part of the cleartext in oracle looks as follows:
//
//                               <------ B ------->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 || 0 0 0 suffix[0] suffix[1]
// <---------- C ------------------------>
//
// We then successively reduce the length of B until the ciphertext of C
// changes. This happens as soon as the cleartext in oracle looks as follows:
//
//                               <-- B -->
// prefix[?] prefix[?] prefix[?] 0 0 ... 0 suffix[0] || suffix[1]
// <---------- C ---------------------------------->
//
// This gives us the length of the prefix in C.
//
// We need to do this with two different constants because suffix[0] might accidentally
// coincide with the constant we have chosen.

pub fn prefix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    let n = prefix_blocks_count(oracle)?;
    let helper = |k: u8| -> Result<usize> {
        let constant_block = vec![k; BLOCK_SIZE];

        let mut prev = oracle.encrypt(&constant_block)?;

        for i in 0..BLOCK_SIZE {
            let cur = oracle.encrypt(&constant_block[i + 1..])?;
            if prev.chunks(BLOCK_SIZE).nth(n) != cur.chunks(BLOCK_SIZE).nth(n) {
                return Ok(i);
            }
            prev = cur;
        }
        Ok(BLOCK_SIZE)
    };

    Ok(n * BLOCK_SIZE + std::cmp::min(helper(0)?, helper(1)?))
}

fn suffix_length<T: DeterministicOracle>(oracle: &T) -> Result<usize> {
    Ok(prefix_plus_suffix_length(oracle)? - prefix_length(oracle)?)
}

#[test]
fn test_length_functions() {
    let key = random_block();
    let mut prefix = Vec::new();
    let mut suffix = Vec::new();
    for _ in 0..64 {
        for _ in 0..64 {
            {
                let oracle = CommonPrefixSuffixOracle {
                    key: key.clone(),
                    prefix: prefix.clone(),
                    suffix: suffix.clone(),
                    mode: MODE::CTR,
                };
                println!("{}", prefix.len());
                println!("{}", prefix_length(&oracle).unwrap());
                assert!(prefix.len() == prefix_length(&oracle).unwrap());
                //assert!(suffix.len() == suffix_length(&oracle).unwrap());
            }
            suffix.push(1);
        }
        suffix.clear();
        prefix.push(0);
    }
}

fn matasano2_11() -> Result<()> {
    let mut oracle = Oracle11::new()?;
    let uses_ecb = uses_ecb(&mut oracle)?;
    oracle.verify_solution(uses_ecb)
}

fn decrypt_suffix<T: DeterministicOracle>(oracle: &T) -> Result<Vec<u8>> {
    // The following input is chosen in such a way that the cleartext in oracle looks as follows:
    //
    //            input start      input end
    //                ↓                ↓
    // <-- prefix --> 0 ... 0 || 0 ... 0 suffix[0] || suffix[1] ...
    //                ↑          ↑
    //            prefix_len  prefix_blocks*BLOCK_SIZE
    //
    // The resulting ciphertext is compared to oracle([input, u]). The u yielding a match is
    // equal to suffix[0].

    let prefix_len = prefix_length(oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let suffix_len = suffix_length(oracle)?;

    let mut suffix = Vec::with_capacity(suffix_len);

    let mut input = vec![0; prefix_padding + BLOCK_SIZE - 1];
    let reference_ciphertexts = (0..BLOCK_SIZE)
        .map(|left_shift| oracle.encrypt(&input[left_shift..]))
        .collect::<Result<Vec<Vec<u8>>>>()?;

    for i in 0..suffix_len {
        let block = prefix_blocks + i / BLOCK_SIZE;
        //let block_range = block*BLOCK_SIZE .. (block + 1)*BLOCK_SIZE;
        let left_shift = i % BLOCK_SIZE;
        for u in all_bytes() {
            input.push(u);
            if reference_ciphertexts[left_shift][block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
                == oracle.encrypt(&input[left_shift..])?
                    [block * BLOCK_SIZE..(block + 1) * BLOCK_SIZE]
            {
                suffix.push(u);
                break;
            }
            input.pop();
        }
    }
    Ok(suffix)
}

fn matasano2_12() -> Result<()> {
    let oracle = Oracle12::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn decode_profile(u: &[u8], sep: u8) -> HashMap<&[u8], &[u8]> {
    let mut p = HashMap::new();
    for pair in u.split(|&x| x == sep) {
        let mut components = pair.split(|&x| x == b'=');
        p.insert(components.next().unwrap(), components.next().unwrap_or(&[]));
    }
    p
}

/* The following function works under the single assumption that the target value "user" (to be
   replaced by "admin") is stored at the very end of the profile. */
fn matasano2_13() -> Result<()> {
    let oracle = Oracle13::new()?;

    let prefix_len = prefix_length(&oracle)?;
    let (prefix_blocks, prefix_padding) = ceil_div(prefix_len, BLOCK_SIZE);
    let target_cleartext = b"admin".pad();
    let mut input = vec![0; prefix_padding];
    input.extend_from_slice(&target_cleartext);
    let target_last_block = oracle
        .encrypt(&input)?
        .split_off(prefix_blocks * BLOCK_SIZE);

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding + "user".len()])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    ciphertext[blocks * BLOCK_SIZE..].move_from2(target_last_block, 0, BLOCK_SIZE);

    oracle.verify_solution(&ciphertext)
}

fn matasano2_14() -> Result<()> {
    let oracle = Oracle14::new()?;
    oracle.verify_suffix(&decrypt_suffix(&oracle)?)
}

pub fn random_block() -> Vec<u8> {
    let mut rng = rand::thread_rng();
    rng.gen_iter().take(BLOCK_SIZE).collect()
}

fn matasano2_15() -> Result<()> {
    compare(true, b"ICE ICE BABY\x04\x04\x04\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x05\x05\x05\x05".padding_valid())?;
    compare(false, b"ICE ICE BABY\x01\x02\x03\x04".padding_valid())?;
    compare(false, b"ICE ICE BABY\x03\x03\x03".padding_valid())?;
    compare(
        true,
        padding_valid(
            b"ICE ICE BABY\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C\x0C",
            12,
        ).unwrap(),
    )?;
    Ok(())
}

fn matasano2_16() -> Result<()> {
    let oracle = Oracle16::new()?;

    let (blocks, padding) = ceil_div(prefix_plus_suffix_length(&oracle)?, BLOCK_SIZE);
    let mut ciphertext = oracle.encrypt(&vec![0; padding])?;
    compare((blocks + 1) * BLOCK_SIZE, ciphertext.len())?;

    let target_last_block = b";admin=true".pad();
    let current_last_block = vec![BLOCK_SIZE as u8; BLOCK_SIZE];
    let attack_bitflip = target_last_block.xor(&current_last_block);

    // Flip the next to last block
    ciphertext[(blocks - 1) * BLOCK_SIZE..blocks * BLOCK_SIZE].xor_inplace(&attack_bitflip);

    oracle.verify_solution(&ciphertext)
}

pub fn run() {
    println!("Set 2");
    run_exercise(matasano2_9, 9);
    run_exercise(matasano2_10, 10);
    run_exercise(matasano2_11, 11);
    run_exercise(matasano2_12, 12);
    run_exercise(matasano2_13, 13);
    run_exercise(matasano2_14, 14);
    run_exercise(matasano2_15, 15);
    run_exercise(matasano2_16, 16);
}
<|endoftext|>