Talaria/native/src/db/mod.rs

mod model;
mod utils;

use model::Session;
use sqlite::{Connection, Error, State};
use std::net::IpAddr;
use utils::{fetch_many, fetch_one};

fn init_schema(conn: &Connection) -> Result<(), Error> {
    let version = match fetch_one(&conn, "SELECT version FROM version LIMIT 1", |stmt| {
        stmt.read::<i64>(0)
    }) {
        Ok(Some(version)) => version,
        _ => 0,
    };

    if version == 0 {
        conn.execute(
            "
            BEGIN TRANSACTION;
            CREATE TABLE version (version INTEGER NOT NULL);
            CREATE TABLE session (
                user_id INTEGER,
                dc_id INTEGER,
                bot INTEGER,
                pts INTEGER,
                qts INTEGER,
                seq INTEGER,
                date INTEGER
            );
            CREATE TABLE channel (
                session_id INTEGER NOT NULL REFERENCES session (rowid),
                id INTEGER NOT NULL,
                hash INTEGER NOT NULL,
                pts INTEGER NOT NULL
            );
            CREATE TABLE datacenter (
                session_id INTEGER NOT NULL REFERENCES session (rowid),
                id INTEGER NOT NULL,
                ipv4 TEXT,
                ipv6 TEXT,
                port INTEGER NOT NULL,
                auth BLOB,
                CONSTRAINT SingleIp CHECK(
                    (ipv4 IS NOT NULL AND ipv6 IS NULL) OR
                    (ipv6 IS NOT NULL AND ipv4 IS NULL))
            );
            INSERT INTO version VALUES (1);
            COMMIT;
            ",
        )?;
    }

    Ok(())
}

pub fn init_connection(db_path: &str) -> Result<Connection, Error> {
    let conn = sqlite::open(db_path)?;
    init_schema(&conn)?;
    Ok(conn)
}

pub fn get_sessions(conn: &Connection) -> Result<Vec<Session>, Error> {
    let query = "
        SELECT s.rowid, s.*, COALESCE(d.ipv4, d.ipv6), d.port, d.auth
        FROM session s
        LEFT JOIN datacenter d ON d.session_id = s.rowid AND d.id = s.dc_id
    ";
    fetch_many(conn, query, |stmt| {
        Ok(Session {
            id: stmt.read(0)?,
            user_id: stmt.read(1)?,
            dc_id: stmt.read::<Option<i64>>(2)?.map(|x| x as _),
            bot: stmt.read::<Option<i64>>(3)?.map(|x| x != 0),
            pts: stmt.read::<Option<i64>>(4)?.map(|x| x as _),
            qts: stmt.read::<Option<i64>>(5)?.map(|x| x as _),
            seq: stmt.read::<Option<i64>>(6)?.map(|x| x as _),
            date: stmt.read::<Option<i64>>(7)?.map(|x| x as _),
            dc_addr: stmt.read::<Option<String>>(8)?,
            dc_port: stmt.read::<Option<i64>>(9)?.map(|x| x as _),
            dc_auth: stmt
                .read::<Option<Vec<u8>>>(10)?
                .map(|x| x.try_into().unwrap()),
        })
    })
}

pub fn create_session(conn: &Connection) -> Result<Session, Error> {
    conn.execute("INSERT INTO session DEFAULT VALUES;")?;
    let id = fetch_one(conn, "SELECT LAST_INSERT_ROWID()", |stmt| {
        stmt.read::<i64>(0)
    })?
    .unwrap();
    Ok(Session {
        id,
        user_id: None,
        dc_id: None,
        bot: None,
        pts: None,
        qts: None,
        seq: None,
        date: None,
        dc_addr: None,
        dc_port: None,
        dc_auth: None,
    })
}

pub fn update_session(conn: &Connection, session: &Session) -> Result<(), Error> {
    let mut stmt = conn
        .prepare(
            "
            UPDATE session SET
                user_id = ?,
                dc_id = ?,
                bot = ?,
                pts = ?,
                qts = ?,
                seq = ?,
                date = ?
            WHERE rowid = ?
            ",
        )?
        .bind(1, session.user_id)?
        .bind(2, session.dc_id.map(|x| x as i64))?
        .bind(3, session.bot.map(|x| x as i64))?
        .bind(4, session.pts.map(|x| x as i64))?
        .bind(5, session.qts.map(|x| x as i64))?
        .bind(6, session.seq.map(|x| x as i64))?
        .bind(7, session.date.map(|x| x as i64))?
        .bind(8, session.id)?;
    while let State::Row = stmt.next()? {}

    match (
        session.dc_id,
        session.dc_addr.as_ref(),
        session.dc_port,
        session.dc_auth,
    ) {
        (Some(id), Some(addr), Some(port), Some(auth)) => {
            let (ipv4, ipv6) = match addr.parse().unwrap() {
                IpAddr::V4(ipv4) => (Some(ipv4.to_string()), None),
                IpAddr::V6(ipv6) => (None, Some(ipv6.to_string())),
            };

            let mut stmt = conn
                .prepare(
                    "
                    DELETE FROM datacenter WHERE session_id = ? AND id = ?
                    ",
                )?
                .bind(1, session.id)?
                .bind(2, id as i64)?;

            while let State::Row = stmt.next()? {}

            let mut stmt = conn
                .prepare("INSERT INTO datacenter VALUES (?, ?, ?, ?, ?, ?)")?
                .bind(1, session.id)?
                .bind(2, id as i64)?
                .bind(3, ipv4.as_deref())?
                .bind(4, ipv6.as_deref())?
                .bind(5, port as i64)?
                .bind(6, auth.as_ref())?;

            while let State::Row = stmt.next()? {}
        }
        _ => {}
    }

    Ok(())
}