2 files changed, 51 insertions, 12 deletions

diff --git a/src/db.rs b/src/db.rs
index 053142a..0c245cc 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -37,16 +37,22 @@ impl Token {
     }
 }
 
+fn default_is_admin() -> bool {
+    false
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct User {
     name: String,
     hashed_password: String,
+    #[serde(default = "default_is_admin")]
+    is_admin: bool,
     id: UID,
     sessions: HashMap<String, String>,
     tokovec: Vec<Token>,
 }
 impl User {
-    pub fn new(name: String, password: String, id: UID) -> User {
+    pub fn new(name: String, password: String, id: UID, admin: bool) -> User {
         let base_tokens = vec![
             Token::new(Color::White, 2),
             Token::new(Color::Red, 2),
@@ -57,7 +63,7 @@ impl User {
 
         let hashed_password = User::hash(&password);
 
-        User { name, hashed_password, id, tokovec: base_tokens, sessions: HashMap::new() }
+        User { name, is_admin: admin, hashed_password, id, tokovec: base_tokens, sessions: HashMap::new() }
     }
 
     fn update_name(&mut self, new_name: String) {
@@ -157,6 +163,7 @@ pub struct DB {
     pub uid_generator: uid::Generator,
     users: Vec<User>,
     config: Config,
+    registration_keys: Vec<String>,
 }
 impl DB {
     async fn save(&self) -> Result<(), String> {
@@ -193,7 +200,7 @@ impl DB {
         }
     }
     pub fn new(config: Config) -> Self {
-        DB { uid_generator: uid::Generator::new(), users: vec![], config }
+        DB { uid_generator: uid::Generator::new(), users: vec![], config, registration_keys: vec!["ADMIN".into()] }
     }
 
     pub fn get_user(&self, id: UID) -> Result<&User, String> {
@@ -272,12 +279,43 @@ impl DB {
         }
     }
 
-    pub async fn new_user(&mut self, name: String, password: String, id: UID) -> Result<User, String> {
-        let user = User::new(name, password, id);
-        self.users.push(user.clone());
-        match self.save().await {
-            Ok(_) => Ok(user),
-            Err(n) => Err(n),
+    pub async fn use_key(&mut self, key: &String) -> Result<(), String> {
+        let mut result = Err("Could not find key".into());
+        let new_vec = self.registration_keys.clone().into_iter().filter_map(|k| {
+            if *key == k {
+                result = Ok(());
+                None
+            } else {
+                Some(k)
+            }
+        }).collect();
+
+        self.registration_keys = new_vec;
+        result
+    }
+
+    pub fn add_key(&mut self, key: &String) {
+        self.registration_keys.push(key.clone())
+    }
+
+    pub async fn new_user(&mut self, name: String, password: String, id: UID, key: &String) -> Result<User, String> {
+        if self.use_key(key).await.is_ok() {
+            let mut is_admin = false;
+            if key == "ADMIN" {
+                is_admin = true;
+            }
+            let user = User::new(name, password, id, is_admin);
+    
+            self.users.push(user.clone());
+            match self.save().await {
+                Ok(_) => Ok(user),
+                Err(n) => {
+                    self.add_key(key);
+                    Err(n)
+                },
+            }
+        } else {
+            Err("Invalid key".into())
         }
         
     }
diff --git a/src/main.rs b/src/main.rs
index 8ddd853..63a497d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -77,14 +77,15 @@ fn default_id() -> UID {
     uid::EMPTY_UID
 }
 #[derive(Deserialize)]
-struct UserForm {
+struct RegisterForm {
     name: String,
     password: String,
     #[serde(default = "default_id")]
     id: UID,
+    key: String,
 }
 #[post("/register", data="<data>", format="json")]
-async fn new_user(data: Json<UserForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<UserOut>, Json<String>>) {
+async fn new_user(data: Json<RegisterForm>, db: &State<Mutex<DB>>) -> (Status, Result<Json<UserOut>, Json<String>>) {
     let mut db = db.lock().await;
     let id;
 
@@ -119,7 +120,7 @@ async fn new_user(data: Json<UserForm>, db: &State<Mutex<DB>>) -> (Status, Resul
     match id {
         None => return (Status::InternalServerError, Err("".to_string().into())),
         Some(n) => {
-            match db.new_user(name.into(), data.password.clone(), n).await {
+            match db.new_user(name.into(), data.password.clone(), n, &data.key).await {
                 // User has been created
                 Ok(u) => return (Status::Created, Ok(Json(u.into()))),
                 // Could not create user