forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
9.5 KiB
1//! HTTP client for the forge runner API (contract section 4).
2//!
3//! Runner calls (register, claim) authenticate with the runner secret. Job calls (logs,
4//! heartbeat, artifacts, finish) authenticate with the per-job token.
5
6use std::path::Path;
7use std::time::Duration;
8
9use forge_ci::proto::{
10 ClaimRequest, FinishRequest, HeartbeatResponse, JobAssignment, RegisterRequest,
11 RegisterResponse,
12};
13use reqwest::StatusCode;
14use reqwest::header::{CONTENT_LENGTH, HeaderMap};
15use serde::de::DeserializeOwned;
16use sha2::{Digest, Sha256};
17use tokio::io::AsyncWriteExt;
18
19/// The claim long-poll is 25s on the forge side, so 35s leaves room for the response.
20pub const CLAIM_TIMEOUT: Duration = Duration::from_secs(35);
21const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
22const LOG_TIMEOUT: Duration = Duration::from_secs(60);
23const ARTIFACT_TIMEOUT: Duration = Duration::from_secs(60 * 60);
24const BINARY_TIMEOUT: Duration = Duration::from_secs(10 * 60);
25const HEADER_LOG_OFFSET: &str = "X-Forge-Log-Offset";
26const HEADER_LOG_SIZE: &str = "X-Forge-Log-Size";
27/// Sent on every claim: `<ARCH>-<OS>` of this runner, for example `x86_64-linux`.
28const HEADER_RUNNER_TARGET: &str = "Forge-Runner-Target";
29/// On a claim response: the sha256 (hex) of the forge's own binary, if the forge advertises one.
30const HEADER_BINARY_SHA256: &str = "Forge-Binary-Sha256";
31/// `GET`: the forge's binary, for a runner that is updating itself.
32const BINARY_PATH: &str = "/api/v1/ci/runners/binary";
33
34/// Errors from talking to the forge.
35#[derive(Debug, thiserror::Error)]
36pub enum ClientError {
37 #[error("http: {0}")]
38 Http(#[from] reqwest::Error),
39 #[error("forge answered {status}: {body}")]
40 Status { status: StatusCode, body: String },
41 #[error("local io: {0}")]
42 Io(#[from] std::io::Error),
43 #[error("protocol: {0}")]
44 Protocol(String),
45}
46
47/// The forge's answer to a log chunk.
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49pub enum LogAck {
50 /// The chunk was stored (or was already stored, which is an idempotent retry).
51 Accepted,
52 /// The forge holds a different number of bytes. The runner must continue from this size.
53 Resync(u64),
54}
55
56/// The forge's answer to a claim.
57#[derive(Clone, Debug)]
58pub struct ClaimReply {
59 /// The job assigned to this runner, or `None` when there is nothing to run (204).
60 pub job: Option<JobAssignment>,
61 /// The sha256 (lowercase hex) of the forge's binary, when the response advertises one.
62 pub binary_sha256: Option<String>,
63}
64
65/// A client for one forge base URL.
66#[derive(Clone, Debug)]
67pub struct ForgeClient {
68 http: reqwest::Client,
69 base: String,
70}
71
72impl ForgeClient {
73 /// Builds a client for `base_url`, for example `https://forge.example`.
74 pub fn new(base_url: &str) -> Result<Self, ClientError> {
75 let http = reqwest::Client::builder()
76 .user_agent(concat!("forge-runner/", env!("CARGO_PKG_VERSION")))
77 .connect_timeout(Duration::from_secs(10))
78 .build()?;
79 Ok(Self {
80 http,
81 base: base_url.trim_end_matches('/').to_owned(),
82 })
83 }
84
85 fn url(&self, path: &str) -> String {
86 format!("{}{}", self.base, path)
87 }
88
89 /// `POST /api/v1/ci/runners/register`.
90 pub async fn register(&self, req: &RegisterRequest) -> Result<RegisterResponse, ClientError> {
91 let resp = self
92 .http
93 .post(self.url("/api/v1/ci/runners/register"))
94 .json(req)
95 .timeout(REQUEST_TIMEOUT)
96 .send()
97 .await?;
98 json_or_error(resp).await
99 }
100
101 /// `POST /api/v1/ci/runners/claim`. A 204 means nothing to run, and the reply has no job.
102 pub async fn claim(
103 &self,
104 runner_secret: &str,
105 labels: &[String],
106 ) -> Result<ClaimReply, ClientError> {
107 let req = ClaimRequest {
108 labels: labels.to_vec(),
109 };
110 let target = format!("{}-{}", std::env::consts::ARCH, std::env::consts::OS);
111 let resp = self
112 .http
113 .post(self.url("/api/v1/ci/runners/claim"))
114 .bearer_auth(runner_secret)
115 .header(HEADER_RUNNER_TARGET, target)
116 .json(&req)
117 .timeout(CLAIM_TIMEOUT)
118 .send()
119 .await?;
120 let binary_sha256 = advertised_sha256(resp.headers());
121 match resp.status() {
122 StatusCode::NO_CONTENT => Ok(ClaimReply {
123 job: None,
124 binary_sha256,
125 }),
126 status if status.is_success() => Ok(ClaimReply {
127 job: Some(resp.json().await?),
128 binary_sha256,
129 }),
130 _ => Err(status_error(resp).await),
131 }
132 }
133
134 /// `GET /api/v1/ci/runners/binary`. Streams the forge's binary into `dest` and returns the
135 /// sha256 (lowercase hex) of the bytes written. The caller checks it against the advertised
136 /// value and syncs the file.
137 pub async fn download_binary(
138 &self,
139 runner_secret: &str,
140 dest: &Path,
141 ) -> Result<String, ClientError> {
142 let mut resp = self
143 .http
144 .get(self.url(BINARY_PATH))
145 .bearer_auth(runner_secret)
146 .timeout(BINARY_TIMEOUT)
147 .send()
148 .await?;
149 if !resp.status().is_success() {
150 return Err(status_error(resp).await);
151 }
152 let mut file = tokio::fs::File::create(dest).await?;
153 let mut hasher = Sha256::new();
154 while let Some(chunk) = resp.chunk().await? {
155 hasher.update(&chunk);
156 file.write_all(&chunk).await?;
157 }
158 file.flush().await?;
159 Ok(hex::encode(hasher.finalize()))
160 }
161
162 /// `POST /api/v1/ci/jobs/{id}/logs` with the byte offset of `body` in the job log.
163 pub async fn post_logs(
164 &self,
165 job_id: &str,
166 job_token: &str,
167 offset: u64,
168 body: Vec<u8>,
169 ) -> Result<LogAck, ClientError> {
170 let resp = self
171 .http
172 .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/logs")))
173 .bearer_auth(job_token)
174 .header(HEADER_LOG_OFFSET, offset.to_string())
175 .header(reqwest::header::CONTENT_TYPE, "text/plain; charset=utf-8")
176 .body(body)
177 .timeout(LOG_TIMEOUT)
178 .send()
179 .await?;
180 match resp.status() {
181 StatusCode::CONFLICT => {
182 let size = resp
183 .headers()
184 .get(HEADER_LOG_SIZE)
185 .and_then(|v| v.to_str().ok())
186 .and_then(|v| v.parse::<u64>().ok())
187 .ok_or_else(|| {
188 ClientError::Protocol("409 without a valid X-Forge-Log-Size".into())
189 })?;
190 Ok(LogAck::Resync(size))
191 }
192 status if status.is_success() => Ok(LogAck::Accepted),
193 _ => Err(status_error(resp).await),
194 }
195 }
196
197 /// `POST /api/v1/ci/jobs/{id}/heartbeat`.
198 pub async fn heartbeat(
199 &self,
200 job_id: &str,
201 job_token: &str,
202 ) -> Result<HeartbeatResponse, ClientError> {
203 let resp = self
204 .http
205 .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/heartbeat")))
206 .bearer_auth(job_token)
207 .timeout(REQUEST_TIMEOUT)
208 .send()
209 .await?;
210 json_or_error(resp).await
211 }
212
213 /// `PUT /api/v1/ci/jobs/{id}/artifacts/{name}`, streaming the file from disk.
214 pub async fn upload_artifact(
215 &self,
216 job_id: &str,
217 job_token: &str,
218 name: &str,
219 path: &Path,
220 ) -> Result<(), ClientError> {
221 let file = tokio::fs::File::open(path).await?;
222 let len = file.metadata().await?.len();
223 let stream = tokio_util::io::ReaderStream::new(file);
224 let resp = self
225 .http
226 .put(self.url(&format!("/api/v1/ci/jobs/{job_id}/artifacts/{name}")))
227 .bearer_auth(job_token)
228 .header(CONTENT_LENGTH, len)
229 .body(reqwest::Body::wrap_stream(stream))
230 .timeout(ARTIFACT_TIMEOUT)
231 .send()
232 .await?;
233 if resp.status().is_success() {
234 Ok(())
235 } else {
236 Err(status_error(resp).await)
237 }
238 }
239
240 /// `POST /api/v1/ci/jobs/{id}/finish`.
241 pub async fn finish(
242 &self,
243 job_id: &str,
244 job_token: &str,
245 req: &FinishRequest,
246 ) -> Result<(), ClientError> {
247 let resp = self
248 .http
249 .post(self.url(&format!("/api/v1/ci/jobs/{job_id}/finish")))
250 .bearer_auth(job_token)
251 .json(req)
252 .timeout(REQUEST_TIMEOUT)
253 .send()
254 .await?;
255 if resp.status().is_success() {
256 Ok(())
257 } else {
258 Err(status_error(resp).await)
259 }
260 }
261}
262
263/// The `Forge-Binary-Sha256` header, trimmed and lowercased. An empty value counts as absent.
264fn advertised_sha256(headers: &HeaderMap) -> Option<String> {
265 headers
266 .get(HEADER_BINARY_SHA256)?
267 .to_str()
268 .ok()
269 .map(|v| v.trim().to_ascii_lowercase())
270 .filter(|v| !v.is_empty())
271}
272
273async fn json_or_error<T: DeserializeOwned>(resp: reqwest::Response) -> Result<T, ClientError> {
274 if resp.status().is_success() {
275 Ok(resp.json().await?)
276 } else {
277 Err(status_error(resp).await)
278 }
279}
280
281async fn status_error(resp: reqwest::Response) -> ClientError {
282 let status = resp.status();
283 let body: String = resp
284 .text()
285 .await
286 .unwrap_or_default()
287 .chars()
288 .take(512)
289 .collect();
290 ClientError::Status { status, body }
291}