Skip to main content

netget/server/libp2p/
mod.rs

1//! libp2p host (server role): accepts TCP connections, upgrades them to Noise and yamux,
2//! answers identify and ping in Rust, and hands every message on the application protocols'
3//! streams to the model, one connection's events at a time.
4pub mod actions;
5pub mod bitswap;
6pub mod host;
7pub mod ipfs;
8pub mod kad;
9pub mod noise;
10mod serve_ipfs;
11pub mod wire;
12pub mod yamux;
13
14use crate::llm::action_helper::call_llm;
15use crate::llm::actions::protocol_trait::ActionResult;
16use crate::logging::emit::Log;
17use crate::protocol::{Event, SpawnContext};
18use crate::server::{
19    accept_bounded::{accept_bounded, ConnectionLimiter, DEFAULT_MAX_CONNECTIONS},
20    connection::ConnectionId,
21};
22use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
23use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
24use anyhow::{Context, Result};
25use host::{Config, Note, Peer, Spawn, Task};
26use serde_json::{json, Value};
27use std::net::SocketAddr;
28use std::sync::Arc;
29use tokio::net::TcpStream;
30use tokio::sync::mpsc;
31
32/// Notes (messages, closes) queued for one connection's model before its streams wait.
33pub const NOTE_QUEUE: usize = 64;
34
35/// The identity a `private_key_seed` names: SHA-256 of the text is the Ed25519 seed, so the
36/// same text always gives the same peer id. None: a new key.
37pub fn identity_from_params(seed: Option<String>) -> Result<noise::Identity> {
38    match seed {
39        Some(text) => {
40            anyhow::ensure!(!text.is_empty(), "private_key_seed must not be empty");
41            use sha2::Digest;
42            Ok(noise::Identity::from_seed(
43                sha2::Sha256::digest(text.as_bytes()).into(),
44            ))
45        }
46        None => Ok(noise::Identity::random()),
47    }
48}
49
50pub fn protocols_from_params(list: Option<&Vec<Value>>) -> Result<Vec<String>> {
51    let Some(list) = list else {
52        return Ok(host::DEFAULT_PROTOCOLS
53            .iter()
54            .map(|s| s.to_string())
55            .collect());
56    };
57    list.iter()
58        .map(|v| {
59            let p = v.as_str().context("protocols must be strings")?;
60            anyhow::ensure!(
61                p.starts_with('/') && p.len() <= 256 && !p.contains('\n'),
62                "{p:?} is not a protocol id such as /netget/chat/1.0.0"
63            );
64            anyhow::ensure!(
65                ![wire::IDENTIFY, wire::IDENTIFY_PUSH, wire::PING].contains(&p),
66                "{p} is answered by NetGet itself"
67            );
68            anyhow::ensure!(
69                !bitswap::is_bitswap(p) && !kad::is_kad(p),
70                "{p} is IPFS: set the ipfs startup parameter instead"
71            );
72            Ok(p.to_string())
73        })
74        .collect()
75}
76
77pub async fn spawn(ctx: SpawnContext) -> Result<SocketAddr> {
78    let params = ctx.startup_params.as_ref();
79    let seed = params
80        .map(|p| p.get_optional_string("private_key_seed"))
81        .transpose()?
82        .flatten();
83    let protocols = protocols_from_params(
84        params
85            .map(|p| p.get_optional_array("protocols"))
86            .transpose()?
87            .flatten(),
88    )?;
89    let identity = identity_from_params(seed)?;
90    let ipfs = params
91        .map(|p| p.get_optional_bool("ipfs"))
92        .transpose()?
93        .flatten()
94        .unwrap_or(actions::IPFS_DEFAULT);
95    let listener = tokio::net::TcpListener::bind(ctx.legacy_listen_addr()).await?;
96    let local = listener.local_addr()?;
97    let cfg = Arc::new(Config {
98        listen_addrs: vec![wire::multiaddr_bytes(&local)],
99        identity,
100        protocols,
101        bitswap: ipfs,
102        kad: ipfs,
103        directory: Default::default(),
104    });
105    Log::new(Some(&ctx.status_tx)).info(format!(
106        "libp2p host listening on {}/p2p/{} ({})",
107        wire::multiaddr_string(&local),
108        cfg.identity.peer_id_string(),
109        cfg.served_protocols().join(", ")
110    ));
111    let server_id = ctx.server_id;
112    let state = ctx.state.clone();
113    let accept =
114        tokio::spawn(async move {
115            let limiter = ConnectionLimiter::new(DEFAULT_MAX_CONNECTIONS);
116            loop {
117                let (stream, peer, permit) =
118                    match accept_bounded(&listener, &limiter, b"", "libp2p", Some(&ctx.status_tx))
119                        .await
120                    {
121                        Ok(v) => v,
122                        Err(_) => break,
123                    };
124                let child = ctx.clone();
125                let cfg = cfg.clone();
126                ctx.state
127                    .spawn_server_task(server_id, async move {
128                        let _permit = permit;
129                        connection(child, cfg, stream, peer, local).await;
130                    })
131                    .await;
132            }
133        });
134    state.register_server_task(server_id, accept).await;
135    Ok(local)
136}
137
138fn spawner(ctx: &SpawnContext) -> Spawn {
139    let state = ctx.state.clone();
140    let server_id = ctx.server_id;
141    Arc::new(move |task: Task| -> Task {
142        let state = state.clone();
143        Box::pin(async move {
144            state.spawn_server_task(server_id, task).await;
145        })
146    })
147}
148
149async fn connection(
150    ctx: SpawnContext,
151    cfg: Arc<Config>,
152    tcp: TcpStream,
153    remote: SocketAddr,
154    local: SocketAddr,
155) {
156    let log = Log::new(Some(&ctx.status_tx));
157    let conn = match host::upgrade_inbound(tcp, &cfg).await {
158        Ok(c) => c,
159        Err(e) => {
160            log.warn(format!("libp2p connection from {remote} refused: {e:#}"));
161            return;
162        }
163    };
164    let remote_peer = wire::peer_id_string(&conn.remote_peer);
165    let id = ConnectionId::new(ctx.state.get_next_unified_id().await);
166    let now = crate::utils::clock::Instant::now();
167    ctx.state
168        .add_connection_to_server(
169            ctx.server_id,
170            ConnectionState {
171                id,
172                remote_addr: remote,
173                local_addr: local,
174                bytes_sent: 0,
175                bytes_received: 0,
176                packets_sent: 0,
177                packets_received: 0,
178                last_activity: now,
179                status: ConnectionStatus::Active,
180                status_changed_at: now,
181                protocol_info: ProtocolConnectionInfo::new(json!({"peer_id": remote_peer})),
182            },
183        )
184        .await;
185    let close = |reason: String| {
186        let ctx = ctx.clone();
187        async move {
188            ctx.state
189                .update_connection_status(ctx.server_id, id, ConnectionStatus::Closed)
190                .await;
191            ctx.state
192                .remove_peer_handle(ctx.server_id, id.as_u32())
193                .await;
194            Log::new(Some(&ctx.status_tx)).info(format!("libp2p connection {id} closed: {reason}"));
195            let _ = ctx.status_tx.send("__UPDATE_UI__".into());
196        }
197    };
198    let session = yamux::session(conn, false);
199    let (notes_tx, mut notes) = mpsc::channel(NOTE_QUEUE);
200    let peer = Arc::new(Peer {
201        opener: session.opener.clone(),
202        remote_peer: remote_peer.clone(),
203        remote_addr: remote,
204        streams: Default::default(),
205        notes: notes_tx.clone(),
206    });
207    let spawn = spawner(&ctx);
208    let reader = session.reader;
209    let ended = notes_tx.clone();
210    spawn(Box::pin(async move {
211        let r = reader.await;
212        let _ = ended
213            .send(Note::Ended(r.err().map(|e| format!("{e:#}"))))
214            .await;
215    }))
216    .await;
217    let writer = session.writer;
218    spawn(Box::pin(async move {
219        let _ = writer.await;
220    }))
221    .await;
222    spawn(Box::pin(host::accept_streams(
223        session.incoming,
224        cfg.clone(),
225        peer.clone(),
226        spawn.clone(),
227    )))
228    .await;
229
230    let remote_info = host::identify(&peer).await;
231    let (agent, protocols, listen) = match &remote_info {
232        Ok(r) => (
233            r.agent_version.clone(),
234            r.protocols.clone(),
235            r.listen_addrs.clone(),
236        ),
237        Err(e) => {
238            log.warn(format!("libp2p connection {id}: identify failed: {e:#}"));
239            (String::new(), Vec::new(), Vec::new())
240        }
241    };
242    log.info(format!(
243        "libp2p connection {id}: peer {remote_peer} ({agent}) connected"
244    ));
245    if cfg.kad {
246        let addrs = listen
247            .iter()
248            .filter_map(|a| wire::multiaddr_from_text(a).ok())
249            .take(kad::MAX_ADDRS)
250            .collect();
251        cfg.directory
252            .lock()
253            .unwrap_or_else(|e| e.into_inner())
254            .insert(remote_peer.clone(), addrs);
255    }
256    let mut peer_rx =
257        crate::server::peer_support::register_peer_channel(&ctx.state, ctx.server_id, id.as_u32())
258            .await;
259    let connected = Event::new(
260        &actions::CONNECTED_EVENT,
261        json!({"peer_id": remote_peer, "remote_addr": wire::multiaddr_string(&remote),
262               "agent_version": agent, "protocols": protocols, "listen_addrs": listen}),
263    );
264    answer(&ctx, id, &peer, &spawn, connected, None).await;
265
266    let reason = loop {
267        tokio::select! {
268            note = notes.recv() => match note {
269                Some(Note::Message { stream_id, protocol, data }) => {
270                    ctx.state
271                        .update_connection_stats(ctx.server_id, id, Some(data.len() as u64), None, Some(1), None)
272                        .await;
273                    if bitswap::is_bitswap(&protocol) {
274                        serve_ipfs::bitswap(&ctx, id, &peer, stream_id, &data).await;
275                        continue;
276                    }
277                    if kad::is_kad(&protocol) {
278                        serve_ipfs::kad(&ctx, id, &peer, &cfg, stream_id, &protocol, &data).await;
279                        continue;
280                    }
281                    let (text, encoding) = wire::shown(&data);
282                    let event = Event::new(
283                        &actions::MESSAGE_EVENT,
284                        json!({"peer_id": remote_peer, "stream_id": stream_id, "protocol": protocol,
285                               "data": text, "encoding": encoding}),
286                    );
287                    answer(&ctx, id, &peer, &spawn, event, Some(stream_id)).await;
288                }
289                Some(Note::Closed { stream_id, error, .. }) => {
290                    // Every message on it has been answered (notes are in order): close ours.
291                    if let Some((_, w)) = peer.writer(stream_id) {
292                        if let Some(e) = error {
293                            log.warn(format!("libp2p connection {id}: stream {stream_id} reset: {e}"));
294                        } else {
295                            w.close();
296                        }
297                    }
298                    peer.forget(stream_id);
299                }
300                Some(Note::Ended(e)) => break e.unwrap_or_else(|| "peer went away".into()),
301                None => break "connection dropped".into(),
302            },
303            command = peer_rx.recv() => {
304                let Some(command) = command else { continue };
305                inject(&ctx, id, &peer, &spawn, command).await;
306            }
307        }
308    };
309    peer.opener.go_away();
310    if cfg.kad {
311        cfg.directory
312            .lock()
313            .unwrap_or_else(|e| e.into_inner())
314            .remove(&remote_peer);
315    }
316    close(reason).await;
317}
318
319/// Ask the model about one event and carry out its answer. `stream` is the stream the
320/// event arrived on, which `libp2p_send` and `libp2p_close_stream` default to.
321async fn answer(
322    ctx: &SpawnContext,
323    id: ConnectionId,
324    peer: &Arc<Peer>,
325    spawn: &Spawn,
326    event: Event,
327    stream: Option<u32>,
328) {
329    let log = Log::new(Some(&ctx.status_tx));
330    let Some(todo) = ask(ctx, id, &event).await else {
331        // Fail closed: the stream being answered is reset, so the peer is not left waiting.
332        if let Some((_, w)) = stream.and_then(|s| peer.writer(s)) {
333            w.reset();
334        }
335        return;
336    };
337    for action in todo {
338        if let Err(e) = apply(ctx, id, peer, spawn, &action, stream).await {
339            log.warn(format!(
340                "libp2p connection {id}: {} failed: {e:#}",
341                action["type"]
342            ));
343        }
344    }
345}
346
347/// Ask the model about one event: the actions it answered with, in order, or None when it
348/// failed (logged `decision=fail_closed_llm_error`; the caller fails closed).
349async fn ask(ctx: &SpawnContext, id: ConnectionId, event: &Event) -> Option<Vec<Value>> {
350    let log = Log::new(Some(&ctx.status_tx));
351    let op = event.id().to_string();
352    let result = match call_llm(
353        &ctx.llm_client,
354        &ctx.state,
355        ctx.server_id,
356        Some(id),
357        event,
358        &actions::Libp2pProtocol,
359    )
360    .await
361    {
362        Ok(r) if r.failures.is_empty() => r,
363        Ok(_) | Err(_) => {
364            log.error(format!(
365                "libp2p connection {id} operation={op} decision=fail_closed_llm_error"
366            ));
367            return None;
368        }
369    };
370    let mut todo = Vec::new();
371    let mut stack = result.protocol_results;
372    while let Some(r) = stack.pop() {
373        match r {
374            ActionResult::Custom { data, .. } => todo.push(data),
375            ActionResult::Multiple(items) => stack.extend(items),
376            _ => {}
377        }
378    }
379    todo.reverse();
380    let decision = if todo.is_empty() {
381        "model_silent"
382    } else {
383        "model_answer"
384    };
385    log.info(format!(
386        "libp2p connection {id} operation={op} decision={decision}"
387    ));
388    Some(todo)
389}
390
391/// Carry out one validated action on this connection; a description of what happened.
392async fn apply(
393    ctx: &SpawnContext,
394    id: ConnectionId,
395    peer: &Arc<Peer>,
396    spawn: &Spawn,
397    action: &Value,
398    stream: Option<u32>,
399) -> Result<String> {
400    let target = || -> Result<u32> {
401        action["stream_id"]
402            .as_u64()
403            .map(|s| s as u32)
404            .or(stream)
405            .context("no stream to use: give stream_id, or open one with libp2p_open_stream")
406    };
407    match action["type"].as_str().unwrap_or_default() {
408        actions::SEND => {
409            let data = wire::data_bytes(action)?;
410            let s = target()?;
411            host::send(peer, s, &data)?;
412            ctx.state
413                .update_connection_stats(
414                    ctx.server_id,
415                    id,
416                    None,
417                    Some(data.len() as u64),
418                    None,
419                    Some(1),
420                )
421                .await;
422            Ok(format!("sent {} bytes on stream {s}", data.len()))
423        }
424        actions::OPEN => {
425            let protocol = action["protocol"].as_str().unwrap_or_default();
426            match host::open_app_stream(peer, protocol).await? {
427                None => anyhow::bail!("the peer does not support {protocol}"),
428                Some((s, reader)) => {
429                    spawn(reader).await;
430                    if !action["data"].is_null() {
431                        host::send(peer, s, &wire::data_bytes(action)?)?;
432                    }
433                    Ok(format!("opened stream {s} on {protocol}"))
434                }
435            }
436        }
437        actions::CLOSE => {
438            let s = target()?;
439            let (_, w) = peer
440                .writer(s)
441                .with_context(|| format!("no open stream {s}"))?;
442            w.close();
443            Ok(format!("closed stream {s}"))
444        }
445        actions::DISCONNECT => {
446            peer.opener.go_away();
447            Ok("disconnecting".into())
448        }
449        t @ (actions::BLOCKS | actions::KAD_REPLY) => {
450            anyhow::bail!("{t} answers an IPFS request, and can only be given as that answer")
451        }
452        other => anyhow::bail!("unknown action {other}"),
453    }
454}
455
456/// An action injected from the dashboard or MCP for this peer.
457async fn inject(
458    ctx: &SpawnContext,
459    id: ConnectionId,
460    peer: &Arc<Peer>,
461    spawn: &Spawn,
462    command: ClientCommand,
463) {
464    let action = command.action.clone();
465    let outcome = match actions::check(&action) {
466        Err(e) => ClientSendOutcome::Rejected {
467            error: e.to_string(),
468        },
469        Ok(()) => match apply(ctx, id, peer, spawn, &action, None).await {
470            Ok(_) if action["type"] == actions::DISCONNECT => ClientSendOutcome::Disconnected,
471            Ok(detail) => ClientSendOutcome::Executed { detail },
472            Err(e) => ClientSendOutcome::Rejected {
473                error: format!("{e:#}"),
474            },
475        },
476    };
477    let _ = command.reply_tx.send(Ok(outcome));
478}