OpenTally/src/sharandom.rs

/*  OpenTally: Open-source election vote counting
 *  Copyright © 2021 Lee Yingtong Li (RunasSudo)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use ibig::UBig;
use sha2::{Digest, Sha256};

/// Deterministic random number generator using SHA256
#[derive(Clone)]
pub struct SHARandom<'r> {
	seed: &'r str,
	counter: usize,
}

impl<'r> SHARandom<'r> {
	/// Return a new [SHARandom] with the given seed
	pub fn new(seed: &'r str) -> Self {
		Self {
			seed: seed,
			counter: 0,
		}
	}
	
	/// Draw a random number *n*, such that 0 ≤ *n* < `max`
	pub fn next(&mut self, max: usize) -> usize {
		let mut hasher = Sha256::new();
		hasher.update(format!("{},{}", self.seed, self.counter).as_bytes());
		
		self.counter += 1;
		
		let hash = hasher.finalize();
		let hash = UBig::from_be_bytes(&hash);
		
		if hash >= UBig::from(2_u32).pow(256) / UBig::from(max) * UBig::from(max) {
			return self.next(max);
		}
		
		return (hash % UBig::from(max)).to_string().parse().unwrap();
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	
	#[test]
	fn sharandom1() {
		let mut random = SHARandom::new(&"foobar");
		assert_eq!(random.next(10), 0);
		assert_eq!(random.next(42), 30);
		assert_eq!(random.next(64), 13);
	}
}