Add support for token prefix masking

gitlab-runner-mock/src/api/request.rs

@@ -107,6 +107,46 @@ impl Respond for JobRequestResponder {
                     value
                 })
                 .collect();
+
+            let failure_reasons = json!([
+                "unknown_failure",
+                "script_failure",
+                "api_failure",
+                "stuck_or_timeout_failure",
+                "runner_system_failure",
+                "missing_dependency_failure",
+                "runner_unsupported",
+                "stale_schedule",
+                "job_execution_timeout",
+                "archived_failure",
+                "unmet_prerequisites",
+                "scheduler_failure",
+                "data_integrity_failure",
+                "forward_deployment_failure",
+                "user_blocked",
+                "project_deleted",
+                "insufficient_bridge_permissions",
+                "downstream_bridge_project_not_found",
+                "invalid_bridge_trigger",
+                "bridge_pipeline_is_child_pipeline",
+                "downstream_pipeline_creation_failed",
+                "secrets_provider_not_found",
+                "reached_max_descendant_pipelines_depth"
+            ]);
+
+            let features = if !job.token_prefixes().is_empty() {
+                json!({
+                    "trace_sections": true,
+                    "token_mask_prefixes": job.token_prefixes(),
+                    "failure_reasons": failure_reasons
+                })
+            } else {
+                json!({
+                    "trace_sections": true,
+                    "failure_reasons": failure_reasons
+                })
+            };
+
             ResponseTemplate::new(201).set_body_json(json!({
                 "id": job.id(),
                 "token": job.token(),
@@ -155,34 +195,7 @@ impl Respond for JobRequestResponder {
                     }
                 ],
                 "dependencies": dependencies,
-                "features": {
-                    "trace_sections": true,
-                    "failure_reasons": [
-                        "unknown_failure",
-                        "script_failure",
-                        "api_failure",
-                        "stuck_or_timeout_failure",
-                        "runner_system_failure",
-                        "missing_dependency_failure",
-                        "runner_unsupported",
-                        "stale_schedule",
-                        "job_execution_timeout",
-                        "archived_failure",
-                        "unmet_prerequisites",
-                        "scheduler_failure",
-                        "data_integrity_failure",
-                        "forward_deployment_failure",
-                        "user_blocked",
-                        "project_deleted",
-                        "insufficient_bridge_permissions",
-                        "downstream_bridge_project_not_found",
-                        "invalid_bridge_trigger",
-                        "bridge_pipeline_is_child_pipeline",
-                        "downstream_pipeline_creation_failed",
-                        "secrets_provider_not_found",
-                        "reached_max_descendant_pipelines_depth"
-                    ]
-                }
+                "features": features
             }))
         } else {
             ResponseTemplate::new(StatusCode::NoContent)

gitlab-runner-mock/src/job.rs

@@ -117,6 +117,7 @@ pub struct MockJob {
     steps: Vec<MockJobStep>,
     dependencies: Vec<MockJob>,
     artifacts: Vec<MockJobArtifact>,
+    token_prefixes: Vec<String>,
     inner: Arc<Mutex<MockJobInner>>,
 }
 
@@ -143,6 +144,7 @@ impl MockJob {
             steps: Vec::new(),
             dependencies: Vec::new(),
             artifacts: Vec::new(),
+            token_prefixes: Vec::new(),
             inner: Arc::new(Mutex::new(MockJobInner {
                 state: MockJobState::Success,
                 state_updates: 2,
@@ -177,6 +179,10 @@ impl MockJob {
         &self.variables
     }
 
+    pub fn token_prefixes(&self) -> &[String] {
+        &self.token_prefixes
+    }
+
     pub fn steps(&self) -> &[MockJobStep] {
         &self.steps
     }
@@ -265,6 +271,7 @@ pub struct MockJobBuilder {
     variables: HashMap<String, MockJobVariable>,
     steps: Vec<MockJobStep>,
     dependencies: Vec<MockJob>,
+    token_prefixes: Vec<String>,
     artifacts: Vec<MockJobArtifact>,
 }
 
@@ -293,6 +300,10 @@ impl MockJobBuilder {
         );
     }
 
+    pub fn add_token_prefix(&mut self, prefix: String) {
+        self.token_prefixes.push(prefix);
+    }
+
     pub fn add_step(
         &mut self,
         name: MockJobStepName,
@@ -373,6 +384,7 @@ impl MockJobBuilder {
             variables: self.variables.into_values().collect(),
             dependencies: self.dependencies,
             artifacts: self.artifacts,
+            token_prefixes: self.token_prefixes,
             inner,
         }
     }

gitlab-runner-mock/src/lib.rs

@@ -134,6 +134,15 @@ impl GitlabRunnerMock {
         jobs.jobs.push(job);
     }
 
+    pub fn get_job_artifact(&self, id: u64) -> Option<Vec<u8>> {
+        let jobs = self.inner.jobs.lock().unwrap();
+
+        jobs.jobs
+            .iter()
+            .find(|j| j.id() == id)
+            .map(|j| j.artifact().as_slice().to_vec())
+    }
+
     fn grab_pending_job(&self) -> Option<MockJob> {
         let jobs = self.inner.jobs.lock().unwrap();
         for job in jobs.jobs.iter() {

gitlab-runner-mock/src/variables.rs

@@ -28,7 +28,7 @@ pub fn default_job_variables(job_id: u64) -> Vec<MockJobVariable> {
         },
         MockJobVariable {
             key: "CI_JOB_TOKEN".to_owned(),
-            value: "tokn".to_owned(),
+            value: "job-token".to_owned(),
             public: false,
             masked: true,
         },
@@ -46,7 +46,7 @@ pub fn default_job_variables(job_id: u64) -> Vec<MockJobVariable> {
         },
         MockJobVariable {
             key: "CI_BUILD_TOKEN".to_owned(),
-            value: "tokn".to_owned(),
+            value: "build-token".to_owned(),
             public: false,
             masked: true,
         },
@@ -58,7 +58,7 @@ pub fn default_job_variables(job_id: u64) -> Vec<MockJobVariable> {
         },
         MockJobVariable {
             key: "CI_REGISTRY_PASSWORD".to_owned(),
-            value: "token".to_owned(),
+            value: "registry-password".to_owned(),
             public: false,
             masked: true,
         },
@@ -77,7 +77,7 @@ pub fn default_job_variables(job_id: u64) -> Vec<MockJobVariable> {
         },
         MockJobVariable {
             key: "CI_DEPENDENCY_PROXY_PASSWORD".to_owned(),
-            value: "token".to_owned(),
+            value: "proxy-password".to_owned(),
             public: false,
             masked: true,
         },

gitlab-runner/Cargo.toml

@@ -29,6 +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.3"
 
 [dev-dependencies]
 tokio = { version = "1.5.0", features = [ "full", "test-util" ] }

gitlab-runner/src/client.rs

@@ -168,6 +168,16 @@ pub(crate) struct JobDependency {
     pub artifacts_file: Option<JobArtifactFile>,
 }
 
+#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
+pub(crate) struct GitlabFeatures {
+    #[serde(default)]
+    pub trace_sections: bool,
+    #[serde(default, deserialize_with = "deserialize_null_default")]
+    pub token_mask_prefixes: Vec<String>,
+    #[serde(default, deserialize_with = "deserialize_null_default")]
+    pub failure_reasons: Vec<String>,
+}
+
 #[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
 pub(crate) struct JobResponse {
     pub id: u64,
@@ -180,6 +190,8 @@ pub(crate) struct JobResponse {
     pub dependencies: Vec<JobDependency>,
     #[serde(deserialize_with = "deserialize_null_default")]
     pub artifacts: Vec<JobArtifact>,
+    #[serde(default)]
+    pub features: Option<GitlabFeatures>,
     #[serde(flatten)]
     unparsed: JsonValue,
 }

gitlab-runner/src/run.rs

@@ -1,4 +1,5 @@
 use bytes::Bytes;
+use masker::{Masker, MatchData};
 use std::future::Future;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -15,10 +16,15 @@ use crate::uploader::Uploader;
 use crate::CancellableJobHandler;
 use crate::{JobResult, Phase};
 
+const GITLAB_MASK: &str = "[MASKED]";
+const GITLAB_TOKEN_SUFFIX_CHARS: &str =
+    "-.0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz=";
+
 async fn run<F, J, Ret>(
     job: Job,
     client: Client,
     response: Arc<JobResponse>,
+    masker: Masker,
     process: F,
     build_dir: PathBuf,
     cancel_token: CancellationToken,
@@ -62,7 +68,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)
@@ -140,7 +146,13 @@ impl Run {
         buf: Bytes,
         cancel_token: &CancellationToken,
     ) -> Option<Duration> {
-        assert!(!buf.is_empty());
+        if buf.is_empty() {
+            // It's convenient to permit this because if we are
+            // masking, the masker may not have produced any output,
+            // so we'd just end up doing the same test in every
+            // caller, rather than once here.
+            return None;
+        }
         let len = buf.len();
 
         match self
@@ -178,6 +190,30 @@ 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 prefixes = self
+            .response
+            .features
+            .iter()
+            .flat_map(|x| x.token_mask_prefixes.iter())
+            // This matches the behaviour of the gitlab runner, which
+            // explicitly supports a maximum of 10 prefixes.
+            .take(10)
+            .map(|p| MatchData {
+                prefix: p.trim().as_bytes(),
+                suffix: GITLAB_TOKEN_SUFFIX_CHARS.as_bytes(),
+                mask_prefix: false,
+            })
+            .collect::<Vec<_>>();
+
+        let masker = Masker::new_with_match_data(&masked_variables, &prefixes, GITLAB_MASK);
+
         let job = Job::new(
             self.client.clone(),
             self.response.clone(),
@@ -189,6 +225,7 @@ impl Run {
                 job,
                 self.client.clone(),
                 self.response.clone(),
+                masker.clone(),
                 process,
                 build_dir,
                 cancel_token.clone(),
@@ -198,6 +235,8 @@ impl Run {
         );
         tokio::pin!(join);
 
+        let mut cm = masker.mask_chunks();
+
         let result = loop {
             tokio::select! {
                 _ = self.interval.tick() => {
@@ -206,6 +245,7 @@ impl Run {
                     let now = Instant::now();
                     if let Some(buf) = self.joblog.split_trace() {
                         // TODO be resiliant against send errors
+                        let buf = cm.mask_chunk(buf).into();
                         if let Some(interval) = self.send_trace(buf, &cancel_token).await {
                             if interval != self.interval.period() {
                                 self.interval = Self::create_interval(now, interval);
@@ -224,8 +264,12 @@ impl Run {
 
         // Send the remaining trace buffer back to gitlab.
         if let Some(buf) = self.joblog.split_trace() {
+            let buf = cm.mask_chunk(buf).into();
             self.send_trace(buf, &cancel_token).await;
         }
+        // Flush anything the masker was holding back
+        let buf = cm.finish().into();
+        self.send_trace(buf, &cancel_token).await;
 
         // Don't bother updating the status if cancelled, since it will just fail.
         if !cancel_token.is_cancelled() {