webcomment/src/db.rs

use base64::engine::Engine;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::{net::IpAddr, path::Path};

pub use sled::transaction::{
	ConflictableTransactionError, ConflictableTransactionResult, TransactionError,
};
pub use typed_sled::Tree;

const DB_DIR: &str = "db";

pub type Time = u64;

pub const BASE64: base64::engine::general_purpose::GeneralPurpose =
	base64::engine::general_purpose::GeneralPurpose::new(
		&base64::alphabet::URL_SAFE,
		base64::engine::general_purpose::NO_PAD,
	);

#[derive(Clone)]
pub struct Dbs {
	pub comment: Tree<CommentId, (Comment, CommentStatus)>,
	pub comment_approved: Tree<(TopicHash, Time, CommentId), ()>,
	/// -> (client_addr, is_edit)
	pub comment_pending: Tree<(TopicHash, Time, CommentId), (Option<IpAddr>, bool)>,
	/// client_addr -> (last_mutation, mutation_count)
	pub client_mutation: Tree<IpAddr, (Time, u32)>,
}

pub fn load_dbs(path: Option<&Path>) -> Dbs {
	let db = sled::Config::new();
	let db = if let Some(path) = path {
		db.path(path.join(DB_DIR))
	} else {
		db.temporary(true)
	}
	.open()
	.expect("Cannot open db");

	Dbs {
		comment: Tree::open(&db, "comment"),
		comment_approved: Tree::open(&db, "comment_approved"),
		comment_pending: Tree::open(&db, "comment_pending"),
		client_mutation: Tree::open(&db, "client_mutation"),
	}
}

#[repr(u8)]
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub enum CommentStatus {
	Pending = 0,
	Approved = 1,
	ApprovedEdited(Comment) = 2,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct Comment {
	pub author: String,
	pub email: Option<String>,
	pub last_edit_time: Option<u64>,
	pub mutation_token: MutationToken,
	pub post_time: u64,
	pub text: String,
	pub topic_hash: TopicHash,
}

#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct MutationToken(pub [u8; 16]);

impl MutationToken {
	pub fn new() -> Self {
		Self(rand::random())
	}

	pub fn to_base64(&self) -> String {
		BASE64.encode(self.0)
	}

	pub fn from_base64(s: &str) -> Result<Self, base64::DecodeError> {
		std::panic::catch_unwind(|| {
			let mut buf = [0; 16];
			BASE64.decode_slice_unchecked(s.as_bytes(), &mut buf)?;
			Ok(Self(buf))
		})
		.map_err(|_| base64::DecodeError::InvalidLength)?
	}
}

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

#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct TopicHash(pub [u8; 32]);

impl TopicHash {
	pub fn from_topic(topic: &str) -> Self {
		let mut hasher = Sha256::new();
		hasher.update(topic.as_bytes());
		Self(hasher.finalize().into())
	}
}

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

#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct CommentId(pub [u8; 16]);

impl CommentId {
	pub fn new() -> Self {
		Self(rand::random())
	}

	pub fn zero() -> Self {
		Self([0; 16])
	}

	pub fn max() -> Self {
		Self([255; 16])
	}

	pub fn to_base64(&self) -> String {
		BASE64.encode(self.0)
	}

	pub fn from_base64(s: &str) -> Result<Self, base64::DecodeError> {
		std::panic::catch_unwind(|| {
			let mut buf = [0; 16];
			BASE64.decode_slice_unchecked(s.as_bytes(), &mut buf)?;
			Ok(Self(buf))
		})
		.map_err(|_| base64::DecodeError::InvalidLength)?
	}
}

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

#[cfg(test)]
mod test {
	use super::*;

	#[test]
	fn test_typed_sled() {
		let db = sled::Config::new().temporary(true).open().unwrap();
		let tree = typed_sled::Tree::<(u32, u32), ()>::open(&db, "test");
		tree.insert(&(123, 456), &()).unwrap();
		tree.flush().unwrap();
		let mut iter = tree.range((123, 0)..(124, 0));
		//let mut iter = tree.iter();
		assert_eq!(iter.next(), Some(Ok(((123, 456), ()))));
	}

	#[test]
	fn test_comment_id_base64() {
		for _ in 0..10 {
			let comment_id = CommentId::new();
			assert_eq!(
				CommentId::from_base64(&comment_id.to_base64()),
				Ok(comment_id)
			);
		}
	}

	#[test]
	fn test_from_base64_dont_panic() {
		assert_eq!(CommentId::from_base64("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"), Err(base64::DecodeError::InvalidLength));
	}
}