Allow client crates to upload masked files

Although we already automatically mask the trace sent to Gitlab, some
runners will produce additional log files or metadata, which must also
be masked to prevent secrets from being revealed in the artifact data.

This permits client crates to identify a file they wish to upload as
additionally requiring masking.

gitlab-runner/Cargo.toml

@@ -29,7 +29,7 @@ tracing-subscriber = "0.3.8"
 tracing = "0.1.30"
 doc-comment = "0.3.3"
 tokio-util = { version = "0.7", features = [ "io" ] }
-masker = "0.0.1"
+masker = { path = "../../masker" }
 
 [dev-dependencies]
 tokio = { version = "1.5.0", features = [ "full", "test-util" ] }

gitlab-runner/src/run.rs

@@ -22,6 +22,7 @@ async fn run<F, J, Ret>(
     job: Job,
     client: Client,
     response: Arc<JobResponse>,
+    masker: Masker,
     process: F,
     build_dir: PathBuf,
     cancel_token: CancellationToken,
@@ -65,7 +66,7 @@ where
         });
 
     let r = if upload {
-        if let Ok(mut uploader) = Uploader::new(client, &build_dir, response) {
+        if let Ok(mut uploader) = Uploader::new(client, &build_dir, response, masker) {
             let r = handler.upload_artifacts(&mut uploader).await;
             if r.is_ok() {
                 uploader.upload().await.and(script_result)
@@ -187,6 +188,15 @@ impl Run {
     {
         let cancel_token = CancellationToken::new();
 
+        let masked_variables = self
+            .response
+            .variables
+            .iter()
+            .filter(|(_, v)| v.masked)
+            .map(|(_, v)| v.value.as_str())
+            .collect::<Vec<_>>();
+        let masker = Masker::new(&masked_variables, GITLAB_MASK);
+
         let job = Job::new(
             self.client.clone(),
             self.response.clone(),
@@ -198,6 +208,7 @@ impl Run {
                 job,
                 self.client.clone(),
                 self.response.clone(),
+                masker.clone(),
                 process,
                 build_dir,
                 cancel_token.clone(),
@@ -207,15 +218,6 @@ impl Run {
         );
         tokio::pin!(join);
 
-        let masked_variables = self
-            .response
-            .variables
-            .iter()
-            .filter(|(_, v)| v.masked)
-            .map(|(_, v)| v.value.as_str())
-            .collect::<Vec<_>>();
-
-        let masker = Masker::new(&masked_variables, GITLAB_MASK);
         let mut cm = masker.mask_chunks();
 
         let result = loop {

gitlab-runner/src/uploader.rs

@@ -8,6 +8,7 @@ use std::thread;
 use std::{sync::Arc, task::Poll};
 
 use futures::{future::BoxFuture, AsyncWrite, FutureExt};
+use masker::{ChunkMasker, Masker};
 use reqwest::Body;
 use tokio::fs::File as AsyncFile;
 use tokio::sync::mpsc::{self, error::SendError};
@@ -71,6 +72,7 @@ fn zip_thread(mut temp: File, mut rx: mpsc::Receiver<UploadRequest>) {
 pub struct UploadFile<'a> {
     tx: &'a mpsc::Sender<UploadRequest>,
     state: UploadFileState<'a>,
+    masker: Option<ChunkMasker<'a>>,
 }
 
 impl<'a> AsyncWrite for UploadFile<'a> {
@@ -84,10 +86,12 @@ impl<'a> AsyncWrite for UploadFile<'a> {
             match this.state {
                 UploadFileState::Idle => {
                     let (tx, rx) = oneshot::channel();
-                    let send = this
-                        .tx
-                        .send(UploadRequest::WriteData(Vec::from(buf), tx))
-                        .boxed();
+                    let buf = if let Some(masker) = &mut this.masker {
+                        masker.mask_chunk(buf)
+                    } else {
+                        Vec::from(buf)
+                    };
+                    let send = this.tx.send(UploadRequest::WriteData(buf, tx)).boxed();
                     this.state = UploadFileState::Writing(Some(send), rx)
                 }
                 UploadFileState::Writing(ref mut send, ref mut rx) => {
@@ -114,9 +118,33 @@ impl<'a> AsyncWrite for UploadFile<'a> {
 
     fn poll_close(
         self: std::pin::Pin<&mut Self>,
-        _cx: &mut std::task::Context<'_>,
+        cx: &mut std::task::Context<'_>,
     ) -> std::task::Poll<std::io::Result<()>> {
-        Poll::Ready(Ok(()))
+        let this = self.get_mut();
+        if let Some(masker) = this.masker.take() {
+            let (tx, mut rx) = oneshot::channel();
+            let buf = masker.finish();
+            let mut send = Some(this.tx.send(UploadRequest::WriteData(buf, tx)).boxed());
+
+            loop {
+                if let Some(mut f) = send {
+                    // Phase 1: Waiting for the send to the writer
+                    // thread to complete.
+
+                    // TODO error handling
+                    let _r = futures::ready!(f.as_mut().poll(cx));
+                    send = None;
+                } else {
+                    // Phase 2: Waiting for the writer thread to
+                    // signal write completion.
+
+                    let _r = futures::ready!(Pin::new(&mut rx).poll(cx));
+                    return Poll::Ready(Ok(()));
+                }
+            }
+        } else {
+            Poll::Ready(Ok(()))
+        }
     }
 }
 
@@ -125,20 +153,27 @@ pub struct Uploader {
     client: Client,
     data: Arc<JobResponse>,
     tx: mpsc::Sender<UploadRequest>,
+    masker: Masker,
 }
 
 impl Uploader {
     pub(crate) fn new(
         client: Client,
         build_dir: &Path,
         data: Arc<JobResponse>,
+        masker: Masker,
     ) -> Result<Self, ()> {
         let temp = tempfile::tempfile_in(build_dir)
             .map_err(|e| warn!("Failed to create artifacts temp file: {:?}", e))?;
 
         let (tx, rx) = mpsc::channel(2);
         thread::spawn(move || zip_thread(temp, rx));
-        Ok(Self { client, data, tx })
+        Ok(Self {
+            client,
+            data,
+            tx,
+            masker,
+        })
     }
 
     /// Create a new file to be uploaded
@@ -149,6 +184,20 @@ impl Uploader {
             .expect("Failed to create file");
         UploadFile {
             tx: &self.tx,
+            masker: None,
+            state: UploadFileState::Idle,
+        }
+    }
+
+    /// Create a new file to be uploaded, which will be masked
+    pub async fn masked_file(&mut self, name: String) -> UploadFile<'_> {
+        self.tx
+            .send(UploadRequest::NewFile(name))
+            .await
+            .expect("Failed to create file");
+        UploadFile {
+            tx: &self.tx,
+            masker: Some(self.masker.mask_chunks()),
             state: UploadFileState::Idle,
         }
     }