kkv/
server.rs

1/**************************************************************************
2 *   src/server.rs  --  This file is part of kkv.                         *
3 *                                                                        *
4 *   Copyright (C) 2025 Mateo Lafalce                                     *
5 *                                                                        *
6 *   kkv is free software: you can redistribute it and/or modify          *
7 *   it under the terms of the GNU General Public License as published    *
8 *   by the Free Software Foundation, either version 3 of the License,    *
9 *   or (at your option) any later version.                               *
10 *                                                                        *
11 *   kkv is distributed in the hope that it will be useful,               *
12 *   but WITHOUT ANY WARRANTY; without even the implied warranty          *
13 *   of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.              *
14 *   See the GNU General Public License for more details.                 *
15 *                                                                        *
16 *   You should have received a copy of the GNU General Public License    *
17 *   along with this program.  If not, see http://www.gnu.org/licenses/.  *
18 *                                                                        *
19 **************************************************************************/
20
21use crate::{
22    axum::ValidatedForm,
23    message::Message,
24    types::{
25        DataCollection, HexPubkey, IndexMsg, KeyExchange, Keypair, OnionPubkey, PoW, RequestPubkey,
26        SentMessage,
27    },
28    utils::{
29        derive_key, deserialize_keys, get_last_elements, get_publickey_bytes, new_nonce,
30        sign_message,
31    },
32};
33use argon2::{Argon2, PasswordHash, PasswordVerifier};
34use axum::{
35    extract::State,
36    http::{header::AUTHORIZATION, HeaderMap, StatusCode},
37    routing::{get, post},
38    Form, Json, Router,
39};
40use blake3_pow::verify;
41use chacha20poly1305::{aead::Aead, ChaCha20Poly1305, KeyInit};
42use colored::*;
43use ed25519_dalek::{Signature, SigningKey, Verifier, VerifyingKey};
44use std::{
45    collections::HashMap,
46    io::Error,
47    sync::{Arc, Mutex},
48    time::{Duration, SystemTime, UNIX_EPOCH},
49};
50use tokio::net::TcpListener;
51use x25519_dalek::{PublicKey, ReusableSecret};
52
53/// Port where kkv works
54pub const PORT: &str = "45069";
55
56/// State of the app
57#[derive(Clone)]
58pub struct AppState {
59    valid_pubkey_for_ms: Arc<Mutex<HashMap<OnionPubkey, Keypair>>>,
60    messages: Arc<Mutex<Vec<Message>>>,
61    max_per_day: usize,
62    cost_pow: u32,
63    expected_time: u64,
64    ed25519_bytes: [u8; 64],
65    hashed_psw: String,
66}
67
68impl AppState {
69    fn new(max_per_day: usize, cost_pow: u32, expected_time: u64, hashed_psw: String) -> Self {
70        let mut csprng = rand::rngs::OsRng;
71        let signing_key: SigningKey = SigningKey::generate(&mut csprng);
72        // it's secure since it created here (it's not a )
73        let ed25519_bytes = signing_key.to_keypair_bytes();
74        AppState {
75            valid_pubkey_for_ms: Arc::new(Mutex::new(HashMap::new())),
76            messages: Arc::new(Mutex::new(Vec::new())),
77            max_per_day,
78            cost_pow,
79            expected_time,
80            ed25519_bytes,
81            hashed_psw,
82        }
83    }
84}
85
86pub async fn server(
87    max_per_day: usize,
88    cost_pow: u32,
89    expected_time: u64,
90    hashed_psw: String,
91) -> Result<(), Error> {
92    let state = AppState::new(max_per_day, cost_pow, expected_time, hashed_psw);
93    let app: Router = Router::new()
94        .route("/get_domain_pk", get(get_domain_pk))
95        .route("/pow", post(pow))
96        .route("/exchange_keys", post(exchange_keys))
97        .route("/msg", post(msg))
98        .route("/status", get(status))
99        .route("/get_msg", post(get_msg))
100        .route("/read_msg", post(read_msg))
101        .route("/del_msg", post(del_msg))
102        .with_state(state);
103
104    let listener: TcpListener = match TcpListener::bind(format!("127.0.0.1:{PORT}")).await {
105        Ok(listener) => listener,
106        Err(e) => {
107            println!("{}", e.to_string().red());
108            return Ok(());
109        }
110    };
111
112    match axum::serve(listener, app).await {
113        Ok(_) => (),
114        Err(e) => {
115            println!("{}", e.to_string().red());
116            return Ok(());
117        }
118    };
119
120    Ok(())
121}
122
123pub async fn status(method: HeaderMap, State(state): State<AppState>) -> StatusCode {
124    let client_psw = &method[AUTHORIZATION].as_bytes();
125    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();
126
127    if Argon2::default()
128        .verify_password(client_psw, &parsed_hash)
129        .is_err()
130    {
131        return StatusCode::UNAUTHORIZED;
132    }
133
134    let mut valids = match state.valid_pubkey_for_ms.lock() {
135        Ok(guard) => guard,
136        Err(_) => return StatusCode::UNAUTHORIZED,
137    };
138    update_senders_hashmap(&mut valids);
139    StatusCode::OK
140}
141
142pub async fn del_msg(
143    method: HeaderMap,
144    State(state): State<AppState>,
145    Form(input): Form<IndexMsg>,
146) -> StatusCode {
147    let client_psw = &method[AUTHORIZATION].as_bytes();
148    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();
149
150    if Argon2::default()
151        .verify_password(client_psw, &parsed_hash)
152        .is_err()
153    {
154        return StatusCode::UNAUTHORIZED;
155    }
156
157    let mut valids = match state.valid_pubkey_for_ms.lock() {
158        Ok(guard) => guard,
159        Err(_) => return StatusCode::UNAUTHORIZED,
160    };
161    update_senders_hashmap(&mut valids);
162
163    let index = input.get_index();
164
165    let actual_size = match state.messages.lock() {
166        Ok(guard) => guard.len(),
167        Err(_) => return StatusCode::UNAUTHORIZED,
168    };
169
170    if index >= actual_size {
171        return StatusCode::UNAUTHORIZED;
172    }
173
174    match state.messages.lock() {
175        Ok(mut guard) => {
176            guard.remove(index);
177            StatusCode::OK
178        }
179        Err(_) => StatusCode::UNAUTHORIZED,
180    }
181}
182
183pub async fn read_msg(
184    method: HeaderMap,
185    State(state): State<AppState>,
186    Form(input): Form<IndexMsg>,
187) -> StatusCode {
188    let client_psw = &method[AUTHORIZATION].as_bytes();
189    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();
190
191    if Argon2::default()
192        .verify_password(client_psw, &parsed_hash)
193        .is_err()
194    {
195        return StatusCode::UNAUTHORIZED;
196    }
197
198    let mut valids = match state.valid_pubkey_for_ms.lock() {
199        Ok(guard) => guard,
200        Err(_) => return StatusCode::UNAUTHORIZED,
201    };
202    update_senders_hashmap(&mut valids);
203
204    let index = input.get_index();
205
206    let actual_size = match state.messages.lock() {
207        Ok(guard) => guard.len(),
208        Err(_) => return StatusCode::UNAUTHORIZED,
209    };
210
211    if index >= actual_size {
212        return StatusCode::UNAUTHORIZED;
213    }
214
215    match state.messages.lock() {
216        Ok(mut guard) => {
217            guard[index].read_message();
218            StatusCode::OK
219        }
220        Err(_) => StatusCode::UNAUTHORIZED,
221    }
222}
223
224pub async fn get_msg(
225    method: HeaderMap,
226    State(state): State<AppState>,
227    ValidatedForm(input): ValidatedForm<IndexMsg>,
228) -> Json<DataCollection> {
229    let client_psw = &method[AUTHORIZATION].as_bytes();
230    let parsed_hash = PasswordHash::new(&state.hashed_psw).unwrap();
231
232    if Argon2::default()
233        .verify_password(client_psw, &parsed_hash)
234        .is_err()
235    {
236        return Json(DataCollection::default());
237    }
238
239    let mut valids = match state.valid_pubkey_for_ms.lock() {
240        Ok(guard) => guard,
241        Err(_) => return Json(DataCollection::default()),
242    };
243    update_senders_hashmap(&mut valids);
244
245    let index = input.get_index();
246
247    match state.messages.lock() {
248        Ok(frame_of_msg) => {
249            let last = get_last_elements(frame_of_msg.to_vec(), index);
250            let data = DataCollection::new(frame_of_msg.len(), last);
251            Json(data)
252        }
253        Err(_) => Json(DataCollection::default()),
254    }
255}
256
257pub async fn msg(
258    State(state): State<AppState>,
259    ValidatedForm(input): ValidatedForm<SentMessage>,
260) -> StatusCode {
261    let mut valid_pubkey_for_ms = match state.valid_pubkey_for_ms.lock() {
262        Ok(guard) => guard,
263        Err(_) => return StatusCode::UNAUTHORIZED,
264    };
265    update_senders_hashmap(&mut valid_pubkey_for_ms);
266
267    let pubkey = match hex::decode(input.get_pubkey()) {
268        Ok(bytes) => bytes,
269        Err(_) => return StatusCode::UNAUTHORIZED,
270    };
271
272    let chiphertext = match hex::decode(input.get_chiphertext()) {
273        Ok(bytes) => bytes,
274        Err(_) => return StatusCode::UNAUTHORIZED,
275    };
276
277    let signature = match hex::decode(input.get_signature()) {
278        Ok(bytes) => bytes,
279        Err(_) => return StatusCode::UNAUTHORIZED,
280    };
281
282    let alice_pk: [u8; 32] = match pubkey.try_into() {
283        Ok(alice_pk) => alice_pk,
284        Err(_) => return StatusCode::UNAUTHORIZED,
285    };
286
287    if valid_pubkey_for_ms.get(&alice_pk).is_none() {
288        return StatusCode::UNAUTHORIZED;
289    }
290
291    let client_verifying_key = match VerifyingKey::from_bytes(&alice_pk) {
292        Ok(client_verifying_key) => client_verifying_key,
293        Err(_) => return StatusCode::UNAUTHORIZED,
294    };
295
296    let chiphertext_sig_bytes: [u8; 64] = match signature.try_into() {
297        Ok(chiphertext_sig) => chiphertext_sig,
298        Err(_) => return StatusCode::UNAUTHORIZED,
299    };
300    let chiphertext_sig = Signature::from_bytes(&chiphertext_sig_bytes);
301
302    if client_verifying_key
303        .verify(&chiphertext, &chiphertext_sig)
304        .is_err()
305    {
306        return StatusCode::UNAUTHORIZED;
307    }
308
309    let secretkey = match valid_pubkey_for_ms.get(&alice_pk) {
310        Some(secretkey) => secretkey.get_bob_diffie_hellman_sk(),
311        None => return StatusCode::UNAUTHORIZED,
312    };
313
314    let alice_pubkey = match valid_pubkey_for_ms.get(&alice_pk) {
315        Some(alice_pubkey) => alice_pubkey.get_alice_diffie_hellman_pk(),
316        None => return StatusCode::UNAUTHORIZED,
317    };
318
319    let shared_secret = secretkey.diffie_hellman(&alice_pubkey);
320    let fixed_12bytes: &[u8; 12] = shared_secret.as_bytes()[0..12]
321        .try_into()
322        .unwrap_or(&[0; 12]);
323    let shared_nonce = chacha20poly1305::Nonce::from_slice(fixed_12bytes);
324
325    let key = derive_key(shared_secret.as_bytes());
326    let cipher = ChaCha20Poly1305::new(&key);
327
328    let decrypted_msg = match cipher.decrypt(shared_nonce, chiphertext.as_ref()) {
329        Ok(result) => result,
330        Err(_) => return StatusCode::UNAUTHORIZED,
331    };
332
333    // if we doesn't have the signature, then revoke the method
334    if decrypted_msg.len() < 64 {
335        return StatusCode::UNAUTHORIZED;
336    }
337
338    let (signature_bytes, message_bytes) = decrypted_msg.split_at(64);
339    let signature_fixed: [u8; 64] = match signature_bytes.try_into() {
340        Ok(signature_fixed) => signature_fixed,
341        Err(_) => return StatusCode::UNAUTHORIZED,
342    };
343    let signature = Signature::from_bytes(&signature_fixed);
344
345    let filtered_bytes: Vec<u8> = message_bytes
346        .iter()
347        .copied()
348        .filter(|&byte| byte != 0)
349        .collect();
350    let text = match String::from_utf8(filtered_bytes.clone()) {
351        Ok(bytes) => bytes,
352        Err(_) => return StatusCode::UNAUTHORIZED,
353    };
354
355    if client_verifying_key
356        .verify(&filtered_bytes, &signature)
357        .is_err()
358    {
359        return StatusCode::UNAUTHORIZED;
360    }
361
362    let new_message = match Message::new(alice_pk, signature_fixed, text) {
363        Ok(result) => result,
364        Err(_) => return StatusCode::UNAUTHORIZED,
365    };
366    let mut messages = match state.messages.lock() {
367        Ok(messages) => messages,
368        Err(_) => return StatusCode::UNAUTHORIZED,
369    };
370    if messages.len() >= state.max_per_day {
371        return StatusCode::INSUFFICIENT_STORAGE;
372    }
373
374    messages.push(new_message.clone());
375    valid_pubkey_for_ms.remove(&alice_pk);
376    StatusCode::OK
377}
378
379pub async fn exchange_keys(
380    State(state): State<AppState>,
381    ValidatedForm(input): ValidatedForm<KeyExchange>,
382) -> Json<KeyExchange> {
383    let mut valids = match state.valid_pubkey_for_ms.lock() {
384        Ok(guard) => guard,
385        Err(_) => return Json(KeyExchange::default()),
386    };
387    update_senders_hashmap(&mut valids);
388
389    let pubkey = match hex::decode(input.get_ed25519_pubkey()) {
390        Ok(pubkey) => pubkey,
391        Err(_) => return Json(KeyExchange::default()),
392    };
393    let nonce = match hex::decode(input.get_nonce()) {
394        Ok(nonce) => nonce,
395        Err(_) => return Json(KeyExchange::default()),
396    };
397    let pow = match hex::decode(input.get_pow()) {
398        Ok(pow) => pow,
399        Err(_) => return Json(KeyExchange::default()),
400    };
401    let pow_sig = match hex::decode(input.get_pow_sig()) {
402        Ok(pow_sig) => pow_sig,
403        Err(_) => return Json(KeyExchange::default()),
404    };
405    let diffie_hellman = match hex::decode(input.get_diffie_hellman()) {
406        Ok(diffie_hellman) => diffie_hellman,
407        Err(_) => return Json(KeyExchange::default()),
408    };
409    let diffie_hellman_sig = match hex::decode(input.get_diffie_hellman_sig()) {
410        Ok(diffie_hellman_sig) => diffie_hellman_sig,
411        Err(_) => return Json(KeyExchange::default()),
412    };
413
414    let alice_pk: [u8; 32] = match pubkey.try_into() {
415        Ok(alice_pk) => alice_pk,
416        Err(_) => return Json(KeyExchange::default()),
417    };
418    let client_verifying_key = match VerifyingKey::from_bytes(&alice_pk) {
419        Ok(verifying_key) => verifying_key,
420        Err(_) => return Json(KeyExchange::default()),
421    };
422
423    let nonce: [u8; 32] = match nonce.try_into() {
424        Ok(nonce) => nonce,
425        Err(_) => return Json(KeyExchange::default()),
426    };
427
428    let bob_pk = get_publickey_bytes(state.ed25519_bytes);
429    let server_verifying_key = match VerifyingKey::from_bytes(&bob_pk) {
430        Ok(verifying_key) => verifying_key,
431        Err(_) => return Json(KeyExchange::default()),
432    };
433
434    let diffie_hellman_sig_bytes: [u8; 64] = match diffie_hellman_sig.try_into() {
435        Ok(diffie_hellman_sig_bytes) => diffie_hellman_sig_bytes,
436        Err(_) => return Json(KeyExchange::default()),
437    };
438    let diffie_hellman_sig = Signature::from_bytes(&diffie_hellman_sig_bytes);
439
440    let pow_sig_bytes: [u8; 64] = match pow_sig.try_into() {
441        Ok(pow_sig_bytes) => pow_sig_bytes,
442        Err(_) => return Json(KeyExchange::default()),
443    };
444    let pow_sig = Signature::from_bytes(&pow_sig_bytes);
445
446    let diffie_hellman_fixed: [u8; 32] = match diffie_hellman.clone().try_into() {
447        Ok(diffie_hellman_fixed) => diffie_hellman_fixed,
448        Err(_) => return Json(KeyExchange::default()),
449    };
450    let alice_diffie_hellman_pk = PublicKey::from(diffie_hellman_fixed);
451
452    let mut dh_to_verify = Vec::new();
453    dh_to_verify.extend_from_slice(&diffie_hellman);
454    dh_to_verify.extend_from_slice(&nonce);
455
456    let pow_fixed: [u8; 76] = match pow.clone().try_into() {
457        Ok(pow) => pow,
458        Err(_) => return Json(KeyExchange::default()),
459    };
460
461    let (pk_bytes, timestamp, bytes, cost) = deserialize_keys(&pow_fixed);
462
463    let now = SystemTime::now();
464    let duration_since_epoch = now
465        .duration_since(UNIX_EPOCH)
466        .unwrap_or(Duration::new(0, 0));
467    let current_timestamp = duration_since_epoch.as_secs();
468
469    if client_verifying_key
470        .verify(&dh_to_verify, &diffie_hellman_sig)
471        .is_err()
472        || server_verifying_key.verify(&pow, &pow_sig).is_err()
473        || pk_bytes != alice_pk
474        || current_timestamp > timestamp
475        || !verify(&bytes, nonce, cost)
476    {
477        return Json(KeyExchange::default());
478    }
479
480    let diffie_hellman_sk = ReusableSecret::random();
481    let diffie_hellman = PublicKey::from(&diffie_hellman_sk).to_bytes();
482    let hex_diffie_hellman = hex::encode(diffie_hellman);
483
484    let diffie_hellman_sig = sign_message(&diffie_hellman, state.ed25519_bytes);
485    let hex_diffie_hellman_sig = hex::encode(diffie_hellman_sig);
486
487    let shared_secret = diffie_hellman_sk.diffie_hellman(&alice_diffie_hellman_pk);
488
489    let fixed_12bytes: &[u8; 12] = shared_secret.as_bytes()[0..12]
490        .try_into()
491        .unwrap_or(&[0; 12]);
492    let shared_nonce = chacha20poly1305::Nonce::from_slice(fixed_12bytes);
493
494    let hex_nonce = hex::encode(nonce);
495
496    let now = SystemTime::now();
497    let duration_since_epoch = now
498        .duration_since(UNIX_EPOCH)
499        .unwrap_or(Duration::new(0, 0));
500    let unix_timestamp_seconds = duration_since_epoch.as_secs();
501    let delta = Duration::new(state.expected_time, 0);
502    let end_timestamp = unix_timestamp_seconds + delta.as_secs();
503
504    let keypair = Keypair::new(
505        diffie_hellman_sk,
506        alice_diffie_hellman_pk,
507        *shared_nonce,
508        hex_nonce,
509        end_timestamp,
510    );
511
512    // this scenario doesnt look posible, but it's better
513    // match every edge case :)
514    if valids.len() == 50 {
515        return Json(KeyExchange::default());
516    }
517
518    valids.insert(alice_pk, keypair);
519
520    let hex_pubkey = hex::encode([0; 32]);
521    let hex_null_nonce = hex::encode([0; 32]);
522    let hex_null_pow = hex::encode([0; 32]);
523    let hex_null_pow_sig = hex::encode([0; 80]);
524
525    Json(KeyExchange::new(
526        hex_pubkey,
527        hex_null_nonce,
528        hex_null_pow,
529        hex_null_pow_sig,
530        hex_diffie_hellman,
531        hex_diffie_hellman_sig,
532    ))
533}
534
535pub async fn get_domain_pk(State(state): State<AppState>) -> HexPubkey {
536    let mut valid_pubkey_for_ms = match state.valid_pubkey_for_ms.lock() {
537        Ok(guard) => guard,
538        Err(_) => return hex::encode([0; 32]),
539    };
540    update_senders_hashmap(&mut valid_pubkey_for_ms);
541    let bob_pubkey_bytes: [u8; 32] = get_publickey_bytes(state.ed25519_bytes);
542    hex::encode(bob_pubkey_bytes)
543}
544
545pub async fn pow(
546    State(state): State<AppState>,
547    ValidatedForm(pubkey): ValidatedForm<RequestPubkey>,
548) -> Json<PoW> {
549    let mut valid_pubkey_for_ms = match state.valid_pubkey_for_ms.lock() {
550        Ok(guard) => guard,
551        Err(_) => return Json(PoW::default()),
552    };
553    update_senders_hashmap(&mut valid_pubkey_for_ms);
554
555    let alice_public_bytes = match hex::decode(pubkey.get_bytes()) {
556        Ok(bytes) => bytes,
557        Err(_) => return Json(PoW::default()),
558    };
559    let alice_pk: [u8; 32] = match alice_public_bytes.try_into() {
560        Ok(alice_pk) => alice_pk,
561        Err(_) => return Json(PoW::default()),
562    };
563
564    let bytes_to_search: [u8; 32] = new_nonce();
565
566    let now = SystemTime::now();
567    let duration_since_epoch = now
568        .duration_since(UNIX_EPOCH)
569        .unwrap_or(Duration::new(0, 0));
570    let unix_timestamp_seconds = duration_since_epoch.as_secs();
571    let delta = Duration::new(state.expected_time, 0);
572    let new_timestamp = unix_timestamp_seconds + delta.as_secs();
573
574    let mut bytes = [].to_vec();
575    bytes.extend(bytes_to_search);
576    bytes.extend(state.cost_pow.to_le_bytes());
577    bytes.extend(alice_pk);
578    bytes.extend(new_timestamp.to_le_bytes());
579
580    let signature = sign_message(&bytes, state.ed25519_bytes);
581
582    let pow = PoW::new(hex::encode(bytes), hex::encode(signature));
583    Json(pow)
584}
585
586pub fn update_senders_hashmap(valid_pubkeys: &mut HashMap<OnionPubkey, Keypair>) {
587    let now = SystemTime::now();
588    let duration_since_epoch = now
589        .duration_since(UNIX_EPOCH)
590        .unwrap_or(Duration::new(0, 0));
591    let current_timestamp = duration_since_epoch.as_secs();
592
593    let keys_to_remove: Vec<OnionPubkey> = valid_pubkeys
594        .iter()
595        .filter(|(_, keypair)| current_timestamp > keypair.get_unix_timestamp_seconds())
596        .map(|(sender, _)| *sender)
597        .collect();
598    for sender in keys_to_remove {
599        valid_pubkeys.remove(&sender);
600    }
601}
kkv/server.rs

kkv/
server.rs
