datadog_agent_commons/ipc/client/
mod.rs

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

datadog_agent_commons/ipc/client/
mod.rs
