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

datadog_agent_commons/ipc/client/
mod.rs
