1use crate::llm::actions::{
2 client_trait::{Client, ClientActionResult},
3 protocol_trait::Protocol,
4 ActionDefinition, ParameterDefinition,
5};
6use crate::protocol::{ConnectContext, EventType};
7use crate::server::nbd::actions::{action, parameter};
8use anyhow::{bail, ensure, Result};
9use serde_json::{json, Value};
10use std::sync::LazyLock;
11
12#[derive(Default)]
13pub struct NbdClientProtocol;
14impl NbdClientProtocol {
15 pub fn new() -> Self {
16 Self
17 }
18}
19
20fn read() -> ActionDefinition {
21 action(
22 "nbd_read",
23 "Read a range of the export; nbd_read_result reports the bytes (text when UTF-8, else hex; the first 4096 shown), their SHA-256, whether they are all zero, or the server's error",
24 vec![
25 parameter("offset", "number", "Byte offset into the export", true),
26 parameter("length", "number", "Bytes to read (1 to 1048576)", true),
27 ],
28 json!({"type": "nbd_read", "offset": 0, "length": 512}),
29 )
30}
31fn block_status() -> ActionDefinition {
32 action(
33 "nbd_block_status",
34 "Ask which parts of a range hold data and which are holes reading as zeroes (needs the server's base:allocation context)",
35 vec![
36 parameter("offset", "number", "Byte offset into the export", true),
37 parameter("length", "number", "Bytes to describe (1 to 4294967295)", true),
38 ],
39 json!({"type": "nbd_block_status", "offset": 0, "length": 1048576}),
40 )
41}
42fn flush() -> ActionDefinition {
43 action(
44 "nbd_flush",
45 "Send NBD_CMD_FLUSH and report the server's answer",
46 vec![],
47 json!({"type": "nbd_flush"}),
48 )
49}
50fn disconnect() -> ActionDefinition {
51 action(
52 "disconnect",
53 "Send NBD_CMD_DISC and close the connection",
54 vec![],
55 json!({"type": "disconnect"}),
56 )
57}
58
59pub fn all_actions() -> Vec<ActionDefinition> {
60 vec![read(), block_status(), flush(), disconnect()]
61}
62
63fn event(id: &str, description: &str, params: Vec<crate::llm::actions::Parameter>) -> EventType {
64 EventType::new(id, description, read().example.clone())
65 .with_parameters(params)
66 .with_actions(all_actions())
67}
68
69pub static CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
70 event(
71 "nbd_connected",
72 "Negotiation finished and the export is open",
73 vec![
74 parameter("export", "string", "The export name", true),
75 parameter("size", "number", "Export size in bytes", true),
76 parameter(
77 "read_only",
78 "boolean",
79 "Whether the server marked the export read-only",
80 true,
81 ),
82 parameter(
83 "description",
84 "string",
85 "The server's description of the export",
86 false,
87 ),
88 parameter(
89 "block_size",
90 "object",
91 "{minimum, preferred, maximum} when the server sent them",
92 false,
93 ),
94 parameter(
95 "structured_replies",
96 "boolean",
97 "Whether structured replies were negotiated",
98 true,
99 ),
100 parameter(
101 "base_allocation",
102 "boolean",
103 "Whether block status is available",
104 true,
105 ),
106 parameter(
107 "exports",
108 "array",
109 "The server's export list, when list_exports was set",
110 false,
111 ),
112 ],
113 )
114});
115pub static READ_EVENT: LazyLock<EventType> = LazyLock::new(|| {
116 event(
117 "nbd_read_result",
118 "The server answered a read",
119 vec![
120 parameter("offset", "number", "Where the read started", true),
121 parameter("length", "number", "How many bytes were asked for", true),
122 parameter(
123 "error",
124 "string",
125 "The server's error (EIO, EINVAL, ...) when the read failed",
126 false,
127 ),
128 parameter(
129 "error_offset",
130 "number",
131 "The first failing byte, when the server named it",
132 false,
133 ),
134 parameter(
135 "data",
136 "string",
137 "The first 4096 bytes: text when UTF-8, else hex per data_encoding",
138 false,
139 ),
140 parameter(
141 "data_encoding",
142 "string",
143 "utf8 when data is text, hex when it is not UTF-8",
144 false,
145 ),
146 parameter("sha256", "string", "SHA-256 of all the bytes read", false),
147 parameter(
148 "all_zero",
149 "boolean",
150 "Whether every byte read was zero",
151 false,
152 ),
153 ],
154 )
155});
156pub static STATUS_EVENT: LazyLock<EventType> = LazyLock::new(|| {
157 event(
158 "nbd_block_status_result",
159 "The server described a range's allocation, or refused",
160 vec![
161 parameter("offset", "number", "Where the range started", true),
162 parameter("extents", "array", "[{offset, length, hole, zero}]", false),
163 parameter(
164 "error",
165 "string",
166 "The server's error when it refused",
167 false,
168 ),
169 ],
170 )
171});
172pub static FLUSH_EVENT: LazyLock<EventType> = LazyLock::new(|| {
173 event(
174 "nbd_flush_result",
175 "The server answered a flush",
176 vec![parameter(
177 "error",
178 "string",
179 "OK, or the server's error",
180 true,
181 )],
182 )
183});
184
185impl Protocol for NbdClientProtocol {
186 fn protocol_name(&self) -> &'static str {
187 "NBD"
188 }
189 fn stack_name(&self) -> &'static str {
190 "ETH>IP>TCP>NBD"
191 }
192 fn keywords(&self) -> Vec<&'static str> {
193 vec!["nbd", "nbd client", "network block device"]
194 }
195 fn description(&self) -> &'static str {
196 "Network Block Device client: negotiates an export and reads it, with block status"
197 }
198 fn get_async_actions(&self, _: &crate::state::app_state::AppState) -> Vec<ActionDefinition> {
199 all_actions()
200 }
201 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
202 vec![]
203 }
204 fn get_event_types(&self) -> Vec<EventType> {
205 vec![
206 CONNECTED_EVENT.clone(),
207 READ_EVENT.clone(),
208 STATUS_EVENT.clone(),
209 FLUSH_EVENT.clone(),
210 ]
211 }
212 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
213 vec![
214 ParameterDefinition {
215 name: "export".into(),
216 type_hint: "string".into(),
217 description: "Export name to open (empty for the server's default)".into(),
218 required: false,
219 example: json!("disk0"),
220 default: Some(json!(super::DEFAULT_EXPORT)),
221 },
222 ParameterDefinition {
223 name: "list_exports".into(),
224 type_hint: "boolean".into(),
225 description: "Ask for the export list (NBD_OPT_LIST) before opening".into(),
226 required: false,
227 example: json!(true),
228 default: Some(json!(super::DEFAULT_LIST_EXPORTS)),
229 },
230 ]
231 }
232 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
233 use crate::protocol::metadata::*;
234 ProtocolMetadataV2::builder()
235 .state(DevelopmentState::Experimental)
236 .privilege_requirement(PrivilegeRequirement::None)
237 .implementation("The server's NBD wire constants; fixed newstyle negotiation (LIST, STRUCTURED_REPLY, SET_META_CONTEXT base:allocation, GO) and one request at a time")
238 .llm_control("Which ranges to read and describe, and when to disconnect")
239 .e2e_testing("tests/client/nbd: nbdkit 1.48.1 (independent, C) serving its data plugin, with its error filter injecting EIO")
240 .notes("Read-only client: no write, trim or zero commands. No TLS or extended headers. Reads up to 1 MiB.")
241 .build()
242 }
243 fn example_prompt(&self) -> &'static str {
244 "Open export disk0 on the NBD server at 127.0.0.1:10809 and read its first sector"
245 }
246 fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
247 let llm = json!({"type":"open_client","protocol":"nbd","remote_addr":"127.0.0.1:10809","instruction":"Read the first 512 bytes of disk0","startup_params":{"export":"disk0"}});
248 let mut static_example = llm.clone();
249 static_example["event_handlers"] = json!([{"event_pattern":"nbd_connected","handler":{"type":"static","actions":[read().example]}}]);
250 let mut scripted = llm.clone();
251 scripted["event_handlers"] = json!([{"event_pattern":"nbd_connected","handler":{"type":"script","language":"python","code":"import json,sys\ne=json.load(sys.stdin)['event']\nprint(json.dumps({'actions':[{'type':'nbd_read','offset':0,'length':min(512,e['size'])}]}))"}}]);
252 crate::llm::actions::StartupExamples::new(llm, scripted, static_example)
253 }
254 fn group_name(&self) -> &'static str {
255 "Web & File"
256 }
257}
258
259impl Client for NbdClientProtocol {
260 fn connect(
261 &self,
262 ctx: ConnectContext,
263 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<std::net::SocketAddr>> + Send>>
264 {
265 Box::pin(super::connect(ctx))
266 }
267 fn execute_action(&self, v: Value) -> Result<ClientActionResult> {
268 match v["type"].as_str() {
269 Some("disconnect") => return Ok(ClientActionResult::Disconnect),
270 Some("nbd_read") => ensure!(
271 v["offset"].as_u64().is_some()
272 && matches!(v["length"].as_u64(), Some(1..=1_048_576)),
273 "offset is a byte offset and length 1 to 1048576"
274 ),
275 Some("nbd_block_status") => ensure!(
276 v["offset"].as_u64().is_some()
277 && matches!(v["length"].as_u64(), Some(1..=4_294_967_295)),
278 "offset is a byte offset and length 1 to 4294967295"
279 ),
280 Some("nbd_flush") => {}
281 _ => bail!("Unknown NBD client action"),
282 }
283 Ok(ClientActionResult::Custom {
284 name: v["type"].as_str().unwrap_or_default().into(),
285 data: v,
286 })
287 }
288}