apeirogond 0.36.0

Daemon for Apeirogon, a cross-platform software KVM
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

// SPDX-License-Identifier: GPL-3.0-or-later

use std::io::{Seek, SeekFrom, Write};
use std::sync::atomic::Ordering;

use a7n_core::files::{EntryKind, ValidManifest};
use a7n_proto::{
    limits,
    v1::{FileChunk, FileEnd, transfer_frame::Frame},
};
use sha2::Digest;

use super::super::{
    Failure, Progress, Result,
    staging::{PartialFile, Staging},
};

pub(in crate::daemon::files) fn apply(
    state: &mut Staging,
    manifest: &ValidManifest,
    frame: Frame,
    progress: &Progress,
) -> Result<bool> {
    match frame {
        Frame::Chunk(chunk) => {
            apply_chunk(state, manifest, &chunk, progress)?;
        }
        Frame::FileEnd(end) => {
            apply_file_end(state, manifest, &end, progress)?;
        }
        Frame::Done(_) if state.files.iter().flatten().all(|file| file.ended) => return Ok(true),
        Frame::Abort(_) => return Err(Failure::Cancelled),
        _ => return Err(Failure::Invalid),
    }
    Ok(false)
}

fn apply_chunk(
    state: &mut Staging,
    manifest: &ValidManifest,
    chunk: &FileChunk,
    progress: &Progress,
) -> Result<()> {
    let index = usize::try_from(chunk.file_index).map_err(|_| Failure::Invalid)?;
    let file = state
        .files
        .get_mut(index)
        .and_then(Option::as_mut)
        .ok_or(Failure::Invalid)?;
    let EntryKind::File { size, .. } = manifest.entries()[index].kind else {
        return Err(Failure::Invalid);
    };
    if file.ended
        || chunk.offset != file.written
        || chunk.data.len() > limits::CHUNK_MAX
        || (chunk.data.len() as u64) > size.saturating_sub(file.written)
    {
        return Err(Failure::Invalid);
    }
    open_partial(&mut state.open, file, chunk.file_index)?;
    state
        .open
        .as_mut()
        .ok_or(Failure::Io)?
        .1
        .write_all(&chunk.data)
        .map_err(|_| Failure::Io)?;
    file.hash.update(&chunk.data);
    file.written += chunk.data.len() as u64;
    progress
        .bytes
        .fetch_add(chunk.data.len() as u64, Ordering::Relaxed);
    Ok(())
}

fn open_partial(
    open: &mut Option<(u32, std::fs::File)>,
    file: &PartialFile,
    file_index: u32,
) -> Result<()> {
    if open.as_ref().is_none_or(|(index, _)| *index != file_index) {
        let mut handle = a7n_platform::transfer_fs::open_file(&file.path, false, true)
            .map_err(|_| Failure::Io)?;
        if handle.metadata().map_err(|_| Failure::Io)?.len() != file.written {
            return Err(Failure::Io);
        }
        handle
            .seek(SeekFrom::Start(file.written))
            .map_err(|_| Failure::Io)?;
        *open = Some((file_index, handle));
    }
    Ok(())
}

fn apply_file_end(
    state: &mut Staging,
    manifest: &ValidManifest,
    end: &FileEnd,
    progress: &Progress,
) -> Result<()> {
    let index = usize::try_from(end.file_index).map_err(|_| Failure::Invalid)?;
    let file = state
        .files
        .get_mut(index)
        .and_then(Option::as_mut)
        .ok_or(Failure::Invalid)?;
    let EntryKind::File { size, sha256 } = manifest.entries()[index].kind else {
        return Err(Failure::Invalid);
    };
    if file.ended {
        return Err(Failure::Invalid);
    }
    if file.written != size
        || end.sha256 != sha256
        || file.hash.clone().finalize().as_slice() != sha256
    {
        return Err(Failure::Integrity);
    }
    if let Some((index, handle)) = state.open.take() {
        handle.sync_all().map_err(|_| Failure::Io)?;
        if index != end.file_index {
            state.open = Some((index, handle));
        }
    }
    file.ended = true;
    progress.files.fetch_add(1, Ordering::Relaxed);
    Ok(())
}