imp(lib): add authentication token support. close #16

Author: omarabid

examples/auth.rs

@@ -0,0 +1,64 @@
+extern crate pyroscope;
+
+use pyroscope::{PyroscopeAgent, Result};
+use pyroscope_pprofrs::{pprof_backend, PprofConfig};
+use std::hash::{Hash, Hasher};
+
+fn hash_rounds(n: u64) -> u64 {
+    let hash_str = "Some string to hash";
+    let mut default_hasher = std::collections::hash_map::DefaultHasher::new();
+
+    for _ in 0..n {
+        for _ in 0..1000 {
+            default_hasher.write(hash_str.as_bytes());
+        }
+        hash_str.hash(&mut default_hasher);
+    }
+
+    n
+}
+
+fn main() -> Result<()> {
+    // TODO: Change this token to your own.
+    let token = String::from("<your-token>");
+
+    let agent = PyroscopeAgent::builder("http://localhost:4040", "example.basic")
+        .backend(pprof_backend(PprofConfig::new().sample_rate(100)))
+        .auth_token(token)
+        .tags([("TagA", "ValueA"), ("TagB", "ValueB")].to_vec())
+        .build()?;
+
+    // Show start time
+    let start = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+    println!("Start Time: {}", start);
+
+    // Start Agent
+    let agent_running = agent.start()?;
+
+    let _result = hash_rounds(300_000);
+
+    // Show stop time
+    let stop = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+    println!("Stop Time: {}", stop);
+
+    // Stop Agent
+    let agent_ready = agent_running.stop()?;
+
+    // Shutdown the Agent
+    agent_ready.shutdown();
+
+    // Show program exit time
+    let exit = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+    println!("Exit Time: {}", exit);
+
+    Ok(())
+}

src/pyroscope.rs

@@ -28,7 +28,7 @@ const LOG_TAG: &str = "Pyroscope::Agent";
 /// use pyroscope::pyroscope::PyroscopeConfig;
 /// let config = PyroscopeConfig::new("http://localhost:8080", "my-app");
 /// ```
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct PyroscopeConfig {
     /// Pyroscope Server Address
     pub url: String,
@@ -40,6 +40,8 @@ pub struct PyroscopeConfig {
     pub sample_rate: u32,
     /// Spy Name
     pub spy_name: String,
+    /// Authentication Token
+    pub auth_token: Option<String>,
 }
 
 impl PyroscopeConfig {
@@ -57,6 +59,7 @@ impl PyroscopeConfig {
             tags: HashMap::new(),         // Empty tags
             sample_rate: 100u32,          // Default sample rate
             spy_name: String::from("undefined"), // Spy Name should be set by the backend
+            auth_token: None,             // No authentication token
         }
     }
 
@@ -73,6 +76,14 @@ impl PyroscopeConfig {
         Self { spy_name, ..self }
     }
 
+    /// Set the Authentication Token.
+    pub fn auth_token(self, auth_token: String) -> Self {
+        Self {
+            auth_token: Some(auth_token),
+            ..self
+        }
+    }
+
     /// Set the tags.
     ///
     /// # Example
@@ -143,6 +154,23 @@ impl PyroscopeAgentBuilder {
         Self { backend, ..self }
     }
 
+    /// Set JWT authentication token.
+    /// This is optional. If not set, the agent will not send any authentication token.
+    ///
+    /// #Example
+    /// ```ignore
+    /// let builder = PyroscopeAgentBuilder::new("http://localhost:8080", "my-app")
+    /// .auth_token("my-token")
+    /// .build()
+    /// ?;
+    /// ```
+    pub fn auth_token(self, auth_token: impl AsRef<str>) -> Self {
+        Self {
+            config: self.config.auth_token(auth_token.as_ref().to_owned()),
+            ..self
+        }
+    }
+
     /// Set tags. Default is empty.
     ///
     /// # Example

src/session.rs

@@ -168,10 +168,22 @@ impl Session {
             report.metadata.tags.clone().into_iter().collect(),
         )?;
 
-        // Create and send the request
-        client
+        // Create Reqwest builder
+        let mut req_builder = client
             .post(format!("{}/ingest", url))
-            .header("Content-Type", "binary/octet-stream")
+            .header("Content-Type", "binary/octet-stream");
+
+        // Set authentication token
+        //if self.config.auth_token.is_some() {
+        //req_builder = req_builder.bearer_auth(self.config.auth_token.clone().unwrap());
+        //}
+        // rewrite with let some
+        if let Some(auth_token) = self.config.auth_token.clone() {
+            req_builder = req_builder.bearer_auth(auth_token);
+        }
+
+        // Send the request
+        req_builder
             .query(&[
                 ("name", application_name.as_str()),
                 ("from", &format!("{}", self.from)),

tests/session.rs

@@ -26,6 +26,7 @@ fn test_session_new() {
         tags: HashMap::new(),
         sample_rate: 100u32,
         spy_name: "test-rs".to_string(),
+        ..Default::default()
     };
 
     let report = vec![Report::new(HashMap::new())];
@@ -44,6 +45,7 @@ fn test_session_send_error() {
         tags: HashMap::new(),
         sample_rate: 100u32,
         spy_name: "test-rs".to_string(),
+        ..Default::default()
     };
 
     let report = vec![Report::new(HashMap::new())];