dosh/tests/protocol_auth.rs

use dosh::auth::{
    build_bootstrap, decode_bootstrap, encode_bootstrap, load_or_create_server_secret,
    open_attach_ticket, verify_attach_ticket, verify_bootstrap,
};
use dosh::config::ServerConfig;
use dosh::crypto;
use dosh::protocol::{self, CLIENT_TO_SERVER, PacketKind, ReplayWindow};

#[test]
fn bootstrap_round_trips_and_verifies() {
    let dir = tempfile::tempdir().unwrap();
    let config = ServerConfig {
        secret_path: dir.path().join("secret").display().to_string(),
        ..ServerConfig::default()
    };
    let secret = load_or_create_server_secret(&config).unwrap();
    let nonce = crypto::random_12();
    let resp = build_bootstrap(
        &config,
        &secret,
        "user".to_string(),
        "default".to_string(),
        "read-write".to_string(),
        (80, 24),
        nonce,
        "127.0.0.1".to_string(),
    )
    .unwrap();
    assert!(verify_bootstrap(&resp, &secret).unwrap());
    let encoded = encode_bootstrap(&resp).unwrap();
    let decoded = decode_bootstrap(&encoded).unwrap();
    assert_eq!(decoded.session, "default");
    assert_eq!(decoded.session_key, resp.session_key);
}

#[test]
fn encrypted_packet_round_trips() {
    let key = crypto::random_32();
    let conn_id = crypto::random_16();
    let packet = protocol::encode_encrypted(
        PacketKind::Input,
        conn_id,
        7,
        0,
        &key,
        CLIENT_TO_SERVER,
        b"hello",
    )
    .unwrap();
    let decoded = protocol::decode(&packet).unwrap();
    assert_eq!(decoded.header.kind, PacketKind::Input);
    assert_eq!(decoded.header.conn_id, conn_id);
    assert_eq!(
        decoded.header.session_key_id,
        protocol::session_key_id(&key)
    );
    let plain = protocol::decrypt_body(&decoded, &key, CLIENT_TO_SERVER).unwrap();
    assert_eq!(plain, b"hello");
}

#[test]
fn packet_decode_rejects_truncated_or_trailing_datagrams() {
    let packet =
        protocol::encode_plain(PacketKind::Ping, crypto::random_16(), 1, 0, b"hello").unwrap();

    let truncated = &packet[..packet.len() - 1];
    let err = protocol::decode(truncated).unwrap_err();
    assert!(err.to_string().contains("truncated packet body"));

    let mut trailing = packet.clone();
    trailing.extend_from_slice(b"junk");
    let err = protocol::decode(&trailing).unwrap_err();
    assert!(err.to_string().contains("trailing packet bytes"));
}

#[test]
fn plaintext_packet_never_decrypts_as_authenticated_body() {
    let key = crypto::random_32();
    let mut decoded = protocol::decode(
        &protocol::encode_plain(PacketKind::Input, crypto::random_16(), 1, 0, b"hello").unwrap(),
    )
    .unwrap();
    decoded.header.session_key_id = protocol::session_key_id(&key);

    let err = protocol::decrypt_body(&decoded, &key, CLIENT_TO_SERVER).unwrap_err();
    assert!(err.to_string().contains("not encrypted"));
}

#[test]
fn encrypted_packet_rejects_wrong_session_key_id_before_decrypt() {
    let key = crypto::random_32();
    let wrong_key = crypto::random_32();
    let packet = protocol::encode_encrypted(
        PacketKind::Input,
        crypto::random_16(),
        1,
        0,
        &key,
        CLIENT_TO_SERVER,
        b"hello",
    )
    .unwrap();
    let decoded = protocol::decode(&packet).unwrap();
    let err = protocol::decrypt_body(&decoded, &wrong_key, CLIENT_TO_SERVER).unwrap_err();
    assert!(err.to_string().contains("session key id"));
}

#[test]
fn attach_ticket_is_sealed_and_verifies_scope() {
    let dir = tempfile::tempdir().unwrap();
    let config = ServerConfig {
        secret_path: dir.path().join("secret").display().to_string(),
        ..ServerConfig::default()
    };
    let secret = load_or_create_server_secret(&config).unwrap();
    let resp = build_bootstrap(
        &config,
        &secret,
        "user".to_string(),
        "work".to_string(),
        "view-only".to_string(),
        (100, 30),
        crypto::random_12(),
        "127.0.0.1".to_string(),
    )
    .unwrap();

    let opened = open_attach_ticket(&secret, &resp.attach_ticket).unwrap();
    assert_eq!(opened.session, "work");
    assert_eq!(opened.mode, "view-only");
    assert_eq!(opened.psk, resp.attach_ticket_psk);

    assert!(
        verify_attach_ticket(
            &secret,
            &resp.attach_ticket,
            &resp.attach_ticket_psk,
            "work",
            "view-only",
        )
        .unwrap()
        .is_some()
    );
    assert!(
        verify_attach_ticket(
            &secret,
            &resp.attach_ticket,
            &resp.attach_ticket_psk,
            "default",
            "view-only",
        )
        .unwrap()
        .is_none()
    );
}

#[test]
fn peek_foreign_wire_version_flags_only_version_skew() {
    let key = crypto::random_32();
    let mut packet = protocol::encode_encrypted(
        PacketKind::Input,
        crypto::random_16(),
        1,
        0,
        &key,
        CLIENT_TO_SERVER,
        b"hi",
    )
    .unwrap();
    // A correctly framed packet for this build is not "foreign".
    assert_eq!(protocol::peek_foreign_wire_version(&packet), None);
    // Bumping the wire version byte makes it undecodable but recognizable.
    packet[4] = protocol::VERSION.wrapping_add(7);
    assert_eq!(
        protocol::peek_foreign_wire_version(&packet),
        Some(protocol::VERSION.wrapping_add(7))
    );
    assert!(protocol::decode(&packet).is_err());
    // Non-Dosh datagrams and runts are ignored.
    assert_eq!(protocol::peek_foreign_wire_version(b"XXXX\x01"), None);
    assert_eq!(protocol::peek_foreign_wire_version(b"DOS"), None);
}

#[test]
fn rekey_key_derivation_agrees_and_is_independent_per_epoch() {
    use dosh::native::derive_rekey_session_key;

    // The handshake/current key both peers already share.
    let current_key = crypto::random_32();
    let current_id = protocol::session_key_id(&current_key);
    // Fresh server-generated material, delivered confidentially in the Rekey.
    let material = crypto::random_32();

    // Both peers derive identically from shared current key + shipped material.
    let server_view = derive_rekey_session_key(&current_key, &material, &current_id, 1).unwrap();
    let client_view = derive_rekey_session_key(&current_key, &material, &current_id, 1).unwrap();
    assert_eq!(server_view, client_view);

    // The rotated key must not equal the handshake/current key.
    assert_ne!(server_view, current_key);

    // A different epoch (or different material) yields a different key.
    let next_epoch = derive_rekey_session_key(&current_key, &material, &current_id, 2).unwrap();
    assert_ne!(server_view, next_epoch);
    let other_material = crypto::random_32();
    let other = derive_rekey_session_key(&current_key, &other_material, &current_id, 1).unwrap();
    assert_ne!(server_view, other);
}

#[test]
fn rekey_round_trip_decrypts_old_and_new_epoch_packets() {
    use dosh::native::derive_rekey_session_key;

    let key_epoch0 = crypto::random_32();
    let id0 = protocol::session_key_id(&key_epoch0);
    let conn_id = crypto::random_16();

    // Pre-rekey packet sealed under epoch-0 key.
    let pre = protocol::encode_encrypted(
        PacketKind::Frame,
        conn_id,
        5,
        0,
        &key_epoch0,
        protocol::SERVER_TO_CLIENT,
        b"before",
    )
    .unwrap();

    // Rotate to epoch 1.
    let material = crypto::random_32();
    let key_epoch1 = derive_rekey_session_key(&key_epoch0, &material, &id0, 1).unwrap();
    let post = protocol::encode_encrypted(
        PacketKind::Frame,
        conn_id,
        6,
        0,
        &key_epoch1,
        protocol::SERVER_TO_CLIENT,
        b"after",
    )
    .unwrap();

    let pre = protocol::decode(&pre).unwrap();
    let post = protocol::decode(&post).unwrap();

    // Each epoch's key carries its own session_key_id; the receiver picks the
    // right one and both decrypt correctly.
    assert_eq!(pre.header.session_key_id, id0);
    assert_eq!(
        pre.header.session_key_id,
        protocol::session_key_id(&key_epoch0)
    );
    assert_eq!(
        post.header.session_key_id,
        protocol::session_key_id(&key_epoch1)
    );
    assert_eq!(
        protocol::decrypt_body(&pre, &key_epoch0, protocol::SERVER_TO_CLIENT).unwrap(),
        b"before"
    );
    assert_eq!(
        protocol::decrypt_body(&post, &key_epoch1, protocol::SERVER_TO_CLIENT).unwrap(),
        b"after"
    );

    // A stale-epoch packet under the wrong key is rejected via session_key_id
    // BEFORE any AEAD work — ignorable, not a fatal decrypt error.
    let err = protocol::decrypt_body(&post, &key_epoch0, protocol::SERVER_TO_CLIENT).unwrap_err();
    assert!(err.to_string().contains("session key id"));
}

#[test]
fn check_native_protocol_version_names_the_mismatch() {
    use dosh::native::{NATIVE_PROTOCOL_VERSION, check_native_protocol_version};
    check_native_protocol_version(NATIVE_PROTOCOL_VERSION, "server").unwrap();
    let err = check_native_protocol_version(NATIVE_PROTOCOL_VERSION.wrapping_add(1), "server")
        .unwrap_err();
    let message = err.to_string();
    assert!(
        message.contains(protocol::VERSION_MISMATCH_REASON),
        "expected actionable upgrade message, got {message:?}"
    );
    assert!(
        message.contains("server"),
        "should name the wrong peer: {message:?}"
    );
}

#[test]
fn replay_window_rejects_duplicates_but_allows_bounded_out_of_order() {
    let mut replay = ReplayWindow::new(8);
    assert!(replay.accept(10));
    assert!(!replay.accept(10));
    assert!(replay.accept(12));
    assert!(replay.accept(11));
    assert!(!replay.accept(11));
    assert!(replay.accept(18));
    assert!(!replay.accept(9));
    assert!(!replay.accept(0));
}