kaptan/firmware/libs/kaptan-file-mavlink/tests/peer_lease.rs

use kaptan_file_mavlink::{
    Content, ControlRequest, ControlResponse, DATA_BYTES, DataRequest, DispatchResult, FileFailure,
    FileMetadata, FileOperation, FilePath, FileProvider, FileRequest, FileResponse, FileService,
    FileStatus, IngressId, Principal, Target, TransferPeer, crc32c, wire,
};

const TARGET: Target = Target::new(1, 1);
const UART: TransferPeer = TransferPeer {
    principal: Principal::new(42, 7),
    ingress: IngressId(0),
};
const USB: TransferPeer = TransferPeer {
    ingress: IngressId(1),
    ..UART
};
const ID: u16 = 9;

#[derive(Default)]
struct Provider {
    bytes: Vec<u8>,
    opens: usize,
    reads: usize,
    writes: usize,
    closes: usize,
    aborts: usize,
    fail_open: bool,
    fail_read: bool,
    fail_write: bool,
    fail_close: bool,
}

fn failure(status: FileStatus) -> FileFailure {
    FileFailure::new(status, Content::try_new(b"provider detail").unwrap()).unwrap()
}

impl FileProvider for Provider {
    fn open_read(&mut self, _: &FilePath) -> Result<FileMetadata, FileFailure> {
        self.opens += 1;
        if self.fail_open {
            return Err(failure(FileStatus::NotFound));
        }
        Ok(FileMetadata {
            size: self.bytes.len() as u32,
            checksum_crc32c: crc32c(&self.bytes),
        })
    }

    fn open_write(&mut self, path: &FilePath, _: FileMetadata) -> Result<(), FileFailure> {
        self.open_read(path)?;
        self.bytes.clear();
        Ok(())
    }

    fn read(&mut self, position: u32, output: &mut [u8]) -> Result<(), FileFailure> {
        self.reads += 1;
        if self.fail_read {
            return Err(failure(FileStatus::InvalidPosition));
        }
        let start = position as usize;
        output.copy_from_slice(&self.bytes[start..start + output.len()]);
        Ok(())
    }

    fn write(&mut self, position: u32, data: &[u8]) -> Result<(), FileFailure> {
        self.writes += 1;
        if self.fail_write {
            return Err(failure(FileStatus::InvalidTransfer));
        }
        assert_eq!(position as usize, self.bytes.len());
        self.bytes.extend_from_slice(data);
        Ok(())
    }

    fn close(&mut self) -> Result<Content, FileFailure> {
        self.closes += 1;
        if self.fail_close {
            return Err(failure(FileStatus::StorageFailure));
        }
        Ok(Content::try_new(b"committed").unwrap())
    }

    fn abort(&mut self) {
        self.aborts += 1;
    }
}

fn control(operation: FileOperation, position: u32) -> FileRequest {
    FileRequest::Control(ControlRequest {
        operation,
        target: TARGET,
        position,
        file_size: 0,
        checksum_crc32c: 0,
        transfer_id: ID,
        content: Content::EMPTY,
    })
}

fn alter(request: FileRequest, change: impl FnOnce(&mut ControlRequest)) -> FileRequest {
    let FileRequest::Control(mut request) = request else {
        panic!("expected control request");
    };
    change(&mut request);
    FileRequest::Control(request)
}

fn open(operation: FileOperation, bytes: &[u8]) -> FileRequest {
    alter(control(operation, 0), |request| {
        request.content = Content::try_new(b"/test/data.bin").unwrap();
        if operation == FileOperation::OpenWrite {
            request.file_size = bytes.len() as u32;
            request.checksum_crc32c = crc32c(bytes);
        }
    })
}

fn data(position: u32, bytes: &[u8]) -> FileRequest {
    let mut data = [0; DATA_BYTES];
    data[..bytes.len()].copy_from_slice(bytes);
    FileRequest::Data(DataRequest {
        target: TARGET,
        transfer_id: ID,
        position,
        data,
    })
}

// Exercise the real message codecs while keeping ingress entirely out of band.
fn exchange(
    service: &mut FileService,
    provider: &mut Provider,
    peer: TransferPeer,
    request: FileRequest,
) -> DispatchResult {
    let decoded = wire::decode_request(&wire::encode_request(request)).unwrap();
    let result = service.dispatch(provider, peer, decoded).unwrap();
    assert_eq!(
        wire::decode_response(&wire::encode_response(result.response)),
        Ok(result.response)
    );
    let target = match result.response {
        FileResponse::Control(response) => response.target,
        FileResponse::Data(response) => response.target,
    };
    assert_eq!(
        target,
        Target::new(peer.principal.system(), peer.principal.component())
    );
    result
}

fn expect_control(result: DispatchResult, status: FileStatus, refresh: bool) -> ControlResponse {
    assert_eq!(result.refresh_lease, refresh);
    let FileResponse::Control(response) = result.response else {
        panic!("expected control response");
    };
    assert_eq!(response.status, status);
    response
}

#[test]
fn equal_wire_identities_cannot_modify_or_replay_another_ingress_upload() {
    for (owner, other_ingress) in [(UART, USB), (USB, UART)] {
        let mut service = FileService::new(TARGET);
        let mut provider = Provider::default();
        let start = open(FileOperation::OpenWrite, b"abc");
        let chunk = data(0, b"abc");
        let close = control(FileOperation::Close, 3);
        let opened = exchange(&mut service, &mut provider, owner, start);
        expect_control(opened, FileStatus::Ok, true);
        assert_eq!(exchange(&mut service, &mut provider, owner, start), opened);
        assert_eq!(provider.opens, 1);

        // Check both components of ownership, including identical transfer IDs.
        for foreign in [
            other_ingress,
            TransferPeer {
                principal: Principal::new(43, 7),
                ..owner
            },
        ] {
            for (request, status) in [
                (start, FileStatus::Busy),
                (alter(start, |r| r.transfer_id += 1), FileStatus::Busy),
                (chunk, FileStatus::InvalidTransfer),
                (control(FileOperation::Read, 0), FileStatus::InvalidTransfer),
                (
                    control(FileOperation::Abort, 0),
                    FileStatus::InvalidTransfer,
                ),
                (close, FileStatus::InvalidTransfer),
                (
                    control(FileOperation::WriteAck, 0),
                    FileStatus::InvalidState,
                ),
            ] {
                expect_control(
                    exchange(&mut service, &mut provider, foreign, request),
                    status,
                    false,
                );
            }
        }
        assert_eq!(
            (
                provider.opens,
                provider.reads,
                provider.writes,
                provider.closes,
                provider.aborts
            ),
            (1, 0, 0, 0, 0)
        );

        let ack = exchange(&mut service, &mut provider, owner, chunk);
        expect_control(ack, FileStatus::Ok, true);
        expect_control(
            exchange(&mut service, &mut provider, other_ingress, chunk),
            FileStatus::InvalidTransfer,
            false,
        );
        assert_eq!(exchange(&mut service, &mut provider, owner, chunk), ack);
        assert_eq!(provider.writes, 1);

        // Lose the durable close reply, then try to steal, evict, or corrupt it.
        let lost_reply = exchange(&mut service, &mut provider, owner, close);
        expect_control(lost_reply, FileStatus::Ok, true);
        for (request, status) in [
            (close, FileStatus::InvalidTransfer),
            (chunk, FileStatus::InvalidTransfer),
            (start, FileStatus::Busy),
            (alter(start, |r| r.transfer_id += 1), FileStatus::Busy),
            (
                control(FileOperation::Abort, 3),
                FileStatus::InvalidTransfer,
            ),
            (
                control(FileOperation::WriteAck, 3),
                FileStatus::InvalidState,
            ),
        ] {
            expect_control(
                exchange(&mut service, &mut provider, other_ingress, request),
                status,
                false,
            );
        }
        expect_control(
            exchange(
                &mut service,
                &mut provider,
                owner,
                control(FileOperation::Close, 2),
            ),
            FileStatus::RequestConflict,
            false,
        );
        expect_control(
            exchange(&mut service, &mut provider, owner, start),
            FileStatus::RequestConflict,
            false,
        );
        for _ in 0..3 {
            assert_eq!(
                exchange(&mut service, &mut provider, owner, close),
                lost_reply
            );
        }
        assert_eq!(provider.bytes, b"abc");
        assert_eq!(
            (
                provider.opens,
                provider.writes,
                provider.closes,
                provider.aborts
            ),
            (1, 1, 1, 0)
        );

        assert!(!service.expire(&mut provider));
        expect_control(
            exchange(&mut service, &mut provider, owner, close),
            FileStatus::InvalidTransfer,
            false,
        );
        expect_control(
            exchange(&mut service, &mut provider, other_ingress, start),
            FileStatus::Ok,
            true,
        );
        expect_control(
            exchange(&mut service, &mut provider, owner, close),
            FileStatus::InvalidTransfer,
            false,
        );
        assert_eq!(provider.closes, 1);
    }
}

#[test]
fn read_and_abort_replay_are_bound_to_the_full_peer() {
    let mut service = FileService::new(TARGET);
    let mut provider = Provider {
        bytes: b"snapshot".to_vec(),
        ..Provider::default()
    };
    let start = open(FileOperation::OpenRead, &[]);
    expect_control(
        exchange(&mut service, &mut provider, UART, start),
        FileStatus::Ok,
        true,
    );
    let read = control(FileOperation::Read, 0);
    expect_control(
        exchange(&mut service, &mut provider, USB, read),
        FileStatus::InvalidTransfer,
        false,
    );
    let first = exchange(&mut service, &mut provider, UART, read);
    assert!(first.refresh_lease);
    let FileResponse::Data(response) = first.response else {
        panic!("expected data")
    };
    assert_eq!(&response.data[..8], b"snapshot");
    expect_control(
        exchange(&mut service, &mut provider, USB, read),
        FileStatus::InvalidTransfer,
        false,
    );
    expect_control(
        exchange(&mut service, &mut provider, UART, data(0, b"changed")),
        FileStatus::InvalidState,
        false,
    );
    expect_control(
        exchange(
            &mut service,
            &mut provider,
            UART,
            control(FileOperation::Read, 8),
        ),
        FileStatus::InvalidPosition,
        false,
    );
    assert_eq!(exchange(&mut service, &mut provider, UART, read), first);
    assert_eq!(provider.reads, 1);

    let abort = control(FileOperation::Abort, 8);
    expect_control(
        exchange(&mut service, &mut provider, USB, abort),
        FileStatus::InvalidTransfer,
        false,
    );
    let aborted = exchange(&mut service, &mut provider, UART, abort);
    expect_control(aborted, FileStatus::Ok, true);
    expect_control(
        exchange(&mut service, &mut provider, USB, abort),
        FileStatus::InvalidTransfer,
        false,
    );
    assert_eq!(exchange(&mut service, &mut provider, UART, abort), aborted);
    assert!(!service.expire(&mut provider));
    assert!(!service.expire(&mut provider));
    expect_control(
        exchange(&mut service, &mut provider, UART, abort),
        FileStatus::InvalidTransfer,
        false,
    );
    expect_control(
        exchange(&mut service, &mut provider, USB, start),
        FileStatus::Ok,
        true,
    );
    assert_eq!(provider.aborts, 1);
}

#[test]
fn rejected_activity_cannot_retain_open_or_completed_lease() {
    const LEASE: u64 = 10;
    for completed in [false, true] {
        let mut service = FileService::new(TARGET);
        let mut provider = Provider::default();
        let start = open(FileOperation::OpenWrite, &[]);
        expect_control(
            exchange(&mut service, &mut provider, UART, start),
            FileStatus::Ok,
            true,
        );
        if completed {
            expect_control(
                exchange(
                    &mut service,
                    &mut provider,
                    UART,
                    control(FileOperation::Close, 0),
                ),
                FileStatus::Ok,
                true,
            );
        }
        let mut lease_until = LEASE;
        for now in 1..LEASE {
            for peer in [UART, USB] {
                for request in [
                    control(FileOperation::WriteAck, 0),
                    alter(start, |r| {
                        r.content = Content::try_new(b"bad path").unwrap()
                    }),
                    alter(start, |r| r.transfer_id = 0),
                    alter(control(FileOperation::Read, 0), |r| r.file_size = 1),
                    alter(control(FileOperation::Abort, 0), |r| r.checksum_crc32c = 1),
                    alter(control(FileOperation::Close, 0), |r| {
                        r.content = Content::try_new(b"bad").unwrap()
                    }),
                    data(1, b"x"),
                ] {
                    let result = exchange(&mut service, &mut provider, peer, request);
                    if result.refresh_lease {
                        lease_until = now + LEASE;
                    }
                    assert_eq!(
                        lease_until, LEASE,
                        "rejection retained the lease: {result:?}"
                    );
                }
            }
            for request in [
                start,
                control(FileOperation::Close, 0),
                control(FileOperation::Abort, 0),
            ] {
                assert!(!exchange(&mut service, &mut provider, USB, request).refresh_lease);
            }
        }
        assert_eq!(provider.opens, 1);
        assert_eq!(provider.closes, usize::from(completed));
        assert_eq!(provider.aborts, 0);

        // A clock-owning caller expires the entire slot at the unchanged deadline.
        assert_eq!(service.expire(&mut provider), !completed);
        assert!(!service.expire(&mut provider));
        assert_eq!(provider.aborts, usize::from(!completed));
        expect_control(
            exchange(
                &mut service,
                &mut provider,
                UART,
                control(FileOperation::Close, 0),
            ),
            FileStatus::InvalidTransfer,
            false,
        );
        expect_control(
            exchange(&mut service, &mut provider, USB, start),
            FileStatus::Ok,
            true,
        );
    }
}

#[test]
fn invalid_owner_requests_do_not_mutate_a_live_upload_or_its_retry() {
    let mut service = FileService::new(TARGET);
    let mut provider = Provider::default();
    let start = open(FileOperation::OpenWrite, b"abc");
    expect_control(
        exchange(&mut service, &mut provider, UART, start),
        FileStatus::Ok,
        true,
    );
    for (request, status) in [
        (control(FileOperation::Read, 0), FileStatus::InvalidState),
        (
            control(FileOperation::Close, 0),
            FileStatus::InvalidPosition,
        ),
        (data(1, b"bc"), FileStatus::InvalidPosition),
        (data(0, b"abc!"), FileStatus::RequestConflict),
        (
            alter(start, |r| r.checksum_crc32c ^= 1),
            FileStatus::RequestConflict,
        ),
        (
            alter(control(FileOperation::Close, 0), |r| r.transfer_id += 1),
            FileStatus::InvalidTransfer,
        ),
    ] {
        expect_control(
            exchange(&mut service, &mut provider, UART, request),
            status,
            false,
        );
    }
    assert_eq!(
        (provider.writes, provider.closes, provider.aborts),
        (0, 0, 0)
    );
    let accepted = exchange(&mut service, &mut provider, UART, data(0, b"abc"));
    expect_control(accepted, FileStatus::Ok, true);
    expect_control(
        exchange(&mut service, &mut provider, UART, data(0, b"abd")),
        FileStatus::RequestConflict,
        false,
    );
    expect_control(
        exchange(&mut service, &mut provider, UART, start),
        FileStatus::RequestConflict,
        false,
    );
    assert_eq!(
        exchange(&mut service, &mut provider, UART, data(0, b"abc")),
        accepted
    );
    expect_control(
        exchange(
            &mut service,
            &mut provider,
            UART,
            control(FileOperation::Close, 3),
        ),
        FileStatus::Ok,
        true,
    );
    assert_eq!(provider.bytes, b"abc");
    assert_eq!(
        (provider.writes, provider.closes, provider.aborts),
        (1, 1, 0)
    );
}

#[test]
fn terminal_failures_refresh_and_replay_only_for_the_accepted_peer() {
    for stage in 0..4 {
        let mut service = FileService::new(TARGET);
        let mut provider = Provider {
            bytes: b"abc".to_vec(),
            fail_read: stage == 0,
            fail_write: stage == 1,
            fail_close: stage == 2,
            ..Provider::default()
        };
        let operation = if stage == 0 {
            FileOperation::OpenRead
        } else {
            FileOperation::OpenWrite
        };
        let start = open(operation, b"abc");
        expect_control(
            exchange(&mut service, &mut provider, UART, start),
            FileStatus::Ok,
            true,
        );
        if stage >= 2 {
            let bytes = if stage == 3 { b"abd" } else { b"abc" };
            expect_control(
                exchange(&mut service, &mut provider, UART, data(0, bytes)),
                FileStatus::Ok,
                true,
            );
        }
        let (request, status) = match stage {
            0 => (control(FileOperation::Read, 0), FileStatus::InvalidPosition),
            1 => (data(0, b"abc"), FileStatus::InvalidTransfer),
            2 => (control(FileOperation::Close, 3), FileStatus::StorageFailure),
            _ => (
                control(FileOperation::Close, 3),
                FileStatus::ChecksumMismatch,
            ),
        };
        let lost_reply = exchange(&mut service, &mut provider, UART, request);
        expect_control(lost_reply, status, true);
        expect_control(
            exchange(&mut service, &mut provider, USB, request),
            FileStatus::InvalidTransfer,
            false,
        );
        let conflict = match request {
            FileRequest::Control(_) => alter(request, |r| r.position += 1),
            FileRequest::Data(_) => data(0, b"abd"),
        };
        expect_control(
            exchange(&mut service, &mut provider, UART, conflict),
            FileStatus::RequestConflict,
            false,
        );
        assert_eq!(
            exchange(&mut service, &mut provider, UART, request),
            lost_reply
        );
        assert_eq!(provider.reads, usize::from(stage == 0));
        assert_eq!(provider.writes, usize::from(stage != 0));
        assert_eq!(provider.closes, usize::from(stage == 2));
        assert_eq!(provider.aborts, 1);
    }
}

#[test]
fn failed_open_does_not_claim_or_renew_the_slot_or_evict_a_completed_reply() {
    let mut service = FileService::new(TARGET);
    let mut provider = Provider {
        fail_open: true,
        ..Provider::default()
    };
    let start = open(FileOperation::OpenWrite, &[]);
    expect_control(
        exchange(&mut service, &mut provider, UART, start),
        FileStatus::NotFound,
        false,
    );
    provider.fail_open = false;
    expect_control(
        exchange(&mut service, &mut provider, USB, start),
        FileStatus::Ok,
        true,
    );
    let close = control(FileOperation::Close, 0);
    let completed = exchange(&mut service, &mut provider, USB, close);
    expect_control(completed, FileStatus::Ok, true);
    provider.fail_open = true;
    expect_control(
        exchange(
            &mut service,
            &mut provider,
            USB,
            alter(start, |r| r.transfer_id += 1),
        ),
        FileStatus::NotFound,
        false,
    );
    assert_eq!(exchange(&mut service, &mut provider, USB, close), completed);
    assert_eq!(provider.closes, 1);
}