1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
use std::sync::Arc;
use std::time::Duration;
use affinidi_tdk::common::TDKSharedState;
use affinidi_tdk::common::config::TDKConfig;
use affinidi_tdk::didcomm::Message;
use affinidi_tdk::messaging::ATM;
use affinidi_tdk::messaging::config::ATMConfig;
use affinidi_tdk::messaging::profiles::ATMProfile;
use affinidi_tdk::secrets_resolver::SecretsResolver;
use tracing::{debug, info, warn};
use crate::error::VtaError;
use crate::protocols::PROBLEM_REPORT_TYPE;
// Per-step ceiling for mediator round-trips during session setup. The
// upstream calls below are otherwise unbounded — when the mediator is
// unreachable a CLI invocation can hang for the full TCP/TLS retry
// window (30–60s on macOS) before failing. 15s is generous for healthy
// mediators and keeps Ctrl-C responsive.
const MEDIATOR_OP_TIMEOUT: Duration = Duration::from_secs(15);
/// Client-side DIDComm session for request-response messaging via ATM.
///
/// Uses WebSocket streaming to receive responses from the mediator.
/// Designed for CLI tools that send a request and wait for a reply.
#[derive(Clone)]
pub struct DIDCommSession {
atm: Arc<ATM>,
profile: Arc<ATMProfile>,
pub(crate) client_did: String,
pub(crate) vta_did: String,
}
impl DIDCommSession {
/// The session's local DID — the one used as the authcrypt
/// sender on outbound messages. Surfaced so SDK helpers can
/// pre-check sender == expected DID before sending (the VTA's
/// `provision-integration` handler enforces sender == VP
/// holder, for instance).
pub fn client_did(&self) -> &str {
&self.client_did
}
/// Connect to a VTA via DIDComm through a mediator.
///
/// Sets up the ATM and profile for REST-based messaging. Does NOT open a
/// WebSocket — all communication goes through the mediator's REST API,
/// avoiding connection storms when the CLI is invoked repeatedly.
pub async fn connect(
client_did: &str,
private_key_multibase: &str,
vta_did: &str,
mediator_did: &str,
) -> Result<Self, Box<dyn std::error::Error>> {
// Decode private key and build DIDComm secrets
let seed = crate::did_key::decode_private_key_multibase(private_key_multibase)?;
let secrets = crate::did_key::secrets_from_did_key(client_did, &seed)?;
// Create TDK shared state and insert secrets
let tdk = TDKSharedState::new(TDKConfig::builder().build()?).await?;
tdk.secrets_resolver().insert(secrets.signing).await;
tdk.secrets_resolver().insert(secrets.key_agreement).await;
// Build ATM (no inbound channel needed — we use REST polling)
let atm_config = ATMConfig::builder().build()?;
let atm = ATM::new(atm_config, Arc::new(tdk)).await?;
// Create profile with mediator
let profile = ATMProfile::new(
&atm,
None,
client_did.to_string(),
Some(mediator_did.to_string()),
)
.await?;
let profile = Arc::new(profile);
let atm = Arc::new(atm);
// Flush stale messages from the inbox (accumulated between CLI runs)
{
use affinidi_tdk::messaging::messages::Folder;
match tokio::time::timeout(
MEDIATOR_OP_TIMEOUT,
atm.list_messages(&profile, Folder::Inbox),
)
.await
{
Ok(Ok(messages)) if !messages.is_empty() => {
let ids: Vec<String> = messages.iter().map(|m| m.msg_id.clone()).collect();
info!(
count = ids.len(),
"flushing stale queued messages from inbox"
);
let delete_req = affinidi_tdk::messaging::messages::DeleteMessageRequest {
message_ids: ids,
};
match tokio::time::timeout(
MEDIATOR_OP_TIMEOUT,
atm.delete_messages_direct(&profile, &delete_req),
)
.await
{
Ok(Ok(resp)) => {
debug!(
deleted = resp.success.len(),
errors = resp.errors.len(),
"inbox flushed"
);
}
Ok(Err(e)) => warn!("failed to flush stale messages (non-fatal): {e}"),
Err(_) => warn!(
"timeout flushing stale messages after {}s (non-fatal)",
MEDIATOR_OP_TIMEOUT.as_secs()
),
}
}
Ok(Ok(_)) => {} // Empty inbox
Ok(Err(e)) => warn!("could not list inbox (non-fatal): {e}"),
Err(_) => warn!(
"timeout listing inbox after {}s (non-fatal)",
MEDIATOR_OP_TIMEOUT.as_secs()
),
}
}
// Enable WebSocket for streaming message delivery from mediator.
// Without this, the ATM can only poll via REST and may miss responses
// that arrive after the initial send_message call returns.
match tokio::time::timeout(MEDIATOR_OP_TIMEOUT, atm.profile_enable_websocket(&profile))
.await
{
Ok(res) => res?,
Err(_) => {
return Err(format!(
"timeout enabling WebSocket to mediator after {}s — \
mediator may be unreachable",
MEDIATOR_OP_TIMEOUT.as_secs()
)
.into());
}
}
debug!("DIDComm session connected via mediator {mediator_did} (WebSocket mode)");
Ok(Self {
atm,
profile,
client_did: client_did.to_string(),
vta_did: vta_did.to_string(),
})
}
/// Send a DIDComm message and wait for a matching response.
///
/// Packs the message, sends it to the mediator, then uses the WebSocket
/// live stream to wait for the response. This handles asynchronous
/// processing where the VTA takes time to respond.
///
/// Problem-report responses are decoded into typed [`VtaError`] variants
/// based on their `e.p.msg.*` code so DIDComm and REST surface the same
/// error taxonomy (conflict, not-found, auth, validation, server).
pub async fn send_and_wait<T: serde::de::DeserializeOwned>(
&self,
msg_type: &str,
body: serde_json::Value,
expected_result_type: &str,
timeout_secs: u64,
) -> Result<T, VtaError> {
let msg_id = uuid::Uuid::new_v4().to_string();
let msg = Message::build(msg_id.clone(), msg_type.to_string(), body)
.from(self.client_did.clone())
.to(self.vta_did.clone())
.finalize();
// Pack encrypted (signed + encrypted to recipient)
let (packed, _) = self
.atm
.pack_encrypted(
&msg,
&self.vta_did,
Some(&self.client_did),
Some(&self.client_did),
)
.await
.map_err(|e| VtaError::DidcommTransport(format!("failed to pack message: {e}")))?;
debug!(msg_type, msg_id, "sending via DIDComm");
// Wrap the inner JWE in a `routing/2.0/forward` envelope addressed
// to the mediator and ship it. Strict mediators
// (`local_direct_delivery_allowed: false`) refuse direct delivery
// — same path the production VTA-side `affinidi-messaging-didcomm-service`
// takes via `forward_and_send_message`.
let mediator_did = self
.profile
.inner
.mediator
.as_ref()
.as_ref()
.map(|m| m.did.clone())
.ok_or_else(|| {
VtaError::DidcommTransport("no mediator configured on profile".into())
})?;
self.atm
.forward_and_send_message(
&self.profile,
false, // authcrypt the forward envelope (mediator policy)
&packed,
Some(&msg_id),
&mediator_did,
&self.vta_did,
None,
None,
false,
)
.await
.map_err(|e| VtaError::DidcommTransport(format!("failed to send message: {e}")))?;
// Wait for the response via WebSocket live stream
let timeout = std::time::Duration::from_secs(timeout_secs);
let wait_duration = std::time::Duration::from_secs(5);
let deadline = tokio::time::Instant::now() + timeout;
loop {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return Err(VtaError::DidcommTransport(
"timeout waiting for DIDComm response".into(),
));
}
let wait = wait_duration.min(remaining);
let next = self
.atm
.message_pickup()
.live_stream_next(&self.profile, Some(wait), true)
.await
.map_err(|e| VtaError::DidcommTransport(format!("message pickup error: {e}")))?;
let (response_msg, _meta) = match next {
Some(pair) => pair,
None => continue, // No message yet, keep waiting
};
// Check if this is the response we're waiting for (matching thread ID)
let response_thid = response_msg.thid.as_deref().unwrap_or("");
if response_thid != msg_id {
debug!(
response_thid,
expected = msg_id,
response_type = %response_msg.typ,
"received message with non-matching thread ID — skipping"
);
continue;
}
debug!(response_type = %response_msg.typ, "received DIDComm response");
// Check for problem report — map the `e.p.msg.*` code to the
// matching VtaError variant so callers can `match` on the same
// error shapes they get from REST (see `VtaError::from_http`).
if response_msg.typ == PROBLEM_REPORT_TYPE
|| response_msg.typ.contains("problem-report")
{
let code = response_msg
.body
.get("code")
.and_then(|v| v.as_str())
.unwrap_or("unknown");
let comment = response_msg
.body
.get("comment")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
return Err(VtaError::from_problem_report(code, comment));
}
// Verify expected type
if response_msg.typ != expected_result_type {
return Err(VtaError::Protocol(format!(
"unexpected response type: expected {expected_result_type}, got {}",
response_msg.typ
)));
}
// Deserialize response body
return serde_json::from_value(response_msg.body).map_err(VtaError::from);
}
}
/// Gracefully shut down the DIDComm session.
pub async fn shutdown(&self) {
self.atm.graceful_shutdown().await;
}
}