<commit_before>extern crate twitter;

use std::io::{Write, Read};

fn main() {
    //IMPORTANT: make an app for yourself at apps.twitter.com and get your
    //key/secret into these files; this example won't work without them
    let consumer_key = include_str!("consumer_key").trim();
    let consumer_secret = include_str!("consumer_secret").trim();

    let token = twitter::Token::new(consumer_key, consumer_secret);

    let mut config = String::new();
    let user_id: i64;
    let username: String;
    let access_token: twitter::Token;

    //look at all this unwrapping! who told you it was my birthday?
    if let Ok(mut f) = std::fs::File::open("twitter_settings") {
        f.read_to_string(&mut config).unwrap();

        let mut iter = config.split('\n');

        username = iter.next().unwrap().to_string();
        user_id = i64::from_str_radix(&iter.next().unwrap(), 10).unwrap();
        access_token = twitter::Token::new(iter.next().unwrap(),
                                                 iter.next().unwrap());

        println!("Welcome back, {}!", username);
    }
    else {
        let request_token = match twitter::request_token(&token, "oob") {
            Ok(token) => token,
            Err(e) => {
                println!("Error: {}", e);
                return;
            },
        };

        println!("Go to the following URL, sign in, and give me the PIN that comes back:");
        println!("{}", twitter::authorize_url(&request_token));

        let mut pin = String::new();
        std::io::stdin().read_line(&mut pin).unwrap();
        println!("");

        let tok_result = twitter::access_token(&token, &request_token, pin).unwrap();

        access_token = tok_result.0;
        user_id = tok_result.1;
        username = tok_result.2;

        config.push_str(&username);
        config.push('\n');
        config.push_str(&format!("{}", user_id));
        config.push('\n');
        config.push_str(&access_token.key);
        config.push('\n');
        config.push_str(&access_token.secret);

        let mut f = std::fs::File::create("twitter_settings").unwrap();
        f.write_all(config.as_bytes()).unwrap();

        println!("Welcome, {}, let's get this show on the road!", username);
    }

    println!("");
    println!("Heterogeneous multi-user lookup:");

    let mut users: Vec<twitter::UserID> = vec![];
    users.push(user_id.into());
    users.push("SwiftOnSecurity".into());

    for user in twitter::user::lookup(&users, &token, &access_token).unwrap().response.iter() {
        print_user(user)
    }

    println!("");
    println!("Searching based on a term: (here, it's 'rustlang')");
    for resp in twitter::user::search("rustlang", &token, &access_token).with_page_size(5).take(5) {
        print_user(&resp.unwrap().response);
    }

    println!("");
    println!("Who do you follow?");
    for resp in twitter::user::friends_of(user_id, &token, &access_token).with_page_size(5).take(5) {
        print_user(&resp.unwrap().response);
    }

    println!("");
    println!("Who follows you?");
    for resp in twitter::user::followers_of(user_id, &token, &access_token).with_page_size(5).take(5) {
        print_user(&resp.unwrap().response);
    }
}

fn print_user(user: &twitter::user::TwitterUser) {
    println!("");
    println!("{} (@{})", user.name, user.screen_name);
    println!("Created at {}", user.created_at);
    println!("Follows {}, followed by {}", user.friends_count, user.followers_count);
    if let Some(ref desc) = user.description {
        println!("{}", desc);
    }
    else {
        println!("(no description provided)");
    }
    match (&user.location, &user.url) {
        (&Some(ref loc), &Some(ref link)) => println!("{} | {}", loc, link),
        (&None, &Some(ref link)) => println!("{}", link),
        (&Some(ref loc), &None) => println!("{}", loc),
        (&None, &None) => (),
    }
}
<commit_msg>update basic to include blocks list<commit_after>extern crate twitter;

use std::io::{Write, Read};

fn main() {
    //IMPORTANT: make an app for yourself at apps.twitter.com and get your
    //key/secret into these files; this example won't work without them
    let consumer_key = include_str!("consumer_key").trim();
    let consumer_secret = include_str!("consumer_secret").trim();

    let token = twitter::Token::new(consumer_key, consumer_secret);

    let mut config = String::new();
    let user_id: i64;
    let username: String;
    let access_token: twitter::Token;

    //look at all this unwrapping! who told you it was my birthday?
    if let Ok(mut f) = std::fs::File::open("twitter_settings") {
        f.read_to_string(&mut config).unwrap();

        let mut iter = config.split('\n');

        username = iter.next().unwrap().to_string();
        user_id = i64::from_str_radix(&iter.next().unwrap(), 10).unwrap();
        access_token = twitter::Token::new(iter.next().unwrap(),
                                                 iter.next().unwrap());

        println!("Welcome back, {}!", username);
    }
    else {
        let request_token = match twitter::request_token(&token, "oob") {
            Ok(token) => token,
            Err(e) => {
                println!("Error: {}", e);
                return;
            },
        };

        println!("Go to the following URL, sign in, and give me the PIN that comes back:");
        println!("{}", twitter::authorize_url(&request_token));

        let mut pin = String::new();
        std::io::stdin().read_line(&mut pin).unwrap();
        println!("");

        let tok_result = twitter::access_token(&token, &request_token, pin).unwrap();

        access_token = tok_result.0;
        user_id = tok_result.1;
        username = tok_result.2;

        config.push_str(&username);
        config.push('\n');
        config.push_str(&format!("{}", user_id));
        config.push('\n');
        config.push_str(&access_token.key);
        config.push('\n');
        config.push_str(&access_token.secret);

        let mut f = std::fs::File::create("twitter_settings").unwrap();
        f.write_all(config.as_bytes()).unwrap();

        println!("Welcome, {}, let's get this show on the road!", username);
    }

    println!("");
    println!("Heterogeneous multi-user lookup:");

    let mut users: Vec<twitter::UserID> = vec![];
    users.push(user_id.into());
    users.push("SwiftOnSecurity".into());

    for user in twitter::user::lookup(&users, &token, &access_token).unwrap().response.iter() {
        print_user(user)
    }

    println!("");
    println!("Searching based on a term: (here, it's 'rustlang')");
    for resp in twitter::user::search("rustlang", &token, &access_token).with_page_size(5).take(5) {
        print_user(&resp.unwrap().response);
    }

    println!("");
    println!("Who do you follow?");
    for resp in twitter::user::friends_of(user_id, &token, &access_token).with_page_size(5).take(5) {
        print_user(&resp.unwrap().response);
    }

    println!("");
    println!("Who follows you?");
    for resp in twitter::user::followers_of(user_id, &token, &access_token).with_page_size(5).take(5) {
        print_user(&resp.unwrap().response);
    }

    println!("");
    println!("Who have you blocked?");
    for resp in twitter::user::blocks(&token, &access_token).take(5) {
        print_user(&resp.unwrap().response);
    }
}

fn print_user(user: &twitter::user::TwitterUser) {
    println!("");
    println!("{} (@{})", user.name, user.screen_name);
    println!("Created at {}", user.created_at);
    println!("Follows {}, followed by {}", user.friends_count, user.followers_count);
    if let Some(ref desc) = user.description {
        println!("{}", desc);
    }
    else {
        println!("(no description provided)");
    }
    match (&user.location, &user.url) {
        (&Some(ref loc), &Some(ref link)) => println!("{} | {}", loc, link),
        (&None, &Some(ref link)) => println!("{}", link),
        (&Some(ref loc), &None) => println!("{}", loc),
        (&None, &None) => (),
    }
}
<|endoftext|>


<commit_before>#![cfg(feature = "postgresql")]

//! PostgreSQL backend,
//! needs citext, hstore and smlar.

extern crate postgres;
extern crate crypto;

use self::crypto::scrypt::{scrypt_simple,scrypt_check,ScryptParams};

use self::postgres::{Connection, TlsMode, Result as SQLResult};
use self::postgres::rows::Row;

pub struct Db(Connection);

use super::{Image,Tag,AnyWith,ImageBuilder,VoteImageError,parse_tag};
use super::super::CONFIG;

impl Default for Db { // Damn clippy
    fn default() -> Self {
        Self::new()
    }
}

lazy_static! {
    static ref SCRYPT_PARAMS: ScryptParams = ScryptParams::new(10, 8, 1); // 10 is really faster than 14
}

impl Db {
    pub fn new() -> Self {
        let conn = Connection::connect(format!("postgres://{}:{}@localhost", config!("postgres-login"), config!("postgres-password")), TlsMode::None).unwrap();
        conn.batch_execute("CREATE EXTENSION IF NOT EXISTS citext;
                            CREATE EXTENSION IF NOT EXISTS hstore;
                            CREATE EXTENSION IF NOT EXISTS smlar;

                            CREATE TABLE IF NOT EXISTS images(
                                id SERIAL PRIMARY KEY,
                                name TEXT NOT NULL UNIQUE,
                                tags TEXT[] NOT NULL,
                                uploader TEXT,
                                score INT NOT NULL DEFAULT 0,

                                got_from TEXT,
                                original_link TEXT,
                                rating CHAR
                            );

                            CREATE TABLE IF NOT EXISTS users(
                                id SERIAL PRIMARY KEY,
                                name CITEXT UNIQUE NOT NULL,
                                pass TEXT NOT NULL,
                                votes HSTORE
                            );").unwrap();
        Db(conn)
    }

    /// Add image, generate name from tags & id
    pub fn add_with_tags_name(&self, tags: &[String], ext: &str, uploader: &str) -> SQLResult<String> {
        let lastnum = self.0.query("SELECT id FROM images ORDER BY id DESC LIMIT 1", &[])?.get(0).get::<_, i32>("id");

        let name = format!("{}_{}.{}", lastnum + 1  , tags.join("_").replace("'","''"),ext);
        self.add_image(&ImageBuilder::new(&name, tags).uploader(uploader).finalize())?;
        Ok(name)
    }

    /// Add image
    pub fn add_image(&self, image: &ImageBuilder) -> SQLResult<()> {
        self.0.execute("INSERT into images (name,
                                            tags,
                                            got_from,
                                            original_link,
                                            rating,
                                            uploader,
                                            score)
                       VALUES ($1,
                               $2,
                               $3,
                               $4,
                               $5,
                               $6,
                               $7) ON CONFLICT (name) DO UPDATE SET tags = $2, score = $7",
        &[&image.name,&image.tags,&image.got_from, &image.original_link,&image.rating.map(|x| x.to_string()), &image.uploader, &image.score]).unwrap();
        Ok(())
    }

    pub fn get_image(&self, id: i32) -> SQLResult<Option<Image>> {
        let rows = self.0.query("SELECT * FROM images WHERE id = $1", &[&id])?;
        Ok(if !rows.is_empty() {
            Some(Db::extract_image(rows.get(0)))
        } else {
            None
        })
    }

    pub fn get_images<T: Into<Option<i32>>>(&self, take: T, skip: usize) -> SQLResult<Vec<Image>>{
        let take = match take.into() {
            Some(x) => x.to_string(),
            None    => "ALL".to_string()
        };

        Ok(self.0.query(&format!("SELECT * FROM images ORDER BY id DESC LIMIT {} OFFSET {}", take, skip as i32),&[])?
           .iter().fold(Vec::new(), |mut acc, row| {
               acc.push(Db::extract_image(row));
               acc
           }))
    }

    /// Search images by tags
    pub fn by_tags<T: Into<Option<i32>>>(&self, take: T, skip: usize, tags: &[String]) -> SQLResult<Vec<Image>> {
        let tags = tags.iter().map(|x| parse_tag(x)).collect::<Vec<_>>();
        let order = tags.iter().filter_map(|x| {
            match *x {
                Tag::OrderBy(ref by, ref ascdesc) => {
                    Some(format!("{:?}", by).to_lowercase() + " " + &format!("{:?}", ascdesc).to_uppercase())
                },
                _   => None
            }
        }).collect::<Vec<_>>();
        let order = match order.last() {
            Some(t) => t,
            None    => "id DESC"
        };

        let q = tags.clone().iter().map(|t| { // TODO: make it better?
            match *t {
                Tag::Include(ref incl) => format!(r"tags @> ARRAY['{}']", incl),
                Tag::Exclude(ref excl) => format!(r"NOT tags @> ARRAY['{}']", excl),
                Tag::AnyWith(ref x) => match *x {
                    AnyWith::Before(ref bef) => format!(r"(SELECT bool_or(tag ~ '^{}') FROM unnest(tags) t (tag))", bef),
                    AnyWith::After(ref aft) => format!(r"(SELECT bool_or(tag ~ '{}$') FROM unnest(tags) t (tag))", aft),
                },
                Tag::Rating(ref r) => {
                    let mut s = "(".to_string();
                    for tg in r {
                        s.push_str(&format!("rating = '{}' OR ", tg))
                    }
                    let _ = (0..4).inspect(|_| {s.pop(); }).collect::<Vec<_>>();
                    s.push_str(")");

                    s
                },

                Tag::From(ref f) => {
                    let mut s = "(".to_string();
                    for tg in f {
                        s.push_str(&format!("got_from = '{}' OR ", tg))
                    }
                    let _ = (0..4).inspect(|_| {s.pop(); }).collect::<Vec<_>>();
                    s.push_str(")");

                    s
                },

                Tag::Uploader(ref u) => {
                    let mut s = "(".to_string();
                    for tg in u {
                        s.push_str(&format!("uploader = '{}' OR ", tg))
                    }
                    let _ = (0..4).inspect(|_| {s.pop(); }).collect::<Vec<_>>();
                    s.push_str(")");

                    s
                },
                Tag::OrderBy(_,_) => {String::new()}
            }
        }).filter(|x| !x.is_empty()).collect::<Vec<_>>().join(" AND ");
        let q = if !q.is_empty() {
            format!("WHERE {}", q)
        } else {
            String::new()
        };

        let take = match take.into() {
            Some(x) => x.to_string(),
            None    => "ALL".to_string()
        };

        Ok(self.0.query(&format!("SELECT * FROM images {} ORDER BY {} LIMIT {} OFFSET {}", q, order, take, skip),&[])?
           .iter().fold(Vec::new(), |mut acc, row| {
               acc.push(Db::extract_image(row));
               acc
           }))
    }

    pub fn delete_image(&self, id: i32) -> SQLResult<String> {
        let name = self.0.query("SELECT * FROM images WHERE id = $1", &[&id])?.get(0).get::<_,String>("name");
        self.0.execute("DELETE FROM images WHERE id = $1", &[&id])?;
        Ok(name)
    }

    // true - all's OK, false - user already exists
    pub fn add_user(&self, login: &str, pass: &str) -> SQLResult<bool> {
        if self.0.query("SELECT * FROM users WHERE name = $1", &[&login])?.is_empty() && login.to_lowercase() != "sync" {
            let pass = scrypt_simple(pass, &SCRYPT_PARAMS).unwrap();

            self.0.execute("INSERT INTO users (name,pass) VALUES ($1,$2)", &[&login, &pass])?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Checks if pass & login match.
    /// Option is used to indicate that user does (not) exist
    pub fn check_user(&self, login: &str, pass: &str) -> SQLResult<Option<bool>> {
        let pass_hash = self.0.query("SELECT * FROM USERS WHERE name = $1", &[&login])?;
        if pass_hash.is_empty() {
            Ok(None)
        } else {
            let pass_hash = pass_hash.get(0).get::<_, String>("pass");
            Ok(Some(scrypt_check(pass, &pass_hash).unwrap()))
        }
    }

    // true - `+`, false - `-`, returns count of votes
    pub fn vote_image(&self, login: &str, image_id: i32 ,vote: bool) -> SQLResult<Result<i32, VoteImageError>> {
        let tr = self.0.transaction()?;
        let votechar =  if vote { "+" } else { "-" }.to_string();
        let previous = tr.query("SELECT votes -> $2 AS vote FROM users WHERE name = $1", &[&login, &image_id.to_string()])?;

        let newcount = if !previous.is_empty() && previous.get(0).get::<_,Option<String>>("vote") == Some(votechar.to_owned()) {
            tr.set_rollback();
            Err(VoteImageError::Already)
        } else {
            let res = if vote {
                tr.query("UPDATE images SET score = score + 1 WHERE id = $1 RETURNING score", &[&image_id])?
            } else {
                tr.query("UPDATE images SET score = score - 1 WHERE id = $1 RETURNING score", &[&image_id])?
            };
            if !res.is_empty() {
                tr.set_commit();
                Ok(res.get(0).get::<_,i32>("score"))
            } else {
                tr.set_rollback();
                Err(VoteImageError::NoImage)
            }
        };

        tr.execute("UPDATE users SET votes = hstore($2,$3) WHERE name = $1", &[&login, &image_id.to_string(), &votechar])?;

        Ok(newcount)
    }

    /// Find similiar images
    pub fn similiar<T: Into<Option<i32>>>(&self, id:i32, take: T, skip: usize) -> SQLResult<Vec<Image>> {
        let take = match take.into() {
            Some(x) => x.to_string(),
            None    => "ALL".to_string()
        };

        Ok(self.0.query(&format!("SELECT * FROM images
                                    WHERE id != $1
                                    ORDER BY smlar(tags, (SELECT tags FROM images WHERE id = $1)) DESC
                                    LIMIT {} OFFSET {}", take, skip as i32),&[&id])?
           .iter().fold(Vec::new(), |mut acc, row| {
               acc.push(Db::extract_image(row));
               acc
           }))
    }

    fn extract_image(row: Row) -> Image {
        Image{
            id: row.get("id"),
            name: row.get("name"),
            tags: row.get("tags"),
            got_from: row.get::<_, Option<String>>("got_from"),
            original_link: row.get::<_, Option<String>>("original_link"),
            rating: row.get::<_,Option<String>>("rating").map(|x| x.to_string().chars().collect::<Vec<_>>()[0]),
            uploader: row.get::<_,Option<String>>("uploader"),
            score: row.get::<_,i32>("score")
        }
    }
}
<commit_msg>Better query generator for tags<commit_after>#![cfg(feature = "postgresql")]

//! PostgreSQL backend,
//! needs citext, hstore and smlar.

extern crate postgres;
extern crate crypto;

use self::crypto::scrypt::{scrypt_simple,scrypt_check,ScryptParams};

use self::postgres::{Connection, TlsMode, Result as SQLResult};
use self::postgres::rows::Row;

pub struct Db(Connection);

use super::{Image,Tag,AnyWith,ImageBuilder,VoteImageError,parse_tag};
use super::super::CONFIG;

impl Default for Db { // Damn clippy
    fn default() -> Self {
        Self::new()
    }
}

lazy_static! {
    static ref SCRYPT_PARAMS: ScryptParams = ScryptParams::new(10, 8, 1); // 10 is really faster than 14
}

impl Db {
    pub fn new() -> Self {
        let conn = Connection::connect(format!("postgres://{}:{}@localhost", config!("postgres-login"), config!("postgres-password")), TlsMode::None).unwrap();
        conn.batch_execute("CREATE EXTENSION IF NOT EXISTS citext;
                            CREATE EXTENSION IF NOT EXISTS hstore;
                            CREATE EXTENSION IF NOT EXISTS smlar;

                            CREATE TABLE IF NOT EXISTS images(
                                id SERIAL PRIMARY KEY,
                                name TEXT NOT NULL UNIQUE,
                                tags TEXT[] NOT NULL,
                                uploader TEXT,
                                score INT NOT NULL DEFAULT 0,

                                got_from TEXT,
                                original_link TEXT,
                                rating CHAR
                            );

                            CREATE TABLE IF NOT EXISTS users(
                                id SERIAL PRIMARY KEY,
                                name CITEXT UNIQUE NOT NULL,
                                pass TEXT NOT NULL,
                                votes HSTORE
                            );").unwrap();
        Db(conn)
    }

    /// Add image, generate name from tags & id
    pub fn add_with_tags_name(&self, tags: &[String], ext: &str, uploader: &str) -> SQLResult<String> {
        let lastnum = self.0.query("SELECT id FROM images ORDER BY id DESC LIMIT 1", &[])?.get(0).get::<_, i32>("id");

        let name = format!("{}_{}.{}", lastnum + 1  , tags.join("_").replace("'","''"),ext);
        self.add_image(&ImageBuilder::new(&name, tags).uploader(uploader).finalize())?;
        Ok(name)
    }

    /// Add image
    pub fn add_image(&self, image: &ImageBuilder) -> SQLResult<()> {
        self.0.execute("INSERT into images (name,
                                            tags,
                                            got_from,
                                            original_link,
                                            rating,
                                            uploader,
                                            score)
                       VALUES ($1,
                               $2,
                               $3,
                               $4,
                               $5,
                               $6,
                               $7) ON CONFLICT (name) DO UPDATE SET tags = $2, score = $7",
        &[&image.name,&image.tags,&image.got_from, &image.original_link,&image.rating.map(|x| x.to_string()), &image.uploader, &image.score]).unwrap();
        Ok(())
    }

    pub fn get_image(&self, id: i32) -> SQLResult<Option<Image>> {
        let rows = self.0.query("SELECT * FROM images WHERE id = $1", &[&id])?;
        Ok(if !rows.is_empty() {
            Some(Db::extract_image(rows.get(0)))
        } else {
            None
        })
    }

    pub fn get_images<T: Into<Option<i32>>>(&self, take: T, skip: usize) -> SQLResult<Vec<Image>>{
        let take = match take.into() {
            Some(x) => x.to_string(),
            None    => "ALL".to_string()
        };

        Ok(self.0.query(&format!("SELECT * FROM images ORDER BY id DESC LIMIT {} OFFSET {}", take, skip as i32),&[])?
           .iter().fold(Vec::new(), |mut acc, row| {
               acc.push(Db::extract_image(row));
               acc
           }))
    }

    /// Search images by tags
    pub fn by_tags<T: Into<Option<i32>>>(&self, take: T, skip: usize, tags: &[String]) -> SQLResult<Vec<Image>> {
        let tags = tags.iter().map(|x| parse_tag(x)).collect::<Vec<_>>();
        let order = tags.iter().filter_map(|x| {
            match *x {
                Tag::OrderBy(ref by, ref ascdesc) => {
                    Some(format!("{:?}", by).to_lowercase() + " " + &format!("{:?}", ascdesc).to_uppercase())
                },
                _   => None
            }
        }).collect::<Vec<_>>();
        let order = match order.last() {
            Some(t) => t,
            None    => "id DESC"
        };

        /// Join tags for tags that can be separated with comma, e.g. rating or uploader
        fn join_tags(kind: &str, values: &[String]) -> String {
            values.iter().map(|s| format!("{} = '{}'", kind, s)).collect::<Vec<_>>().join(" OR ")
        }

        let q = tags.iter().map(|t| { // TODO: do something with OrderBy
            match *t {
                Tag::Include(ref incl)  => format!(r"tags @> ARRAY['{}']", incl),
                Tag::Exclude(ref excl)  => format!(r"NOT tags @> ARRAY['{}']", excl),
                Tag::AnyWith(ref x)     => match *x {
                    AnyWith::Before(ref bef) => format!(r"(SELECT bool_or(tag ~ '^{}') FROM unnest(tags) t (tag))", bef),
                    AnyWith::After(ref aft) => format!(r"(SELECT bool_or(tag ~ '{}$') FROM unnest(tags) t (tag))", aft),
                },
                Tag::Rating(ref r)      => join_tags("rating", r),
                Tag::From(ref f)        => join_tags("got_from", f),
                Tag::Uploader(ref u)    => join_tags("uploader", u),
                Tag::OrderBy(_,_)       => String::new() // <- This one
            }
        }).filter(|x| !x.is_empty()).collect::<Vec<_>>().join(" AND ");
        let q = if !q.is_empty() { format!("WHERE {}", q) } else { String::new() };

        let take = match take.into() {
            Some(x) => x.to_string(),
            None    => "ALL".to_string()
        };

        Ok(self.0.query(&format!("SELECT * FROM images {} ORDER BY {} LIMIT {} OFFSET {}", q, order, take, skip),&[])?
           .iter().fold(Vec::new(), |mut acc, row| {
               acc.push(Db::extract_image(row));
               acc
           }))
    }

    pub fn delete_image(&self, id: i32) -> SQLResult<String> {
        let name = self.0.query("SELECT * FROM images WHERE id = $1", &[&id])?.get(0).get::<_,String>("name");
        self.0.execute("DELETE FROM images WHERE id = $1", &[&id])?;
        Ok(name)
    }

    // true - all's OK, false - user already exists
    pub fn add_user(&self, login: &str, pass: &str) -> SQLResult<bool> {
        if self.0.query("SELECT * FROM users WHERE name = $1", &[&login])?.is_empty() && login.to_lowercase() != "sync" {
            let pass = scrypt_simple(pass, &SCRYPT_PARAMS).unwrap();

            self.0.execute("INSERT INTO users (name,pass) VALUES ($1,$2)", &[&login, &pass])?;
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Checks if pass & login match.
    /// Option is used to indicate that user does (not) exist
    pub fn check_user(&self, login: &str, pass: &str) -> SQLResult<Option<bool>> {
        let pass_hash = self.0.query("SELECT * FROM USERS WHERE name = $1", &[&login])?;
        if pass_hash.is_empty() {
            Ok(None)
        } else {
            let pass_hash = pass_hash.get(0).get::<_, String>("pass");
            Ok(Some(scrypt_check(pass, &pass_hash).unwrap()))
        }
    }

    // true - `+`, false - `-`, returns count of votes
    pub fn vote_image(&self, login: &str, image_id: i32 ,vote: bool) -> SQLResult<Result<i32, VoteImageError>> {
        let tr = self.0.transaction()?;
        let votechar =  if vote { "+" } else { "-" }.to_string();
        let previous = tr.query("SELECT votes -> $2 AS vote FROM users WHERE name = $1", &[&login, &image_id.to_string()])?;

        let newcount = if !previous.is_empty() && previous.get(0).get::<_,Option<String>>("vote") == Some(votechar.to_owned()) {
            tr.set_rollback();
            Err(VoteImageError::Already)
        } else {
            let res = if vote {
                tr.query("UPDATE images SET score = score + 1 WHERE id = $1 RETURNING score", &[&image_id])?
            } else {
                tr.query("UPDATE images SET score = score - 1 WHERE id = $1 RETURNING score", &[&image_id])?
            };
            if !res.is_empty() {
                tr.set_commit();
                Ok(res.get(0).get::<_,i32>("score"))
            } else {
                tr.set_rollback();
                Err(VoteImageError::NoImage)
            }
        };

        tr.execute("UPDATE users SET votes = hstore($2,$3) WHERE name = $1", &[&login, &image_id.to_string(), &votechar])?;

        Ok(newcount)
    }

    /// Find similiar images
    pub fn similiar<T: Into<Option<i32>>>(&self, id:i32, take: T, skip: usize) -> SQLResult<Vec<Image>> {
        let take = match take.into() {
            Some(x) => x.to_string(),
            None    => "ALL".to_string()
        };

        Ok(self.0.query(&format!("SELECT * FROM images
                                    WHERE id != $1
                                    ORDER BY smlar(tags, (SELECT tags FROM images WHERE id = $1)) DESC
                                    LIMIT {} OFFSET {}", take, skip as i32),&[&id])?
           .iter().fold(Vec::new(), |mut acc, row| {
               acc.push(Db::extract_image(row));
               acc
           }))
    }

    fn extract_image(row: Row) -> Image {
        Image{
            id: row.get("id"),
            name: row.get("name"),
            tags: row.get("tags"),
            got_from: row.get::<_, Option<String>>("got_from"),
            original_link: row.get::<_, Option<String>>("original_link"),
            rating: row.get::<_,Option<String>>("rating").map(|x| x.to_string().chars().collect::<Vec<_>>()[0]),
            uploader: row.get::<_,Option<String>>("uploader"),
            score: row.get::<_,i32>("score")
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, SocketAddr},
    time::Duration,
};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    sync::mpsc,
};

#[cfg(target_os = "android")]
use std::path::PathBuf;
#[cfg(target_os = "android")]
use tokio::net::UnixStream;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use byteorder::{BigEndian, ByteOrder};
use log::{debug, error, info};
use rand::Rng;
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{Name, RData, RecordType},
};

use crate::{
    config::{ConfigType, ServerConfig},
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        socks5::Address,
        sys::create_udp_socket,
        tcprelay::ProxyStream,
        utils::try_timeout,
    },
};

async fn stream_lookup<T>(qname: &Name, qtype: RecordType, stream: &mut T) -> io::Result<Message>
where
    T: AsyncReadExt + AsyncWriteExt + Unpin,
{
    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    let size = req_buffer.len();
    let mut send_buffer = vec![0; size + 2];

    BigEndian::write_u16(&mut send_buffer[0..2], size as u16);
    send_buffer[2..size + 2].copy_from_slice(&req_buffer[0..size]);
    stream.write_all(&send_buffer[0..size + 2]).await?;

    let mut res_buffer = vec![0; 2];
    stream.read_exact(&mut res_buffer[0..2]).await?;

    let size = BigEndian::read_u16(&res_buffer[0..2]) as usize;
    let mut res_buffer = vec![0; size];
    stream.read_exact(&mut res_buffer[0..size]).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

#[cfg(target_os = "android")]
async fn local_lookup(qname: &Name, qtype: RecordType, path: &PathBuf) -> io::Result<Message> {
    let mut stream = UnixStream::connect(path).await?;
    stream_lookup(qname, qtype, &mut stream).await
}

#[cfg(not(target_os = "android"))]
async fn local_lookup(qname: &Name, qtype: RecordType, server: &SocketAddr) -> io::Result<Message> {
    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0));
    let mut socket = UdpSocket::bind(bind_addr).await?;

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    socket.send_to(&req_buffer, server).await?;

    let mut res_buffer = vec![0; 512];
    socket.recv_from(&mut res_buffer).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

async fn proxy_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    ns: &Address,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<Message> {
    let mut stream = ProxyStream::connect_proxied(context, svr_cfg, ns).await?;
    stream_lookup(qname, qtype, &mut stream).await
}

async fn acl_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    #[cfg(target_os = "android")] local: &PathBuf,
    #[cfg(not(target_os = "android"))] local: &SocketAddr,
    remote: &Address,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<(Message, bool)> {
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        qtype, qname, local, remote
    );

    // remove the last dot from fqdn name
    let mut name = qname.to_ascii();
    name.pop();
    let addr = Address::DomainNameAddress(name, 0);
    let qname_in_proxy_list = context.check_qname_in_proxy_list(&addr);

    let remote_response_fut = async {
        match qname_in_proxy_list {
            Some(false) => None,
            _ => {
                let timeout = Some(Duration::new(3, 0));
                try_timeout(proxy_lookup(context.clone(), svr_cfg, remote, qname, qtype), timeout).await.ok()
            }
        }
    };

    let local_response = match qname_in_proxy_list {
        Some(true) => None,
        _ => {
            let timeout = Some(Duration::new(3, 0));
            try_timeout(local_lookup(qname, qtype, local), timeout).await.ok()
        }
    }.unwrap_or_else(|| Message::new());

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(|| Message::new());
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response.clone(), true));
        }
        Some(false) => {
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response.clone(), false));
        }
        None => (),
    }

    if local_response.answer_count() == 0 {
        let remote_response = remote_response_fut.await.unwrap_or_else(|| Message::new());
        return Ok((remote_response.clone(), true));
    }

    for rec in local_response.answers() {
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            _ => true,
        };
        if !forward {
            debug!("pick local response (ip): {:?}", local_response);
            return Ok((local_response.clone(), false))
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(|| Message::new());
    debug!("pick remote response (ip): {:?}", remote_response);
    Ok((remote_response.clone(), true))
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let mut qtx = qtx.clone();
        let server = balancer.pick_server();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let config = context.config();

                #[cfg(target_os = "android")]
                let local_addr = config.local_dns_path.as_ref().expect("local query DNS path");
                #[cfg(not(target_os = "android"))]
                let local_addr = config.local_dns_addr.as_ref().expect("local query DNS address");

                let remote_addr = config.remote_dns_addr.as_ref().expect("remote query DNS address");

                let question = &request.queries()[0];

                let qname = question.name();
                let qtype = question.query_type();
                let svr_cfg = server.server_config();

                let r = acl_lookup(context.clone(), svr_cfg, &local_addr, &remote_addr, qname, qtype).await;

                if let Ok((result, forward)) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Make clippy happy<commit_after>use std::{
    io,
    net::{IpAddr, SocketAddr},
    time::Duration,
};

use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    sync::mpsc,
};

#[cfg(target_os = "android")]
use std::path::PathBuf;
#[cfg(target_os = "android")]
use tokio::net::UnixStream;

#[cfg(not(target_os = "android"))]
use std::net::{Ipv4Addr, SocketAddrV4};
#[cfg(not(target_os = "android"))]
use tokio::net::UdpSocket;

use byteorder::{BigEndian, ByteOrder};
use log::{debug, error, info};
use rand::Rng;
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{Name, RData, RecordType},
};

use crate::{
    config::{ConfigType, ServerConfig},
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        socks5::Address,
        sys::create_udp_socket,
        tcprelay::ProxyStream,
        utils::try_timeout,
    },
};

async fn stream_lookup<T>(qname: &Name, qtype: RecordType, stream: &mut T) -> io::Result<Message>
where
    T: AsyncReadExt + AsyncWriteExt + Unpin,
{
    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    let size = req_buffer.len();
    let mut send_buffer = vec![0; size + 2];

    BigEndian::write_u16(&mut send_buffer[0..2], size as u16);
    send_buffer[2..size + 2].copy_from_slice(&req_buffer[0..size]);
    stream.write_all(&send_buffer[0..size + 2]).await?;

    let mut res_buffer = vec![0; 2];
    stream.read_exact(&mut res_buffer[0..2]).await?;

    let size = BigEndian::read_u16(&res_buffer[0..2]) as usize;
    let mut res_buffer = vec![0; size];
    stream.read_exact(&mut res_buffer[0..size]).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

#[cfg(target_os = "android")]
async fn local_lookup(qname: &Name, qtype: RecordType, path: &PathBuf) -> io::Result<Message> {
    let mut stream = UnixStream::connect(path).await?;
    stream_lookup(qname, qtype, &mut stream).await
}

#[cfg(not(target_os = "android"))]
async fn local_lookup(qname: &Name, qtype: RecordType, server: &SocketAddr) -> io::Result<Message> {
    let bind_addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0));
    let mut socket = UdpSocket::bind(bind_addr).await?;

    let mut message = Message::new();
    let mut query = Query::new();

    query.set_query_type(qtype);
    query.set_name(qname.clone());

    let id = rand::thread_rng().gen();
    message.set_id(id);
    message.set_recursion_desired(true);
    message.add_query(query);

    let req_buffer = message.to_vec()?;
    socket.send_to(&req_buffer, server).await?;

    let mut res_buffer = vec![0; 512];
    socket.recv_from(&mut res_buffer).await?;

    Ok(Message::from_vec(&res_buffer)?)
}

async fn proxy_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    ns: &Address,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<Message> {
    let mut stream = ProxyStream::connect_proxied(context, svr_cfg, ns).await?;
    stream_lookup(qname, qtype, &mut stream).await
}

async fn acl_lookup(
    context: SharedContext,
    svr_cfg: &ServerConfig,
    #[cfg(target_os = "android")] local: &PathBuf,
    #[cfg(not(target_os = "android"))] local: &SocketAddr,
    remote: &Address,
    qname: &Name,
    qtype: RecordType,
) -> io::Result<(Message, bool)> {
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        qtype, qname, local, remote
    );

    // remove the last dot from fqdn name
    let mut name = qname.to_ascii();
    name.pop();
    let addr = Address::DomainNameAddress(name, 0);
    let qname_in_proxy_list = context.check_qname_in_proxy_list(&addr);

    let remote_response_fut = async {
        match qname_in_proxy_list {
            Some(false) => None,
            _ => {
                let timeout = Some(Duration::new(3, 0));
                try_timeout(proxy_lookup(context.clone(), svr_cfg, remote, qname, qtype), timeout).await.ok()
            }
        }
    };

    let local_response = match qname_in_proxy_list {
        Some(true) => None,
        _ => {
            let timeout = Some(Duration::new(3, 0));
            try_timeout(local_lookup(qname, qtype, local), timeout).await.ok()
        }
    }.unwrap_or_else(Message::new);

    match qname_in_proxy_list {
        Some(true) => {
            let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
            debug!("pick remote response (qname): {:?}", remote_response);
            return Ok((remote_response, true));
        }
        Some(false) => {
            debug!("pick local response (qname): {:?}", local_response);
            return Ok((local_response, false));
        }
        None => (),
    }

    if local_response.answer_count() == 0 {
        let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
        return Ok((remote_response, true));
    }

    for rec in local_response.answers() {
        let forward = match rec.rdata() {
            RData::A(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            RData::AAAA(ref ip) => context.check_ip_in_proxy_list(&IpAddr::from(*ip)),
            _ => true,
        };
        if !forward {
            debug!("pick local response (ip): {:?}", local_response);
            return Ok((local_response, false))
        }
    }

    let remote_response = remote_response_fut.await.unwrap_or_else(Message::new);
    debug!("pick remote response (ip): {:?}", remote_response);
    Ok((remote_response, true))
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let mut qtx = qtx.clone();
        let server = balancer.pick_server();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let config = context.config();

                #[cfg(target_os = "android")]
                let local_addr = config.local_dns_path.as_ref().expect("local query DNS path");
                #[cfg(not(target_os = "android"))]
                let local_addr = config.local_dns_addr.as_ref().expect("local query DNS address");

                let remote_addr = config.remote_dns_addr.as_ref().expect("remote query DNS address");

                let question = &request.queries()[0];

                let qname = question.name();
                let qtype = question.query_type();
                let svr_cfg = server.server_config();

                let r = acl_lookup(context.clone(), svr_cfg, &local_addr, &remote_addr, qname, qtype).await;

                if let Ok((result, forward)) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   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.
//

//!
//! # DFA builder
//!
//! The DFA builder trait is implemented by classes that initialise DFAs.
//!

use super::state_machine::*;
use super::matcher::*;

///
/// Class that can build a particular type of DFA
///
pub trait DfaBuilder<InputSymbol, OutputSymbol, DfaType: Matcher<InputSymbol, OutputSymbol>> {
    ///
    /// Starts the next state for this DFA
    ///
    /// When this is first called, the DFA will enter state 0, then state 1, etc. If this hasn't been called yet then the DFA is not
    /// in a valid state and the other calls cannot be made.
    ///
    fn start_state(&mut self);

    ///
    /// Adds a transition to the current state
    ///
    /// Any input symbol can appear exactly once in a state, and must not overlap any other input symbol.
    ///
    fn transition(&mut self, symbol: InputSymbol, target_state: StateId);

    ///
    /// Sets the current state as an accepting state and sets the output symbol that will be produced if this is the longest match
    ///
    fn accept(&mut self, symbol: OutputSymbol);

    ///
    /// Finishes building the DFA and returns the matcher for the pattern it represents
    ///
    fn build(self) -> DfaType;
}
<commit_msg>Add a new condition to the DFA builder: callers are expected to sort transitions<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   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.
//

//!
//! # DFA builder
//!
//! The DFA builder trait is implemented by classes that initialise DFAs.
//!

use super::state_machine::*;
use super::matcher::*;

///
/// Class that can build a particular type of DFA
///
pub trait DfaBuilder<InputSymbol, OutputSymbol, DfaType: Matcher<InputSymbol, OutputSymbol>> {
    ///
    /// Starts the next state for this DFA
    ///
    /// When this is first called, the DFA will enter state 0, then state 1, etc. If this hasn't been called yet then the DFA is not
    /// in a valid state and the other calls cannot be made.
    ///
    fn start_state(&mut self);

    ///
    /// Adds a transition to the current state
    ///
    /// Any input symbol can appear exactly once in a state, and must not overlap any other input symbol. Transitions must be in input
    /// symbol order.
    ///
    fn transition(&mut self, symbol: InputSymbol, target_state: StateId);

    ///
    /// Sets the current state as an accepting state and sets the output symbol that will be produced if this is the longest match
    ///
    fn accept(&mut self, symbol: OutputSymbol);

    ///
    /// Finishes building the DFA and returns the matcher for the pattern it represents
    ///
    fn build(self) -> DfaType;
}
<|endoftext|>


<commit_before>use common::take_cell::TakeCell;
use core::cell::Cell;
use hil::{AppId,Driver,Callback,AppSlice,Shared,NUM_PROCS};
use hil::spi_master::{SpiMaster,SpiCallback};
use core::cmp;
use hil::spi_master::ClockPolarity;
use hil::spi_master::ClockPhase;

/* SPI operations are handled by coping into a kernel buffer for
 * writes and copying out of a kernel buffer for reads.
 *
 * If the application buffer is larger than the kernel buffer,
 * the driver issues multiple HAL operations. The len field
 * of an application keeps track of the length of the desired
 * operation, while the index variable keeps track of the
 * index an ongoing operation is at in the buffers. */

struct App {
    callback:  Option<Callback>,
    app_read:  Option<AppSlice<Shared, u8>>,
    app_write: Option<AppSlice<Shared, u8>>,
    len:       Cell<usize>,
    index:     Cell<usize>,
}

pub struct Spi<'a, S: SpiMaster + 'a> {
    spi_master:   &'a mut S,
    busy:         Cell<bool>,
    apps:         [TakeCell<App>; NUM_PROCS],
    kernel_read:  TakeCell<&'static mut [u8]>,
    kernel_write: TakeCell<&'static mut [u8]>,
    kernel_len:   Cell<usize>
}

impl<'a, S: SpiMaster> Spi<'a, S> {
    pub fn new(spi_master: &'a mut S) -> Spi<S> {
        Spi {
            spi_master: spi_master,
            busy: Cell::new(false),
            apps: [TakeCell::empty()],
            kernel_len: Cell::new(0),
            kernel_read : TakeCell::empty(),
            kernel_write : TakeCell::empty()
        }
    }

    pub fn config_buffers(&mut self,
                          read: &'static mut [u8],
                          write: &'static mut [u8]) {
        let len = cmp::min(read.len(), write.len());
        self.kernel_len.set(len);
        self.kernel_read.replace(read);
        self.kernel_write.replace(write);
    }

    // Assumes checks for busy/etc. already done
    // Updates app.index to be index + length of op
    fn do_next_read_write(&self, app: &mut App) {
        let start = app.index.get();
        let len = cmp::min(app.len.get() - start, self.kernel_len.get());
        let end = start + len;
        app.index.set(end);

        self.kernel_write.map(|kwbuf| {
            app.app_write.as_mut().map(|src| {
                for (i, c) in src.as_ref()[start .. end].iter().enumerate() {
                    kwbuf[i] = *c;
                }
            });
        });

        self.spi_master.read_write_bytes(self.kernel_write.take(),
                                         self.kernel_read.take(), len);
    }
}

impl<'a, S: SpiMaster> Driver for Spi<'a, S> {
    fn allow(&self, appid: AppId,
             allow_num: usize, slice: AppSlice<Shared, u8>) -> isize {
        let app = appid.idx();
        match allow_num {
            0 => {
                let appc = match self.apps[app].take() {
                    None => App {
                        callback: None,
                        app_read: Some(slice),
                        app_write: None,
                        len: Cell::new(0),
                        index: Cell::new(0),
                    },
                    Some(mut appc) => {
                        appc.app_read = Some(slice);
                        appc
                    }
                };
                self.apps[app].replace(appc);
                0
            },
            1 => {
                let appc = match self.apps[app].take() {
                    None => App {
                        callback: None,
                        app_read: None,
                        app_write: Some(slice),
                        len: Cell::new(0),
                        index: Cell::new(0),
                    },
                    Some(mut appc) => {
                        appc.app_write = Some(slice);
                        appc
                    }
                };
                self.apps[app].replace(appc);
                0
            }
            _ => -1
        }
    }

    #[inline(never)]
    fn subscribe(&self, subscribe_num: usize, callback: Callback) -> isize {
        match subscribe_num {
            0 /* read_write */ => {
                let appc = match self.apps[0].take() {
                    None => App {
                        callback: Some(callback),
                        app_read: None,
                        app_write: None,
                        len: Cell::new(0),
                        index: Cell::new(0),
                    },
                    Some(mut appc) => {
                        appc.callback = Some(callback);
                        appc
                    }
                };
                self.apps[0].replace(appc);
                0
            },
            _ => -1
        }
    }
    /*
     * 0: read/write a single byte (blocking)
     * 1: read/write buffers
     *   - requires write buffer registered with allow
     *   - read buffer optional
     * 2: set chip select
     *   - selects which peripheral (CS line) the SPI should
     *     activate
     *   - valid values are 0-3 for SAM4L
     *   - invalid value will result in CS 0
     * 3: get chip select
     *   - returns current selected peripheral
     *   - If none selected, returns 255
     * 4: set rate on current peripheral
     *   - parameter in bps
     * 5: get rate on current peripheral
     *   - value in bps
     * 6: set clock phase on current peripheral
     *   - 0 is sample leading
     *   - non-zero is sample trailing
     * 7: get clock phase on current peripheral
     *   - 0 is sample leading
     *   - non-zero is sample trailing
     * 8: set clock polarity on current peripheral
     *   - 0 is idle low
     *   - non-zero is idle high
     * 9: get clock polarity on current peripheral
     *   - 0 is idle low
     *   - non-zero is idle high
     * x: lock spi
     *   - if you perform an operation without the lock,
     *     it implicitly acquires the lock before the
     *     operation and releases it after
     *   - while an app holds the lock no other app can issue
     *     operations on SPI (they are buffered)
     * x+1: unlock spi
     *   - does nothing if lock not held
     */

    fn command(&self, cmd_num: usize, arg1: usize) -> isize {
        match cmd_num {
            0 /* read_write_byte */ => {
                self.spi_master.read_write_byte(arg1 as u8) as isize
            },
            1 /* read_write_bytes */ => {
                if self.busy.get() {
                    return -1;
                }
                let mut result = -1;
                self.apps[0].map(|app| {
                    let mut mlen = 0;
                    // If write buffer too small, return
                    app.app_write.as_mut().map(|w| {
                        mlen = w.len();
                    });
                    app.app_read.as_mut().map(|r| {
                        mlen = cmp::min(mlen, r.len());
                    });
                    if mlen >= arg1 {
                        app.len.set(arg1);
                        app.index.set(0);
                        self.busy.set(true);
                        self.do_next_read_write(app);
                        result = 0;
                    }
                });
                return result;
            }
            2 /* set chip select */ => {
                let cs = arg1 as u8;
                if cs <= 3 {
                    self.spi_master.set_chip_select(cs);
                    0
                } else {
                    -1
                }
            }
            3 /* get chip select */ => {
                self.spi_master.get_chip_select() as isize
            }
            4 /* set baud rate */ => {
                self.spi_master.set_rate(arg1 as u32) as isize
            }
            5 /* get baud rate */ => {
                self.spi_master.get_rate() as isize
            }
            6 /* set phase */ => {
                match arg1 {
                    0 => self.spi_master.set_phase(ClockPhase::SampleLeading),
                    _ => self.spi_master.set_phase(ClockPhase::SampleTrailing),
                };
                0
            }
            7 /* get phase */ => {
                self.spi_master.get_phase() as isize
            }
            8 /* set polarity */ => {
                match arg1 {
                    0 => self.spi_master.set_clock(ClockPolarity::IdleLow),
                    _ => self.spi_master.set_clock(ClockPolarity::IdleHigh),
                };
                0
            }
            9 /* get polarity */ => {
                self.spi_master.get_clock() as isize
            }
            _ => -1
        }
    }
}

#[allow(dead_code)]
fn each_some<'a, T, I, F>(lst: I, f: F)
        where T: 'static, I: Iterator<Item=&'a TakeCell<T>>, F: Fn(&mut T) {
    for item in lst {
        item.map(|i| f(i));
    }
}

impl<'a, S: SpiMaster> SpiCallback for Spi<'a, S> {
    fn read_write_done(&self,
                       writebuf: Option<&'static mut [u8]>,
                       readbuf:  Option<&'static mut [u8]>,
                       length: usize) {
        self.apps[0].map(|app| {
            if app.app_read.is_some() {
                let src = readbuf.as_ref().unwrap();
                let dest = app.app_read.as_mut().unwrap();
                let start = app.index.get() - length;
                let end = start + length;

                let d = &mut dest.as_mut()[start .. end];
                for (i, c) in src[0 .. length].iter().enumerate() {
                    d[i] = *c;
                }
            }

            self.kernel_read.put(readbuf);
            self.kernel_write.put(writebuf);

            if app.index.get() == app.len.get() {
                self.busy.set(false);
                app.len.set(0);
                app.index.set(0);
                app.callback.take().map(|mut cb| {
                    cb.schedule(app.len.get(), 0, 0);
                });
            } else {
                self.do_next_read_write(app);
            }
        });
    }
}

<commit_msg>SPI: unnecessary cells inside TakeCell<commit_after>use common::take_cell::TakeCell;
use core::cell::Cell;
use hil::{AppId,Driver,Callback,AppSlice,Shared,NUM_PROCS};
use hil::spi_master::{SpiMaster,SpiCallback};
use core::cmp;
use hil::spi_master::ClockPolarity;
use hil::spi_master::ClockPhase;

/* SPI operations are handled by coping into a kernel buffer for
 * writes and copying out of a kernel buffer for reads.
 *
 * If the application buffer is larger than the kernel buffer,
 * the driver issues multiple HAL operations. The len field
 * of an application keeps track of the length of the desired
 * operation, while the index variable keeps track of the
 * index an ongoing operation is at in the buffers. */

struct App {
    callback:  Option<Callback>,
    app_read:  Option<AppSlice<Shared, u8>>,
    app_write: Option<AppSlice<Shared, u8>>,
    len:       usize,
    index:     usize,
}

pub struct Spi<'a, S: SpiMaster + 'a> {
    spi_master:   &'a mut S,
    busy:         Cell<bool>,
    apps:         [TakeCell<App>; NUM_PROCS],
    kernel_read:  TakeCell<&'static mut [u8]>,
    kernel_write: TakeCell<&'static mut [u8]>,
    kernel_len:   Cell<usize>
}

impl<'a, S: SpiMaster> Spi<'a, S> {
    pub fn new(spi_master: &'a mut S) -> Spi<S> {
        Spi {
            spi_master: spi_master,
            busy: Cell::new(false),
            apps: [TakeCell::empty()],
            kernel_len: Cell::new(0),
            kernel_read : TakeCell::empty(),
            kernel_write : TakeCell::empty()
        }
    }

    pub fn config_buffers(&mut self,
                          read: &'static mut [u8],
                          write: &'static mut [u8]) {
        let len = cmp::min(read.len(), write.len());
        self.kernel_len.set(len);
        self.kernel_read.replace(read);
        self.kernel_write.replace(write);
    }

    // Assumes checks for busy/etc. already done
    // Updates app.index to be index + length of op
    fn do_next_read_write(&self, app: &mut App) {
        let start = app.index;
        let len = cmp::min(app.len - start, self.kernel_len.get());
        let end = start + len;
        app.index = end;

        self.kernel_write.map(|kwbuf| {
            app.app_write.as_mut().map(|src| {
                for (i, c) in src.as_ref()[start .. end].iter().enumerate() {
                    kwbuf[i] = *c;
                }
            });
        });

        self.spi_master.read_write_bytes(self.kernel_write.take(),
                                         self.kernel_read.take(), len);
    }
}

impl<'a, S: SpiMaster> Driver for Spi<'a, S> {
    fn allow(&self, appid: AppId,
             allow_num: usize, slice: AppSlice<Shared, u8>) -> isize {
        let app = appid.idx();
        match allow_num {
            0 => {
                let appc = match self.apps[app].take() {
                    None => App {
                        callback: None,
                        app_read: Some(slice),
                        app_write: None,
                        len: 0,
                        index: 0,
                    },
                    Some(mut appc) => {
                        appc.app_read = Some(slice);
                        appc
                    }
                };
                self.apps[app].replace(appc);
                0
            },
            1 => {
                let appc = match self.apps[app].take() {
                    None => App {
                        callback: None,
                        app_read: None,
                        app_write: Some(slice),
                        len: 0,
                        index: 0,
                    },
                    Some(mut appc) => {
                        appc.app_write = Some(slice);
                        appc
                    }
                };
                self.apps[app].replace(appc);
                0
            }
            _ => -1
        }
    }

    #[inline(never)]
    fn subscribe(&self, subscribe_num: usize, callback: Callback) -> isize {
        match subscribe_num {
            0 /* read_write */ => {
                let appc = match self.apps[0].take() {
                    None => App {
                        callback: Some(callback),
                        app_read: None,
                        app_write: None,
                        len: 0,
                        index: 0,
                    },
                    Some(mut appc) => {
                        appc.callback = Some(callback);
                        appc
                    }
                };
                self.apps[0].replace(appc);
                0
            },
            _ => -1
        }
    }
    /*
     * 0: read/write a single byte (blocking)
     * 1: read/write buffers
     *   - requires write buffer registered with allow
     *   - read buffer optional
     * 2: set chip select
     *   - selects which peripheral (CS line) the SPI should
     *     activate
     *   - valid values are 0-3 for SAM4L
     *   - invalid value will result in CS 0
     * 3: get chip select
     *   - returns current selected peripheral
     *   - If none selected, returns 255
     * 4: set rate on current peripheral
     *   - parameter in bps
     * 5: get rate on current peripheral
     *   - value in bps
     * 6: set clock phase on current peripheral
     *   - 0 is sample leading
     *   - non-zero is sample trailing
     * 7: get clock phase on current peripheral
     *   - 0 is sample leading
     *   - non-zero is sample trailing
     * 8: set clock polarity on current peripheral
     *   - 0 is idle low
     *   - non-zero is idle high
     * 9: get clock polarity on current peripheral
     *   - 0 is idle low
     *   - non-zero is idle high
     * x: lock spi
     *   - if you perform an operation without the lock,
     *     it implicitly acquires the lock before the
     *     operation and releases it after
     *   - while an app holds the lock no other app can issue
     *     operations on SPI (they are buffered)
     * x+1: unlock spi
     *   - does nothing if lock not held
     */

    fn command(&self, cmd_num: usize, arg1: usize) -> isize {
        match cmd_num {
            0 /* read_write_byte */ => {
                self.spi_master.read_write_byte(arg1 as u8) as isize
            },
            1 /* read_write_bytes */ => {
                if self.busy.get() {
                    return -1;
                }
                let mut result = -1;
                self.apps[0].map(|app| {
                    let mut mlen = 0;
                    // If write buffer too small, return
                    app.app_write.as_mut().map(|w| {
                        mlen = w.len();
                    });
                    app.app_read.as_mut().map(|r| {
                        mlen = cmp::min(mlen, r.len());
                    });
                    if mlen >= arg1 {
                        app.len = arg1;
                        app.index = 0;
                        self.busy.set(true);
                        self.do_next_read_write(app);
                        result = 0;
                    }
                });
                return result;
            }
            2 /* set chip select */ => {
                let cs = arg1 as u8;
                if cs <= 3 {
                    self.spi_master.set_chip_select(cs);
                    0
                } else {
                    -1
                }
            }
            3 /* get chip select */ => {
                self.spi_master.get_chip_select() as isize
            }
            4 /* set baud rate */ => {
                self.spi_master.set_rate(arg1 as u32) as isize
            }
            5 /* get baud rate */ => {
                self.spi_master.get_rate() as isize
            }
            6 /* set phase */ => {
                match arg1 {
                    0 => self.spi_master.set_phase(ClockPhase::SampleLeading),
                    _ => self.spi_master.set_phase(ClockPhase::SampleTrailing),
                };
                0
            }
            7 /* get phase */ => {
                self.spi_master.get_phase() as isize
            }
            8 /* set polarity */ => {
                match arg1 {
                    0 => self.spi_master.set_clock(ClockPolarity::IdleLow),
                    _ => self.spi_master.set_clock(ClockPolarity::IdleHigh),
                };
                0
            }
            9 /* get polarity */ => {
                self.spi_master.get_clock() as isize
            }
            _ => -1
        }
    }
}

impl<'a, S: SpiMaster> SpiCallback for Spi<'a, S> {
    fn read_write_done(&self,
                       writebuf: Option<&'static mut [u8]>,
                       readbuf:  Option<&'static mut [u8]>,
                       length: usize) {
        self.apps[0].map(|app| {
            if app.app_read.is_some() {
                let src = readbuf.as_ref().unwrap();
                let dest = app.app_read.as_mut().unwrap();
                let start = app.index - length;
                let end = start + length;

                let d = &mut dest.as_mut()[start .. end];
                for (i, c) in src[0 .. length].iter().enumerate() {
                    d[i] = *c;
                }
            }

            self.kernel_read.put(readbuf);
            self.kernel_write.put(writebuf);

            if app.index == app.len {
                self.busy.set(false);
                app.len = 0;
                app.index = 0;
                app.callback.take().map(|mut cb| {
                    cb.schedule(app.len, 0, 0);
                });
            } else {
                self.do_next_read_write(app);
            }
        });
    }
}

<|endoftext|>


<commit_before>use std::sync::Arc;

use serde::{Serialize, Deserialize};
use tokio::sync::{RwLock, Mutex, MutexGuard};

use super::state::{TurtleState, DrawingState};
use super::renderer::display_list::PrimHandle;

/// The unique ID of a particular turtle
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TurtleId(usize);

#[derive(Default, Debug)]
pub struct TurtleDrawings {
    pub state: TurtleState,
    pub drawings: Vec<PrimHandle>,
    /// If the turtle is currently filling, this is the handle to the fill polygon that new points
    /// should be appended to
    ///
    /// This handle will already be present in `drawings`, so it does not need to be added after
    /// the fill has begun.
    pub current_fill_polygon: Option<PrimHandle>,
}

/// The entire state of the application, shared between threads in the server
#[derive(Default, Debug)]
pub struct App {
    // NOTE: To avoid deadlocking, all of the code in this crate follows a consistent locking
    // order as specified below:
    //
    // 1. drawing.lock()
    // 2. turtle.lock()
    // 3. display_list.lock()
    // 4. event_loop.lock()
    //
    // Any of these steps may be omitted, but the order must always be consistent.

    /// The current state of the drawing
    drawing: Arc<Mutex<DrawingState>>,
    /// Each `TurtleId` indexes into this field
    ///
    /// Need to be very careful deleting from this field because the `TurtleId` returned from
    /// `add_turtle()` must remain unique and thus can never be repeated. Also, we wouldn't want to
    /// invalidate any clones of the `Arc<Mutex<TurtleDrawings>>` returned from the `turtle()`
    /// method.
    ///
    /// The outer `RwLock` in this type makes it possible to 1) push into the `Vec` using `write()`
    /// and 2) `clone` multiple items in the `Vec` concurrently using `read()`.
    turtles: RwLock<Vec<Arc<Mutex<TurtleDrawings>>>>,
}

impl App {
    /// Adds a new turtle to the application state, returning its `TurtleId`
    pub async fn add_turtle(&self) -> TurtleId {
        let mut turtles = self.turtles.write().await;
        let id = TurtleId(turtles.len());
        turtles.push(Default::default());
        id
    }

    pub fn drawing(&self) -> &Arc<Mutex<DrawingState>> {
        &self.drawing
    }

    /// Returns a mutable handle to the drawing state
    pub async fn drawing_mut(&self) -> MutexGuard<'_, DrawingState> {
        self.drawing.lock().await
    }

    /// Returns the total number of turtles currently stored in the application state
    pub async fn turtles_len(&self) -> usize {
        self.turtles.read().await.len()
    }

    /// Returns the IDs of all turtles currently stored in the application state
    pub async fn turtle_ids(&self) -> impl Iterator<Item=TurtleId> + Clone {
        let len = self.turtles_len().await;
        (0..len).map(TurtleId)
    }

    /// Returns a handle to a the state and drawings of the given turtle
    ///
    /// The data is not locked, so multiple callers of this method may race to lock the data after
    /// the mutex is returned.
    pub async fn turtle(&self, id: TurtleId) -> Arc<Mutex<TurtleDrawings>> {
        let TurtleId(index) = id;
        let turtles = self.turtles.read().await;
        turtles[index].clone()
    }
}
<commit_msg>Ignoring dead code warning for `drawing_mut` since it is still used in the renderer<commit_after>use std::sync::Arc;

use serde::{Serialize, Deserialize};
use tokio::sync::{RwLock, Mutex, MutexGuard};

use super::state::{TurtleState, DrawingState};
use super::renderer::display_list::PrimHandle;

/// The unique ID of a particular turtle
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TurtleId(usize);

#[derive(Default, Debug)]
pub struct TurtleDrawings {
    pub state: TurtleState,
    pub drawings: Vec<PrimHandle>,
    /// If the turtle is currently filling, this is the handle to the fill polygon that new points
    /// should be appended to
    ///
    /// This handle will already be present in `drawings`, so it does not need to be added after
    /// the fill has begun.
    pub current_fill_polygon: Option<PrimHandle>,
}

/// The entire state of the application, shared between threads in the server
#[derive(Default, Debug)]
pub struct App {
    // NOTE: To avoid deadlocking, all of the code in this crate follows a consistent locking
    // order as specified below:
    //
    // 1. drawing.lock()
    // 2. turtle.lock()
    // 3. display_list.lock()
    // 4. event_loop.lock()
    //
    // Any of these steps may be omitted, but the order must always be consistent.

    /// The current state of the drawing
    drawing: Arc<Mutex<DrawingState>>,
    /// Each `TurtleId` indexes into this field
    ///
    /// Need to be very careful deleting from this field because the `TurtleId` returned from
    /// `add_turtle()` must remain unique and thus can never be repeated. Also, we wouldn't want to
    /// invalidate any clones of the `Arc<Mutex<TurtleDrawings>>` returned from the `turtle()`
    /// method.
    ///
    /// The outer `RwLock` in this type makes it possible to 1) push into the `Vec` using `write()`
    /// and 2) `clone` multiple items in the `Vec` concurrently using `read()`.
    turtles: RwLock<Vec<Arc<Mutex<TurtleDrawings>>>>,
}

impl App {
    /// Adds a new turtle to the application state, returning its `TurtleId`
    pub async fn add_turtle(&self) -> TurtleId {
        let mut turtles = self.turtles.write().await;
        let id = TurtleId(turtles.len());
        turtles.push(Default::default());
        id
    }

    pub fn drawing(&self) -> &Arc<Mutex<DrawingState>> {
        &self.drawing
    }

    /// Returns a mutable handle to the drawing state
    #[cfg_attr(feature = "test", allow(dead_code))] // Used in renderer, but not for tests
    pub async fn drawing_mut(&self) -> MutexGuard<'_, DrawingState> {
        self.drawing.lock().await
    }

    /// Returns the total number of turtles currently stored in the application state
    pub async fn turtles_len(&self) -> usize {
        self.turtles.read().await.len()
    }

    /// Returns the IDs of all turtles currently stored in the application state
    pub async fn turtle_ids(&self) -> impl Iterator<Item=TurtleId> + Clone {
        let len = self.turtles_len().await;
        (0..len).map(TurtleId)
    }

    /// Returns a handle to a the state and drawings of the given turtle
    ///
    /// The data is not locked, so multiple callers of this method may race to lock the data after
    /// the mutex is returned.
    pub async fn turtle(&self, id: TurtleId) -> Arc<Mutex<TurtleDrawings>> {
        let TurtleId(index) = id;
        let turtles = self.turtles.read().await;
        turtles[index].clone()
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018 Stefan Lankes, RWTH Aachen University
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

use crate::arch::drop_user_space;
use crate::arch::irq::{irq_nested_disable, irq_nested_enable};
use crate::arch::switch;
use crate::consts::*;
use crate::errno::*;
use crate::logging::*;
use crate::scheduler::task::*;
use crate::synch::spinlock::*;
use alloc::collections::{BTreeMap, VecDeque};
use alloc::rc::Rc;
use core::cell::RefCell;
use core::sync::atomic::{AtomicU32, Ordering};

static NO_TASKS: AtomicU32 = AtomicU32::new(0);
static TID_COUNTER: AtomicU32 = AtomicU32::new(0);

pub struct Scheduler {
	/// task id which is currently running
	current_task: Rc<RefCell<Task>>,
	/// task id of the idle task
	idle_task: Rc<RefCell<Task>>,
	/// queue of tasks, which are ready
	ready_queue: SpinlockIrqSave<PriorityTaskQueue>,
	/// queue of tasks, which are finished and can be released
	finished_tasks: SpinlockIrqSave<VecDeque<TaskId>>,
	// map between task id and task controll block
	tasks: SpinlockIrqSave<BTreeMap<TaskId, Rc<RefCell<Task>>>>,
}

impl Scheduler {
	pub fn new() -> Scheduler {
		let tid = TaskId::from(TID_COUNTER.fetch_add(1, Ordering::SeqCst));
		let idle_task = Rc::new(RefCell::new(Task::new_idle(tid)));
		let tasks = SpinlockIrqSave::new(BTreeMap::new());

		tasks.lock().insert(tid, idle_task.clone());

		Scheduler {
			current_task: idle_task.clone(),
			idle_task: idle_task.clone(),
			ready_queue: SpinlockIrqSave::new(PriorityTaskQueue::new()),
			finished_tasks: SpinlockIrqSave::new(VecDeque::<TaskId>::new()),
			tasks: tasks,
		}
	}

	fn get_tid(&self) -> TaskId {
		loop {
			let id = TaskId::from(TID_COUNTER.fetch_add(1, Ordering::SeqCst));

			if self.tasks.lock().contains_key(&id) == false {
				return id;
			}
		}
	}

	pub fn spawn(&mut self, func: extern "C" fn(), prio: TaskPriority) -> Result<TaskId> {
		let prio_number = prio.into() as usize;

		if prio_number >= NO_PRIORITIES {
			return Err(Error::BadPriority);
		}

		// Create the new task.
		let tid = self.get_tid();
		let task = Rc::new(RefCell::new(Task::new(tid, TaskStatus::TaskReady, prio)));

		task.borrow_mut().create_stack_frame(func);

		// Add it to the task lists.
		self.ready_queue.lock().push(task.clone());
		self.tasks.lock().insert(tid, task);
		NO_TASKS.fetch_add(1, Ordering::SeqCst);

		info!("Creating task {}", tid);

		Ok(tid)
	}

	fn cleanup(&mut self) {
		// destroy user space
		drop_user_space();

		self.current_task.borrow_mut().status = TaskStatus::TaskFinished;

		// update the number of tasks
		NO_TASKS.fetch_sub(1, Ordering::SeqCst);
	}

	pub fn exit(&mut self) -> ! {
		if self.current_task.borrow().status != TaskStatus::TaskIdle {
			info!("finish task with id {}", self.current_task.borrow().id);
			self.cleanup();
		} else {
			panic!("unable to terminate idle task");
		}

		self.reschedule();

		// we should never reach this point
		panic!("exit failed!");
	}

	pub fn abort(&mut self) -> ! {
		if self.current_task.borrow().status != TaskStatus::TaskIdle {
			info!("abort task with id {}", self.current_task.borrow().id);
			self.cleanup();
		} else {
			panic!("unable to terminate idle task");
		}

		self.reschedule();

		// we should never reach this point
		panic!("abort failed!");
	}

	pub fn block_current_task(&mut self) -> Rc<RefCell<Task>> {
		if self.current_task.borrow().status == TaskStatus::TaskRunning {
			debug!("block task {}", self.current_task.borrow().id);

			self.current_task.borrow_mut().status = TaskStatus::TaskBlocked;
			self.current_task.clone()
		} else {
			panic!("unable to block task {}", self.current_task.borrow().id);
		}
	}

	pub fn wakeup_task(&mut self, task: Rc<RefCell<Task>>) {
		if task.borrow().status == TaskStatus::TaskBlocked {
			debug!("wakeup task {}", task.borrow().id);

			task.borrow_mut().status = TaskStatus::TaskReady;
			self.ready_queue.lock().push(task.clone());
		}
	}

	pub fn get_current_taskid(&self) -> TaskId {
		self.current_task.borrow().id
	}

	/// Determines the start address of the stack
	#[no_mangle]
	pub fn get_current_stack(&self) -> usize {
		unsafe { (*self.current_task.borrow().stack).bottom() }
	}

	pub fn get_root_page_table(&self) -> usize {
		self.current_task.borrow().root_page_table
	}

	pub fn set_root_page_table(&self, addr: usize) {
		self.current_task.borrow_mut().root_page_table = addr;
	}

	pub fn schedule(&mut self) {
		// do we have finished tasks? => drop tasks => deallocate implicitly the stack
		match self.finished_tasks.lock().pop_front() {
			Some(id) => {
				if self.tasks.lock().remove(&id).is_none() == true {
					info!("Unable to drop task {}", id);
				}
			}
			_ => {}
		}

		// Get information about the current task.
		let (current_id, current_stack_pointer, current_prio, current_status) = {
			let mut borrowed = self.current_task.borrow_mut();
			(
				borrowed.id,
				&mut borrowed.last_stack_pointer as *mut usize,
				borrowed.prio,
				borrowed.status,
			)
		};

		// do we have a task, which is ready?
		let mut next_task;
		if current_status == TaskStatus::TaskRunning {
			next_task = self.ready_queue.lock().pop_with_prio(current_prio);
		} else {
			next_task = self.ready_queue.lock().pop();
		}

		if next_task.is_none() == true {
			if current_status != TaskStatus::TaskRunning && current_status != TaskStatus::TaskIdle {
				debug!("Switch to idle task");
				// current task isn't able to run and no other task available
				// => switch to the idle task
				next_task = Some(self.idle_task.clone());
			}
		}

		match next_task {
			Some(new_task) => {
				let (new_id, new_stack_pointer) = {
					let mut borrowed = new_task.borrow_mut();
					borrowed.status = TaskStatus::TaskRunning;
					(borrowed.id, borrowed.last_stack_pointer)
				};

				if current_status == TaskStatus::TaskRunning {
					debug!("Add task {} to ready queue", current_id);
					self.current_task.borrow_mut().status = TaskStatus::TaskReady;
					self.ready_queue.lock().push(self.current_task.clone());
				} else if current_status == TaskStatus::TaskFinished {
					debug!("Task {} finished", current_id);
					self.current_task.borrow_mut().status = TaskStatus::TaskInvalid;
					// release the task later, because the stack is required
					// to call the function "switch"
					// => push id to a queue and release the task later
					self.finished_tasks.lock().push_back(current_id);
				}

				debug!(
					"Switching task from {} to {} (stack {:#X} => {:#X})",
					current_id,
					new_id,
					unsafe { *current_stack_pointer },
					new_stack_pointer
				);

				self.current_task = new_task;

				switch(current_stack_pointer, new_stack_pointer);
			}
			_ => {}
		}
	}

	pub fn reschedule(&mut self) {
		let flags = irq_nested_disable();
		self.schedule();
		irq_nested_enable(flags);
	}
}
<commit_msg>protect access to a task by the closure irqsave<commit_after>// Copyright (c) 2017-2018 Stefan Lankes, RWTH Aachen University
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

use crate::arch::drop_user_space;
use crate::arch::switch;
use crate::collections::irqsave;
use crate::consts::*;
use crate::errno::*;
use crate::logging::*;
use crate::scheduler::task::*;
use crate::synch::spinlock::*;
use alloc::collections::{BTreeMap, VecDeque};
use alloc::rc::Rc;
use core::cell::RefCell;
use core::sync::atomic::{AtomicU32, Ordering};

static NO_TASKS: AtomicU32 = AtomicU32::new(0);
static TID_COUNTER: AtomicU32 = AtomicU32::new(0);

pub struct Scheduler {
	/// task id which is currently running
	current_task: Rc<RefCell<Task>>,
	/// task id of the idle task
	idle_task: Rc<RefCell<Task>>,
	/// queue of tasks, which are ready
	ready_queue: SpinlockIrqSave<PriorityTaskQueue>,
	/// queue of tasks, which are finished and can be released
	finished_tasks: SpinlockIrqSave<VecDeque<TaskId>>,
	// map between task id and task controll block
	tasks: SpinlockIrqSave<BTreeMap<TaskId, Rc<RefCell<Task>>>>,
}

impl Scheduler {
	pub fn new() -> Scheduler {
		let tid = TaskId::from(TID_COUNTER.fetch_add(1, Ordering::SeqCst));
		let idle_task = Rc::new(RefCell::new(Task::new_idle(tid)));
		let tasks = SpinlockIrqSave::new(BTreeMap::new());

		tasks.lock().insert(tid, idle_task.clone());

		Scheduler {
			current_task: idle_task.clone(),
			idle_task: idle_task.clone(),
			ready_queue: SpinlockIrqSave::new(PriorityTaskQueue::new()),
			finished_tasks: SpinlockIrqSave::new(VecDeque::<TaskId>::new()),
			tasks: tasks,
		}
	}

	fn get_tid(&self) -> TaskId {
		loop {
			let id = TaskId::from(TID_COUNTER.fetch_add(1, Ordering::SeqCst));

			if self.tasks.lock().contains_key(&id) == false {
				return id;
			}
		}
	}

	pub fn spawn(&mut self, func: extern "C" fn(), prio: TaskPriority) -> Result<TaskId> {
		let closure = || {
			let prio_number = prio.into() as usize;

			if prio_number >= NO_PRIORITIES {
				return Err(Error::BadPriority);
			}

			// Create the new task.
			let tid = self.get_tid();
			let task = Rc::new(RefCell::new(Task::new(tid, TaskStatus::TaskReady, prio)));

			task.borrow_mut().create_stack_frame(func);

			// Add it to the task lists.
			self.ready_queue.lock().push(task.clone());
			self.tasks.lock().insert(tid, task);
			NO_TASKS.fetch_add(1, Ordering::SeqCst);

			info!("Creating task {}", tid);

			Ok(tid)
		};

		irqsave(closure)
	}

	fn cleanup(&mut self) {
		// destroy user space
		drop_user_space();

		self.current_task.borrow_mut().status = TaskStatus::TaskFinished;

		// update the number of tasks
		NO_TASKS.fetch_sub(1, Ordering::SeqCst);
	}

	pub fn exit(&mut self) -> ! {
		let closure = || {
			if self.current_task.borrow().status != TaskStatus::TaskIdle {
				info!("finish task with id {}", self.current_task.borrow().id);
				self.cleanup();
			} else {
				panic!("unable to terminate idle task");
			}
		};

		irqsave(closure);

		self.reschedule();

		// we should never reach this point
		panic!("exit failed!");
	}

	pub fn abort(&mut self) -> ! {
		let closure = || {
			if self.current_task.borrow().status != TaskStatus::TaskIdle {
				info!("abort task with id {}", self.current_task.borrow().id);
				self.cleanup();
			} else {
				panic!("unable to terminate idle task");
			}
		};

		irqsave(closure);

		self.reschedule();

		// we should never reach this point
		panic!("abort failed!");
	}

	pub fn block_current_task(&mut self) -> Rc<RefCell<Task>> {
		let closure = || {
			if self.current_task.borrow().status == TaskStatus::TaskRunning {
				debug!("block task {}", self.current_task.borrow().id);

				self.current_task.borrow_mut().status = TaskStatus::TaskBlocked;
				self.current_task.clone()
			} else {
				panic!("unable to block task {}", self.current_task.borrow().id);
			}
		};

		irqsave(closure)
	}

	pub fn wakeup_task(&mut self, task: Rc<RefCell<Task>>) {
		let closure = || {
			if task.borrow().status == TaskStatus::TaskBlocked {
				debug!("wakeup task {}", task.borrow().id);

				task.borrow_mut().status = TaskStatus::TaskReady;
				self.ready_queue.lock().push(task.clone());
			}
		};

		irqsave(closure);
	}

	pub fn get_current_taskid(&self) -> TaskId {
		irqsave(|| self.current_task.borrow().id)
	}

	/// Determines the start address of the stack
	#[no_mangle]
	pub fn get_current_stack(&self) -> usize {
		irqsave(|| unsafe { (*self.current_task.borrow().stack).bottom() })
	}

	pub fn get_root_page_table(&self) -> usize {
		self.current_task.borrow().root_page_table
	}

	pub fn set_root_page_table(&self, addr: usize) {
		self.current_task.borrow_mut().root_page_table = addr;
	}

	pub fn schedule(&mut self) {
		// do we have finished tasks? => drop tasks => deallocate implicitly the stack
		match self.finished_tasks.lock().pop_front() {
			Some(id) => {
				if self.tasks.lock().remove(&id).is_none() == true {
					info!("Unable to drop task {}", id);
				}
			}
			_ => {}
		}

		// Get information about the current task.
		let (current_id, current_stack_pointer, current_prio, current_status) = {
			let mut borrowed = self.current_task.borrow_mut();
			(
				borrowed.id,
				&mut borrowed.last_stack_pointer as *mut usize,
				borrowed.prio,
				borrowed.status,
			)
		};

		// do we have a task, which is ready?
		let mut next_task;
		if current_status == TaskStatus::TaskRunning {
			next_task = self.ready_queue.lock().pop_with_prio(current_prio);
		} else {
			next_task = self.ready_queue.lock().pop();
		}

		if next_task.is_none() == true {
			if current_status != TaskStatus::TaskRunning && current_status != TaskStatus::TaskIdle {
				debug!("Switch to idle task");
				// current task isn't able to run and no other task available
				// => switch to the idle task
				next_task = Some(self.idle_task.clone());
			}
		}

		match next_task {
			Some(new_task) => {
				let (new_id, new_stack_pointer) = {
					let mut borrowed = new_task.borrow_mut();
					borrowed.status = TaskStatus::TaskRunning;
					(borrowed.id, borrowed.last_stack_pointer)
				};

				if current_status == TaskStatus::TaskRunning {
					debug!("Add task {} to ready queue", current_id);
					self.current_task.borrow_mut().status = TaskStatus::TaskReady;
					self.ready_queue.lock().push(self.current_task.clone());
				} else if current_status == TaskStatus::TaskFinished {
					debug!("Task {} finished", current_id);
					self.current_task.borrow_mut().status = TaskStatus::TaskInvalid;
					// release the task later, because the stack is required
					// to call the function "switch"
					// => push id to a queue and release the task later
					self.finished_tasks.lock().push_back(current_id);
				}

				debug!(
					"Switching task from {} to {} (stack {:#X} => {:#X})",
					current_id,
					new_id,
					unsafe { *current_stack_pointer },
					new_stack_pointer
				);

				self.current_task = new_task;

				switch(current_stack_pointer, new_stack_pointer);
			}
			_ => {}
		}
	}

	pub fn reschedule(&mut self) {
		irqsave(|| self.schedule());
	}
}
<|endoftext|>


<commit_before><commit_msg>more cautious shooting<commit_after><|endoftext|>


<commit_before><commit_msg>chown: Fix clippy warning to fix CI<commit_after><|endoftext|>