Skip to main content

made_client/
artifact_export.rs

1use std::path::{Path, PathBuf};
2
3use made_proto::v1::ReadArtifactChunkRequest;
4use sha2::{Digest, Sha256};
5use tokio::io::AsyncWriteExt;
6use uuid::Uuid;
7
8use crate::{MadeClient, MadeClientError};
9
10const READ_LIMIT: u32 = 1024 * 1024;
11
12impl MadeClient {
13    pub async fn export_artifact(
14        &self,
15        artifact_id: &str,
16        destination: &Path,
17        overwrite: bool,
18    ) -> Result<(), MadeClientError> {
19        let parent = destination.parent().unwrap_or_else(|| Path::new("."));
20        tokio::fs::create_dir_all(parent)
21            .await
22            .map_err(|error| MadeClientError::io(parent, error))?;
23        let record = self.get_artifact(artifact_id).await?;
24        let reference = record.artifact.ok_or_else(|| {
25            MadeClientError::ProtocolViolation("artifact record has no reference".to_owned())
26        })?;
27        if reference.artifact_id != artifact_id {
28            return Err(MadeClientError::ProtocolViolation(format!(
29                "requested artifact {artifact_id}, got {}",
30                reference.artifact_id
31            )));
32        }
33        let temporary = temporary_path(destination);
34        let result = self
35            .write_verified_artifact(
36                &reference.digest,
37                reference.size_bytes,
38                &temporary,
39                artifact_id,
40            )
41            .await
42            .and_then(|()| install_file(&temporary, destination, overwrite));
43        if result.is_err() {
44            let _ = tokio::fs::remove_file(&temporary).await;
45        }
46        result
47    }
48
49    async fn write_verified_artifact(
50        &self,
51        expected_digest: &str,
52        expected_size: u64,
53        temporary: &Path,
54        artifact_id: &str,
55    ) -> Result<(), MadeClientError> {
56        let mut file = tokio::fs::OpenOptions::new()
57            .create_new(true)
58            .write(true)
59            .open(temporary)
60            .await
61            .map_err(|error| MadeClientError::io(temporary, error))?;
62        let mut aggregate = Sha256::new();
63        let mut offset = 0_u64;
64        loop {
65            let response = self
66                .rpc()
67                .read_artifact_chunk(Self::request(
68                    &self.context(),
69                    "/underpass.made.v1.MadeService/ReadArtifactChunk",
70                    ReadArtifactChunkRequest {
71                        artifact_id: artifact_id.to_owned(),
72                        offset,
73                        max_bytes: READ_LIMIT,
74                    },
75                ))
76                .await
77                .map_err(MadeClientError::from_status)?
78                .into_inner();
79            if response.bytes.is_empty() && !response.eof {
80                return Err(MadeClientError::ProtocolViolation(
81                    "artifact read made no progress before EOF".to_owned(),
82                ));
83            }
84            let observed_chunk_digest = sha256_digest(&response.bytes);
85            if observed_chunk_digest != response.chunk_digest {
86                return Err(MadeClientError::ArtifactIntegrityMismatch {
87                    expected: response.chunk_digest,
88                    observed: observed_chunk_digest,
89                });
90            }
91            let next_offset = offset
92                .checked_add(response.bytes.len() as u64)
93                .ok_or_else(|| {
94                    MadeClientError::ProtocolViolation("artifact offset overflow".to_owned())
95                })?;
96            if response.next_offset != next_offset {
97                return Err(MadeClientError::ProtocolViolation(format!(
98                    "artifact next offset {}, expected {next_offset}",
99                    response.next_offset
100                )));
101            }
102            file.write_all(&response.bytes)
103                .await
104                .map_err(|error| MadeClientError::io(temporary, error))?;
105            aggregate.update(&response.bytes);
106            offset = next_offset;
107            if response.eof {
108                break;
109            }
110        }
111        if offset != expected_size {
112            return Err(MadeClientError::ArtifactSizeMismatch {
113                expected: expected_size,
114                observed: offset,
115            });
116        }
117        let observed = format!("sha256:{:x}", aggregate.finalize());
118        if observed != expected_digest {
119            return Err(MadeClientError::ArtifactIntegrityMismatch {
120                expected: expected_digest.to_owned(),
121                observed,
122            });
123        }
124        file.sync_all()
125            .await
126            .map_err(|error| MadeClientError::io(temporary, error))
127    }
128}
129
130fn sha256_digest(bytes: &[u8]) -> String {
131    format!("sha256:{:x}", Sha256::digest(bytes))
132}
133
134fn temporary_path(destination: &Path) -> PathBuf {
135    let parent = destination.parent().unwrap_or_else(|| Path::new("."));
136    let name = destination
137        .file_name()
138        .and_then(|name| name.to_str())
139        .unwrap_or("artifact");
140    parent.join(format!(".{name}.{}.part", Uuid::new_v4()))
141}
142
143fn install_file(
144    temporary: &Path,
145    destination: &Path,
146    overwrite: bool,
147) -> Result<(), MadeClientError> {
148    if overwrite {
149        std::fs::rename(temporary, destination)
150            .map_err(|error| MadeClientError::io(destination, error))
151    } else {
152        std::fs::hard_link(temporary, destination)
153            .map_err(|error| MadeClientError::io(destination, error))?;
154        std::fs::remove_file(temporary).map_err(|error| MadeClientError::io(temporary, error))
155    }
156}