feat: add basic auth (#101)

korniltsev · web-flow · commit fc5c0f1eb7ce · 2023-04-23T21:24:55.000+07:00
diff --git a/pyroscope_cli/src/cli/lib.rs b/pyroscope_cli/src/cli/lib.rs
@@ -64,6 +64,20 @@ enum Commands {
             help = "Authentication token used when uploading profiling data"
         )]
         auth_token: Option<String>,
+        #[clap(
+            name = "basic_auth_username",
+            long = "basic-auth-username",
+            value_name = "BASIC_AUTH_USERNAME",
+            help = "HTTP Basic Authentication username used when uploading profiling data"
+        )]
+        basic_auth_username: Option<String>,
+        #[clap(
+            name = "basic_auth_password",
+            long = "basic-auth-password",
+            value_name = "BASIC_AUTH_PASSWORD",
+            help = "HTTP Basic Authentication password used when uploading profiling data"
+        )]
+        basic_auth_password: Option<String>,
         #[clap(
             name = "scope_org_id",
             long = "scope_org_id",
@@ -197,6 +211,20 @@ enum Commands {
             help = "Authentication token used when uploading profiling data"
         )]
         auth_token: Option<String>,
+        #[clap(
+            name = "basic_auth_username",
+            long = "basic-auth-username",
+            value_name = "BASIC_AUTH_USERNAME",
+            help = "HTTP Basic Authentication username used when uploading profiling data"
+        )]
+        basic_auth_username: Option<String>,
+        #[clap(
+            name = "basic_auth_password",
+            long = "basic-auth-password",
+            value_name = "BASIC_AUTH_PASSWORD",
+            help = "HTTP Basic Authentication password used when uploading profiling data"
+        )]
+        basic_auth_password: Option<String>,
         #[clap(
             name = "scope_org_id",
             long = "scope_org_id",
diff --git a/pyroscope_cli/src/core/profiler.rs b/pyroscope_cli/src/core/profiler.rs
@@ -25,6 +25,8 @@ impl Profiler {
         let app_name: String = AppConfig::get::<String>("application_name")?;
 
         let auth_token: String = AppConfig::get::<String>("auth_token")?;
+        let basic_auth_username: String = AppConfig::get::<String>("basic_auth_username")?;
+        let basic_auth_password: String = AppConfig::get::<String>("basic_auth_password")?;
 
         let scope_org_id: String = AppConfig::get::<String>("scope_org_id").unwrap_or("".to_string());
 
@@ -70,6 +72,8 @@ impl Profiler {
                 // There must be a better way to do this, hopefully as clap supports Option<String>
                 if auth_token.len() > 0 {
                     builder = builder.auth_token(auth_token);
+                } else if basic_auth_username != "" && basic_auth_password != "" {
+                    builder = builder.basic_auth(basic_auth_username, basic_auth_password);
                 }
 
                 builder.backend(backend).tags(tags).build()?
diff --git a/pyroscope_cli/src/resources/default_config.toml b/pyroscope_cli/src/resources/default_config.toml
@@ -8,4 +8,6 @@ pyspy_native = false
 sample_rate = 100
 server_address = "http://localhost:4040"
 auth_token = ""
+basic_auth_username = ""
+basic_auth_password = ""
 debug = false
diff --git a/pyroscope_cli/src/utils/app_config.rs b/pyroscope_cli/src/utils/app_config.rs
@@ -21,6 +21,8 @@ pub struct AppConfig {
     pub spy_name: Spy,
     pub application_name: Option<String>,
     pub auth_token: Option<String>,
+    pub basic_auth_username: Option<String>,
+    pub basic_auth_password: Option<String>,
     pub detect_subprocesses: Option<bool>,
     pub no_logging: Option<bool>,
     pub log_level: LogLevel,
@@ -119,6 +121,16 @@ impl AppConfig {
                     AppConfig::set("auth_token", auth_token)?;
                 }
             }
+            if sub_connect.is_present("basic_auth_username") {
+                if let Some(basic_auth_username) = sub_connect.value_of("basic_auth_username") {
+                    AppConfig::set("basic_auth_username", basic_auth_username)?;
+                }
+            }
+            if sub_connect.is_present("basic_auth_password") {
+                if let Some(basic_auth_password) = sub_connect.value_of("basic_auth_password") {
+                    AppConfig::set("basic_auth_password", basic_auth_password)?;
+                }
+            }
             if sub_connect.is_present("scope_org_id") {
                 if let Some(scope_org_id) = sub_connect.value_of("scope_org_id") {
                     AppConfig::set("scope_org_id", scope_org_id)?;
@@ -210,6 +222,16 @@ impl AppConfig {
                     AppConfig::set("auth_token", auth_token)?;
                 }
             }
+            if sub_exec.is_present("basic_auth_username") {
+                if let Some(basic_auth_username) = sub_exec.value_of("basic_auth_username") {
+                    AppConfig::set("basic_auth_username", basic_auth_username)?;
+                }
+            }
+            if sub_exec.is_present("basic_auth_password") {
+                if let Some(basic_auth_password) = sub_exec.value_of("basic_auth_password") {
+                    AppConfig::set("basic_auth_password", basic_auth_password)?;
+                }
+            }
             if sub_exec.is_present("scope_org_id") {
                 if let Some(scope_org_id) = sub_exec.value_of("scope_org_id") {
                     AppConfig::set("scope_org_id", scope_org_id)?;
diff --git a/pyroscope_ffi/python/lib/include/pyroscope_ffi.h b/pyroscope_ffi/python/lib/include/pyroscope_ffi.h
@@ -17,6 +17,8 @@ bool initialize_logging(uint32_t logging_level);
 bool initialize_agent(const char *application_name,
                       const char *server_address,
                       const char *auth_token,
+                      const char *basic_auth_username,
+                      const char *basic_auth_password,
                       uint32_t sample_rate,
                       bool detect_subprocesses,
                       bool oncpu,
diff --git a/pyroscope_ffi/python/lib/src/lib.rs b/pyroscope_ffi/python/lib/src/lib.rs
@@ -40,9 +40,20 @@ pub extern "C" fn initialize_logging(logging_level: u32) -> bool {
 
 #[no_mangle]
 pub extern "C" fn initialize_agent(
-    application_name: *const c_char, server_address: *const c_char, auth_token: *const c_char,
-    sample_rate: u32, detect_subprocesses: bool, oncpu: bool, native: bool, gil_only: bool,
-    report_pid: bool, report_thread_id: bool, report_thread_name: bool, tags: *const c_char,
+    application_name: *const c_char,
+    server_address: *const c_char,
+    auth_token: *const c_char,
+    basic_auth_username: *const c_char,
+    basic_auth_password: *const c_char,
+    sample_rate: u32,
+    detect_subprocesses: bool,
+    oncpu: bool,
+    native: bool,
+    gil_only: bool,
+    report_pid: bool,
+    report_thread_id: bool,
+    report_thread_name: bool,
+    tags: *const c_char,
     scope_org_id: *const c_char,
     http_headers_json: *const c_char,
 ) -> bool {
@@ -61,12 +72,21 @@ pub extern "C" fn initialize_agent(
         .unwrap()
         .to_string();
 
-    // auth_token
     let auth_token = unsafe { CStr::from_ptr(auth_token) }
         .to_str()
         .unwrap()
         .to_string();
 
+    let basic_auth_username = unsafe { CStr::from_ptr(basic_auth_username) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
+    let basic_auth_password = unsafe { CStr::from_ptr(basic_auth_password) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
     // tags
     let tags_string = unsafe { CStr::from_ptr(tags) }
         .to_str()
@@ -124,6 +144,8 @@ pub extern "C" fn initialize_agent(
     // Add the auth token if it is not empty.
     if auth_token != "" {
         agent_builder = agent_builder.auth_token(auth_token);
+    } else if basic_auth_username != "" && basic_auth_password != "" {
+        agent_builder = agent_builder.basic_auth(basic_auth_username, basic_auth_password);
     }
     if scope_org_id != "" {
         agent_builder = agent_builder.scope_org_id(scope_org_id);
diff --git a/pyroscope_ffi/python/pyroscope/__init__.py b/pyroscope_ffi/python/pyroscope/__init__.py
@@ -12,6 +12,8 @@ def configure(
         application_name=None,
         server_address="http://localhost:4040",
         auth_token="",
+        basic_auth_username="",
+        basic_auth_password="",
         enable_logging=False,
         sample_rate=100,
         detect_subprocesses=False,
@@ -39,6 +41,8 @@ def configure(
         application_name.encode("UTF-8"),
         server_address.encode("UTF-8"),
         auth_token.encode("UTF-8"),
+        basic_auth_username.encode("UTF-8"),
+        basic_auth_password.encode("UTF-8"),
         sample_rate,
         detect_subprocesses,
         oncpu,
diff --git a/pyroscope_ffi/ruby/ext/rbspy/src/lib.rs b/pyroscope_ffi/ruby/ext/rbspy/src/lib.rs
@@ -110,9 +110,18 @@ pub extern "C" fn initialize_logging(logging_level: u32) -> bool {
 
 #[no_mangle]
 pub extern "C" fn initialize_agent(
-    application_name: *const c_char, server_address: *const c_char, auth_token: *const c_char,
-    sample_rate: u32, detect_subprocesses: bool, oncpu: bool, report_pid: bool,
-    report_thread_id: bool, tags: *const c_char, compression: *const c_char,
+    application_name: *const c_char,
+    server_address: *const c_char,
+    auth_token: *const c_char,
+    basic_auth_user: *const c_char,
+    basic_auth_password: *const c_char,
+    sample_rate: u32,
+    detect_subprocesses: bool,
+    oncpu: bool,
+    report_pid: bool,
+    report_thread_id: bool,
+    tags: *const c_char,
+    compression: *const c_char,
     report_encoding: *const c_char,
     scope_org_id: *const c_char,
     http_headers_json: *const c_char,
@@ -135,6 +144,16 @@ pub extern "C" fn initialize_agent(
         .unwrap()
         .to_string();
 
+    let basic_auth_user = unsafe { CStr::from_ptr(basic_auth_user) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
+    let basic_auth_password = unsafe { CStr::from_ptr(basic_auth_password) }
+        .to_str()
+        .unwrap()
+        .to_string();
+
     let tags_string = unsafe { CStr::from_ptr(tags) }
         .to_str()
         .unwrap()
@@ -186,6 +205,8 @@ pub extern "C" fn initialize_agent(
 
     if auth_token != "" {
         agent_builder = agent_builder.auth_token(auth_token);
+    } else if basic_auth_user != "" && basic_auth_password != "" {
+        agent_builder = agent_builder.basic_auth(basic_auth_user, basic_auth_password);
     }
 
     if scope_org_id != "" {
diff --git a/pyroscope_ffi/ruby/lib/pyroscope.rb b/pyroscope_ffi/ruby/lib/pyroscope.rb
@@ -9,7 +9,7 @@ module Rust
     extend FFI::Library
     ffi_lib File.expand_path(File.dirname(__FILE__)) + "/rbspy/rbspy.#{RbConfig::CONFIG["DLEXT"]}"
     attach_function :initialize_logging, [:int], :bool
-    attach_function :initialize_agent, [:string, :string, :string, :int, :bool, :bool, :bool, :bool, :string, :string, :string, :string, :string], :bool
+    attach_function :initialize_agent, [:string, :string, :string, :string, :string, :int, :bool, :bool, :bool, :bool, :string, :string, :string, :string, :string], :bool
     attach_function :add_thread_tag, [:uint64, :string, :string], :bool
     attach_function :remove_thread_tag, [:uint64, :string, :string], :bool
     attach_function :add_global_tag, [:string, :string], :bool
@@ -38,6 +38,8 @@ class Engine < ::Rails::Engine
     :app_name,
     :server_address,
     :auth_token,
+    :basic_auth_username,
+    :basic_auth_password,
     :log_level,
     :sample_rate,
     :detect_subprocesses,
@@ -57,6 +59,8 @@ def initialize(*)
       self.application_name = ''
       self.server_address = 'http://localhost:4040'
       self.auth_token = ''
+      self.basic_auth_username = ''
+      self.basic_auth_password = ''
       self.sample_rate = 100
       self.detect_subprocesses = false
       self.oncpu = true
@@ -106,6 +110,8 @@ def configure
         @config.app_name || @config.application_name || "",
         @config.server_address || "",
         @config.auth_token || "",
+        @config.basic_auth_username || "",
+        @config.basic_auth_password || "",
         @config.sample_rate || 100,
         @config.detect_subprocesses || false,
         @config.oncpu || false,
diff --git a/src/pyroscope.rs b/src/pyroscope.rs
@@ -47,6 +47,7 @@ pub struct PyroscopeConfig {
     pub spy_name: String,
     /// Authentication Token
     pub auth_token: Option<String>,
+    pub basic_auth: Option<BasicAuth>,
     /// Function to apply
     pub func: Option<fn(Report) -> Report>,
     /// Pyroscope http request body compression
@@ -56,6 +57,12 @@ pub struct PyroscopeConfig {
     pub http_headers: HashMap<String, String>,
 }
 
+#[derive(Clone, Debug)]
+pub struct BasicAuth {
+    pub username: String,
+    pub password: String,
+}
+
 impl Default for PyroscopeConfig {
     fn default() -> Self {
         Self {
@@ -68,6 +75,7 @@ impl Default for PyroscopeConfig {
             sample_rate: 100u32,
             spy_name: "undefined".to_string(),
             auth_token: None,
+            basic_auth: None,
             func: None,
             compression: None,
             report_encoding: ReportEncoding::FOLDED,
@@ -93,6 +101,7 @@ impl PyroscopeConfig {
             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
+            basic_auth: None,
             func: None,                   // No function
             compression: None,
             report_encoding: ReportEncoding::FOLDED,
@@ -130,14 +139,21 @@ 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
         }
     }
 
+    pub fn basic_auth(self, username: String, password: String) -> Self {
+        Self {
+            basic_auth: Some(BasicAuth { username, password }),
+            ..self
+        }
+    }
+
     /// Set the Function.
     pub fn func(self, func: fn(Report) -> Report) -> Self {
         Self {
@@ -310,6 +326,13 @@ impl PyroscopeAgentBuilder {
         }
     }
 
+    pub fn basic_auth(self, username: impl AsRef<str>, password: impl AsRef<str>) -> Self {
+        Self {
+            config: self.config.basic_auth(username.as_ref().to_owned(), password.as_ref().to_owned()),
+            ..self
+        }
+    }
+
     /// Set the Function.
     /// This is optional. If not set, the agent will not apply any function.
     /// #Example
diff --git a/src/session.rs b/src/session.rs
@@ -220,6 +220,8 @@ impl Session {
 
         if let Some(auth_token) = &self.config.auth_token {
             req_builder = req_builder.bearer_auth(auth_token);
+        } else if let Some(basic_auth) = &self.config.basic_auth {
+            req_builder = req_builder.basic_auth(basic_auth.username.clone(), Some(basic_auth.password.clone()));
         }
         if report.content_encoding != "" {
             req_builder = req_builder.header("Content-Encoding", report.content_encoding.as_str());
