datadog_agent_commons/ipc/client/
mod.rs

1//! Helpers for interacting with the Datadog Agent.
2
3use std::time::Duration;
4
5use backon::Retryable as _;
6use datadog_protos::agent::v1::{RefreshRemoteAgentRequest, RegisterRemoteAgentRequest, RegisterRemoteAgentResponse};
7use datadog_protos::agent::{
8    AgentClient, AgentSecureClient, AutodiscoveryStreamResponse, ConfigEvent, ConfigStreamRequest, EntityId,
9    FetchEntityRequest, HostTagReply, HostTagRequest, HostnameRequest, StreamTagsRequest, StreamTagsResponse,
10    TagCardinality, WorkloadmetaEventType, WorkloadmetaFilter, WorkloadmetaKind, WorkloadmetaSource,
11    WorkloadmetaStreamRequest, WorkloadmetaStreamResponse,
12};
13use saluki_config::GenericConfiguration;
14use saluki_error::{generic_error, ErrorContext as _, GenericError};
15use saluki_io::net::client::http::HttpsCapableConnectorBuilder;
16use tonic::{
17    service::interceptor::InterceptedService,
18    transport::{Channel, Endpoint},
19    Code, Request, Response,
20};
21use tracing::warn;
22
23use crate::ipc::{config::RemoteAgentClientConfiguration, session::SessionId, tls::build_ipc_client_ipc_tls_config};
24
25mod bearer_auth;
26use self::bearer_auth::BearerAuthInterceptor;
27
28mod streaming;
29pub use self::streaming::StreamingResponse;
30
31/// A client for interacting with the Datadog Agent's internal gRPC-based API.
32#[derive(Clone)]
33pub struct RemoteAgentClient {
34    client: AgentClient<InterceptedService<Channel, BearerAuthInterceptor>>,
35    secure_client: AgentSecureClient<InterceptedService<Channel, BearerAuthInterceptor>>,
36}
37
38impl RemoteAgentClient {
39    /// Creates a new `RemoteAgentClient` from the given configuration.
40    ///
41    /// # Errors
42    ///
43    /// If the Agent gRPC client can't be created (invalid API endpoint, missing authentication token, etc), or if the
44    /// authentication token is invalid, an error will be returned.
45    pub async fn from_configuration(config: &GenericConfiguration) -> Result<Self, GenericError> {
46        let config = RemoteAgentClientConfiguration::from_configuration(config)?;
47        Self::from_client_configuration(&config).await
48    }
49
50    /// Creates a new `RemoteAgentClient` from the given client configuration.
51    ///
52    /// # Errors
53    ///
54    /// If the Agent gRPC client can't be created (invalid API endpoint, missing authentication token, etc), or if the
55    /// authentication token is invalid, an error will be returned.
56    pub async fn from_client_configuration(config: &RemoteAgentClientConfiguration) -> Result<Self, GenericError> {
57        // TODO: We need to write a Tower middleware service that allows applying a backoff between failed calls,
58        // specifically so that we can throttle reconnection attempts.
59        //
60        // When the remote Agent endpoint is not available -- Agent isn't running, etc -- the gRPC client will
61        // essentially freewheel, trying to reconnect as quickly as possible, which spams the logs, wastes resources, so
62        // on and so forth. We would want to essentially apply a backoff like any other client would for the RPC calls
63        // themselves, but use it with the _connector_ instead.
64        //
65        // We could potentially just use a retry middleware, but Tonic does have its own reconnection logic, so we'd
66        // have to test it out to make sure it behaves sensibly.
67        let service_builder = || async {
68            let auth_interceptor = BearerAuthInterceptor::from_file(&config.auth().auth_token_file_path()).await?;
69            let ipc_cert_file_path = config.auth().ipc_cert_file_path();
70            let client_tls_config = build_ipc_client_ipc_tls_config(ipc_cert_file_path).await?;
71            let connector_builder = HttpsCapableConnectorBuilder::default();
72            #[cfg(target_os = "linux")]
73            let connector_builder = if let Some(addr) = config.vsock_addr()? {
74                connector_builder.with_vsock_addr(addr)
75            } else {
76                connector_builder
77            };
78            let https_connector = connector_builder.build(client_tls_config)?;
79            let endpoint = config.endpoint()?;
80            let channel = Endpoint::from(endpoint.clone())
81                .connect_timeout(Duration::from_secs(2))
82                .connect_with_connector(https_connector)
83                .await
84                .with_error_context(|| format!("Failed to connect to Datadog Agent API at '{}'.", endpoint))?;
85
86            Ok::<_, GenericError>(InterceptedService::new(channel, auth_interceptor))
87        };
88
89        let service = service_builder
90            .retry(config)
91            .notify(|e, delay| {
92                warn!(error = %e, "Failed to create Datadog Agent API client. Retrying in {:?}...", delay);
93            })
94            .await
95            .error_context("Failed to create Datadog Agent API client.")?;
96
97        let client = AgentClient::new(service.clone()).max_decoding_message_size(config.grpc_max_message_size());
98        let mut secure_client =
99            AgentSecureClient::new(service).max_decoding_message_size(config.grpc_max_message_size());
100
101        // Try and do a basic health check to make sure we can connect and that our authentication token is valid.
102        try_query_agent_api(&mut secure_client).await?;
103
104        Ok(Self { client, secure_client })
105    }
106
107    /// Gets the detected hostname from the Agent.
108    ///
109    /// # Errors
110    ///
111    /// If there is an error querying the Agent API, an error will be returned.
112    pub async fn get_hostname(&mut self) -> Result<String, GenericError> {
113        let response = self
114            .client
115            .get_hostname(HostnameRequest {})
116            .await
117            .map(|r| r.into_inner())?;
118
119        Ok(response.hostname)
120    }
121
122    /// Gets a stream of tagger entities at the given cardinality.
123    ///
124    /// If there is an error with the initial request, or an error occurs while streaming, the next message in the
125    /// stream will be `Some(Err(status))`, where the status indicates the underlying error.
126    pub fn get_tagger_stream(&mut self, cardinality: TagCardinality) -> StreamingResponse<StreamTagsResponse> {
127        let mut client = self.secure_client.clone();
128        StreamingResponse::from_response_future(async move {
129            client
130                .tagger_stream_entities(StreamTagsRequest {
131                    cardinality: cardinality.into(),
132                    ..Default::default()
133                })
134                .await
135        })
136    }
137
138    /// Gets a stream of all workloadmeta entities.
139    ///
140    /// If there is an error with the initial request, or an error occurs while streaming, the next message in the
141    /// stream will be `Some(Err(status))`, where the status indicates the underlying error.
142    pub fn get_workloadmeta_stream(&mut self) -> StreamingResponse<WorkloadmetaStreamResponse> {
143        let mut client = self.secure_client.clone();
144        StreamingResponse::from_response_future(async move {
145            client
146                .workloadmeta_stream_entities(WorkloadmetaStreamRequest {
147                    filter: Some(WorkloadmetaFilter {
148                        kinds: vec![
149                            WorkloadmetaKind::Container.into(),
150                            WorkloadmetaKind::KubernetesPod.into(),
151                            WorkloadmetaKind::EcsTask.into(),
152                        ],
153                        source: WorkloadmetaSource::All.into(),
154                        event_type: WorkloadmetaEventType::EventTypeAll.into(),
155                    }),
156                })
157                .await
158        })
159    }
160
161    /// Registers a Remote Agent with the Agent.
162    ///
163    /// # Errors
164    ///
165    /// If there is an error sending the request to the Agent API, an error will be returned.
166    pub async fn register_remote_agent_request(
167        &mut self, pid: u32, display_name: &str, flavor: &str, api_endpoint: &str, services: Vec<String>,
168    ) -> Result<Response<RegisterRemoteAgentResponse>, GenericError> {
169        let mut client = self.secure_client.clone();
170        let response = client
171            .register_remote_agent(RegisterRemoteAgentRequest {
172                pid: pid.to_string(),
173                flavor: flavor.to_string(),
174                display_name: display_name.to_string(),
175                api_endpoint_uri: api_endpoint.to_string(),
176                services,
177            })
178            .await?;
179        Ok(response)
180    }
181
182    /// Refreshes the given remote agent session with the Agent.
183    ///
184    /// # Errors
185    ///
186    /// If there is an error sending the request to the Agent API, an error will be returned.
187    pub async fn refresh_remote_agent_request(&mut self, session_id: &SessionId) -> Result<Response<()>, GenericError> {
188        let mut client = self.secure_client.clone();
189        let response = client
190            .refresh_remote_agent(RefreshRemoteAgentRequest {
191                session_id: session_id.to_string(),
192            })
193            .await?
194            .map(|_| ());
195        Ok(response)
196    }
197
198    /// Gets the host tags from the Agent.
199    ///
200    /// # Errors
201    ///
202    /// If there is an error querying the Agent API, an error will be returned.
203    pub async fn get_host_tags(&self) -> Result<Response<HostTagReply>, GenericError> {
204        let mut client = self.secure_client.clone();
205        let response = client.get_host_tags(HostTagRequest {}).await?;
206        Ok(response)
207    }
208
209    /// Gets a stream of autodiscovery config updates.
210    ///
211    /// If there is an error with the initial request, or an error occurs while streaming, the next message in the
212    /// stream will be `Some(Err(status))`, where the status indicates the underlying error.
213    pub fn get_autodiscovery_stream(&mut self) -> StreamingResponse<AutodiscoveryStreamResponse> {
214        let mut client = self.secure_client.clone();
215        StreamingResponse::from_response_future(async move { client.autodiscovery_stream_config(()).await })
216    }
217
218    /// Gets a stream of config events.
219    ///
220    /// If there is an error with the initial request, or an error occurs while streaming, the next message in the
221    /// stream will be `Some(Err(status))`, where the status indicates the underlying error.
222    pub fn stream_config_events(&mut self, session_id: &SessionId) -> StreamingResponse<ConfigEvent> {
223        let mut client = self.secure_client.clone();
224        let app_details = saluki_metadata::get_app_details();
225        let formatted_full_name = app_details
226            .full_name()
227            .replace(" ", "-")
228            .replace("_", "-")
229            .to_lowercase();
230
231        let mut request = Request::new(ConfigStreamRequest {
232            name: formatted_full_name,
233        });
234
235        request
236            .metadata_mut()
237            .insert("session_id", session_id.to_grpc_header_value());
238
239        StreamingResponse::from_response_future(async move { client.stream_config_events(request).await })
240    }
241}
242
243async fn try_query_agent_api(
244    client: &mut AgentSecureClient<InterceptedService<Channel, BearerAuthInterceptor>>,
245) -> Result<(), GenericError> {
246    let noop_fetch_request = FetchEntityRequest {
247        id: Some(EntityId {
248            prefix: "container_id".to_string(),
249            uid: "nonexistent".to_string(),
250        }),
251        cardinality: TagCardinality::High.into(),
252    };
253    match client.tagger_fetch_entity(noop_fetch_request).await {
254        Ok(_) => Ok(()),
255        Err(e) => match e.code() {
256            Code::Unauthenticated => Err(generic_error!(
257                "Failed to authenticate to Datadog Agent API. Check that the configured authentication token is correct."
258            )),
259            _ => Err(e.into()),
260        },
261    }
262}
datadog_agent_commons/ipc/client/mod.rs

datadog_agent_commons/ipc/client/
mod.rs
