lewis.moe
Nothing to see here, move along
1#![no_std]
2#![no_main]
3
4mod rng;
5mod ssh;
6
7use lancer_core::ipc::{MAX_IPC_BYTES, unpack_bytes};
8use lancer_core::packet_ring::{PacketRingReader, PacketRingWriter};
9use lancer_user::path;
10use lancer_user::syscall;
11use ssh::{SshProgress, SshSession};
12
13const SHELL_RING_BASE_SLOT: u64 = 64;
14const SHELL_RING_FRAME_COUNT: u64 = 4;
15const OWN_NOTIF_SLOT: u64 = 3;
16const NETSTACK_NOTIF_SLOT: u64 = 4;
17
18const MAX_SESSIONS: usize = 4;
19
20const TAG_ENDPOINT: u64 = 1;
21const TAG_NOTIFICATION: u64 = 2;
22const TAG_SCHED_CONTEXT: u64 = 5;
23const TAG_FRAME: u64 = 11;
24const RIGHTS_ALL: u64 = 0b1111;
25
26const SLOT_UNTYPED: u64 = 28;
27
28const VFS_RING_BASE: u64 = 128;
29const VFS_RING_FRAMES: u64 = 16;
30const VFS_CLIENT_NOTIF_SLOT: u64 = 12;
31const VFS_NOTIF_SLOT: u64 = 13;
32const VFS_PROC_SLOT: u64 = 15;
33
34const BOOT_MODULE_COUNT: u64 = 7;
35
36const RING_BASE: u64 = 0x4000_0000;
37const RING_HALF: usize = 8192;
38
39const MSG_DATA: u8 = 0;
40const MSG_CONNECT: u8 = 1;
41const MSG_DISCONNECT: u8 = 2;
42const MSG_DNS_RESULT: u8 = 4;
43const MSG_PING_RESULT: u8 = 6;
44const MSG_PING_DONE: u8 = 7;
45const MSG_UDP_RECV: u8 = 10;
46const MSG_UDP_CLOSE: u8 = 11;
47const MSG_UDP_BOUND: u8 = 12;
48const MSG_DNS_RESULT6: u8 = 14;
49const MSG_IFCONFIG_RESULT: u8 = 17;
50const MSG_ARP_ENTRY: u8 = 19;
51const MSG_ARP_DONE: u8 = 20;
52const MSG_NETSTAT_ENTRY: u8 = 22;
53const MSG_NETSTAT_ENTRY_UDP: u8 = 23;
54const MSG_NETSTAT_DONE: u8 = 24;
55
56const IDLE_TIMEOUT_MS: i64 = 300_000;
57const NOTIFY_BIT_SHELL_DEATH_BASE: u64 = 0x10;
58const NOTIFY_BIT_SHELL_NETSOCK: u64 = 0x04;
59
60const INPUT_RING_PAGES: u64 = 1;
61const INPUT_RING_SIZE: usize = 4096;
62const INPUT_RING_SLOT_SIZE: u32 = 64;
63
64const NETSOCK_RING_PAGES: u64 = 2;
65const NETSOCK_RING_HALF: usize = 4096;
66const NETSOCK_RING_SLOT_SIZE: u32 = 128;
67
68static mut SSH_INBUFS: [[u8; 35000]; MAX_SESSIONS] = [[0u8; 35000]; MAX_SESSIONS];
69static mut SSH_OUTBUFS: [[u8; 35000]; MAX_SESSIONS] = [[0u8; 35000]; MAX_SESSIONS];
70
71fn slot_proc(i: usize) -> u64 {
72 30 + i as u64 * 7
73}
74fn slot_sched(i: usize) -> u64 {
75 31 + i as u64 * 7
76}
77fn slot_output_ep(i: usize) -> u64 {
78 32 + i as u64 * 7
79}
80fn slot_input_ring_mem(i: usize) -> u64 {
81 33 + i as u64 * 7
82}
83fn slot_netsock_mem(i: usize) -> u64 {
84 34 + i as u64 * 7
85}
86fn slot_shell_notif(i: usize) -> u64 {
87 36 + i as u64 * 7
88}
89
90fn input_ring_vaddr(i: usize) -> u64 {
91 0x7000_0000 + i as u64 * 0x1000
92}
93fn netsock_ring_vaddr(i: usize) -> u64 {
94 0x8000_0000 + i as u64 * 0x2000
95}
96
97struct Session<'a> {
98 ssh: Option<SshSession<'a>>,
99 shell_running: bool,
100 input_tx: Option<PacketRingWriter>,
101 netsock_tx: Option<PacketRingWriter>,
102 netsock_rx: Option<PacketRingReader>,
103 last_activity_ms: i64,
104}
105
106impl<'a> Session<'a> {
107 const fn new() -> Self {
108 Self {
109 ssh: None,
110 shell_running: false,
111 input_tx: None,
112 netsock_tx: None,
113 netsock_rx: None,
114 last_activity_ms: 0,
115 }
116 }
117}
118
119fn find_module_by_name(name: &[u8]) -> Option<u64> {
120 let mut dummy = [0u8; 0];
121 let total = syscall::module_info(u64::MAX, &mut dummy);
122 if total < 0 {
123 return None;
124 }
125
126 let mut path_buf = [0u8; 128];
127 (BOOT_MODULE_COUNT..total as u64).find(|&i| {
128 path_buf = [0u8; 128];
129 let path_len = syscall::module_info(i, &mut path_buf);
130 match path_len < 0 {
131 true => false,
132 false => {
133 let used = (path_len as usize).min(path_buf.len());
134 path::bytes_eq(path::extract_filename(&path_buf[..used]), name)
135 }
136 }
137 })
138}
139
140fn send_ssh_bytes(ssh: &mut SshSession, tx: &PacketRingWriter, conn_id: u8, data: &[u8]) {
141 let mut offset = 0usize;
142 core::iter::from_fn(|| match offset < data.len() {
143 true => {
144 let written = ssh.write_plaintext(&data[offset..]);
145 offset += written;
146 flush_ssh_output(ssh, tx, conn_id);
147 Some(())
148 }
149 false => None,
150 })
151 .take(64)
152 .count();
153}
154
155fn flush_ssh_output(ssh: &mut SshSession, tx: &PacketRingWriter, conn_id: u8) -> bool {
156 let mut flushed = false;
157 let mut tmp = [0u8; 126];
158 core::iter::from_fn(|| {
159 let n = ssh.peek_output(&mut tmp[2..]);
160 match n > 0 {
161 true => {
162 tmp[0] = conn_id;
163 tmp[1] = MSG_DATA;
164 match tx.try_push(&tmp[..2 + n]) {
165 true => {
166 ssh.consume_output(n);
167 flushed = true;
168 Some(())
169 }
170 false => None,
171 }
172 }
173 false => None,
174 }
175 })
176 .take(32)
177 .count();
178 if flushed {
179 syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
180 }
181 flushed
182}
183
184fn drive_ssh(ssh: &mut SshSession, tx: &PacketRingWriter, conn_id: u8) -> bool {
185 let mut any_flushed = false;
186 let mut keep_going = true;
187 core::iter::from_fn(|| match keep_going {
188 true => {
189 any_flushed |= flush_ssh_output(ssh, tx, conn_id);
190 let r = ssh.progress();
191 any_flushed |= flush_ssh_output(ssh, tx, conn_id);
192 match r {
193 SshProgress::Continue => Some(()),
194 SshProgress::ShellReady => {
195 any_flushed = true;
196 None
197 }
198 SshProgress::Idle => {
199 keep_going = false;
200 None
201 }
202 SshProgress::Disconnected => {
203 keep_going = false;
204 None
205 }
206 }
207 }
208 false => None,
209 })
210 .take(128)
211 .count();
212 any_flushed
213}
214
215fn spawn_shell(conn_id: usize, module_idx: u64, session: &mut Session) -> bool {
216 let proc_slot = slot_proc(conn_id);
217 let sched_slot = slot_sched(conn_id);
218 let output_ep = slot_output_ep(conn_id);
219 let input_mem = slot_input_ring_mem(conn_id);
220 let netsock_mem = slot_netsock_mem(conn_id);
221 let shell_notif = slot_shell_notif(conn_id);
222
223 let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_ENDPOINT, 0, output_ep, 1);
224 if r < 0 {
225 return false;
226 }
227
228 let child_pid = syscall::proc_create(proc_slot);
229 if child_pid < 0 {
230 syscall::cap_revoke(output_ep);
231 return false;
232 }
233
234 let r = syscall::proc_load_module(proc_slot, module_idx);
235 if r < 0 {
236 syscall::proc_destroy(proc_slot);
237 syscall::cap_revoke(output_ep);
238 return false;
239 }
240
241 let r = syscall::cap_grant(output_ep, proc_slot, 2, RIGHTS_ALL);
242 if r < 0 {
243 syscall::proc_destroy(proc_slot);
244 syscall::cap_revoke(output_ep);
245 return false;
246 }
247
248 let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_FRAME, 0, input_mem, INPUT_RING_PAGES);
249 if r < 0 {
250 syscall::proc_destroy(proc_slot);
251 syscall::cap_revoke(output_ep);
252 return false;
253 }
254 let ir_vaddr = input_ring_vaddr(conn_id);
255 let map_ok = (0..INPUT_RING_PAGES).all(|i| {
256 syscall::frame_map(input_mem + i, ir_vaddr + i * 4096, 1) >= 0
257 });
258 if !map_ok {
259 (0..INPUT_RING_PAGES).for_each(|i| {
260 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
261 });
262 (0..INPUT_RING_PAGES).for_each(|i| {
263 syscall::cap_revoke(input_mem + i);
264 });
265 syscall::proc_destroy(proc_slot);
266 syscall::cap_revoke(output_ep);
267 return false;
268 }
269 let input_tx = unsafe {
270 PacketRingWriter::init(ir_vaddr as *mut u8, INPUT_RING_SIZE, INPUT_RING_SLOT_SIZE)
271 };
272 let r = syscall::cap_grant(input_mem, proc_slot, 1, RIGHTS_ALL);
273 if r < 0 {
274 (0..INPUT_RING_PAGES).for_each(|i| {
275 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
276 });
277 (0..INPUT_RING_PAGES).for_each(|i| {
278 syscall::cap_revoke(input_mem + i);
279 });
280 syscall::proc_destroy(proc_slot);
281 syscall::cap_revoke(output_ep);
282 return false;
283 }
284
285 let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_FRAME, 0, netsock_mem, NETSOCK_RING_PAGES);
286 if r < 0 {
287 (0..INPUT_RING_PAGES).for_each(|i| {
288 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
289 });
290 (0..INPUT_RING_PAGES).for_each(|i| {
291 syscall::cap_revoke(input_mem + i);
292 });
293 syscall::proc_destroy(proc_slot);
294 syscall::cap_revoke(output_ep);
295 return false;
296 }
297 let ns_vaddr = netsock_ring_vaddr(conn_id);
298 let ns_map_ok = (0..NETSOCK_RING_PAGES).all(|i| {
299 syscall::frame_map(netsock_mem + i, ns_vaddr + i * 4096, 1) >= 0
300 });
301 if !ns_map_ok {
302 (0..NETSOCK_RING_PAGES).for_each(|i| {
303 syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
304 });
305 (0..NETSOCK_RING_PAGES).for_each(|i| {
306 syscall::cap_revoke(netsock_mem + i);
307 });
308 (0..INPUT_RING_PAGES).for_each(|i| {
309 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
310 });
311 (0..INPUT_RING_PAGES).for_each(|i| {
312 syscall::cap_revoke(input_mem + i);
313 });
314 syscall::proc_destroy(proc_slot);
315 syscall::cap_revoke(output_ep);
316 return false;
317 }
318 let netsock_tx = unsafe {
319 PacketRingWriter::init(
320 ns_vaddr as *mut u8,
321 NETSOCK_RING_HALF,
322 NETSOCK_RING_SLOT_SIZE,
323 )
324 };
325 let _ = unsafe {
326 PacketRingWriter::init(
327 (ns_vaddr + NETSOCK_RING_HALF as u64) as *mut u8,
328 NETSOCK_RING_HALF,
329 NETSOCK_RING_SLOT_SIZE,
330 )
331 };
332 let netsock_rx = unsafe {
333 PacketRingReader::attach(
334 (ns_vaddr + NETSOCK_RING_HALF as u64) as *mut u8,
335 NETSOCK_RING_HALF,
336 )
337 };
338 let netsock_grant_ok = (0..NETSOCK_RING_PAGES).all(|i| {
339 syscall::cap_grant(netsock_mem + i, proc_slot, 3 + i, RIGHTS_ALL) >= 0
340 });
341 if !netsock_grant_ok {
342 (0..NETSOCK_RING_PAGES).for_each(|i| {
343 syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
344 });
345 (0..NETSOCK_RING_PAGES).for_each(|i| {
346 syscall::cap_revoke(netsock_mem + i);
347 });
348 (0..INPUT_RING_PAGES).for_each(|i| {
349 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
350 });
351 (0..INPUT_RING_PAGES).for_each(|i| {
352 syscall::cap_revoke(input_mem + i);
353 });
354 syscall::proc_destroy(proc_slot);
355 syscall::cap_revoke(output_ep);
356 return false;
357 }
358
359 let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_NOTIFICATION, 0, shell_notif, 1);
360 if r < 0 {
361 (0..NETSOCK_RING_PAGES).for_each(|i| {
362 syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
363 });
364 (0..NETSOCK_RING_PAGES).for_each(|i| {
365 syscall::cap_revoke(netsock_mem + i);
366 });
367 (0..INPUT_RING_PAGES).for_each(|i| {
368 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
369 });
370 (0..INPUT_RING_PAGES).for_each(|i| {
371 syscall::cap_revoke(input_mem + i);
372 });
373 syscall::proc_destroy(proc_slot);
374 syscall::cap_revoke(output_ep);
375 return false;
376 }
377 let r = syscall::cap_grant(shell_notif, proc_slot, 9, RIGHTS_ALL);
378 if r < 0 {
379 syscall::cap_revoke(shell_notif);
380 (0..NETSOCK_RING_PAGES).for_each(|i| {
381 syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
382 });
383 (0..NETSOCK_RING_PAGES).for_each(|i| {
384 syscall::cap_revoke(netsock_mem + i);
385 });
386 (0..INPUT_RING_PAGES).for_each(|i| {
387 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
388 });
389 (0..INPUT_RING_PAGES).for_each(|i| {
390 syscall::cap_revoke(input_mem + i);
391 });
392 syscall::proc_destroy(proc_slot);
393 syscall::cap_revoke(output_ep);
394 return false;
395 }
396
397 let _ = syscall::cap_grant(OWN_NOTIF_SLOT, proc_slot, 10, 0b0010);
398
399 (0..VFS_RING_FRAMES).for_each(|i| {
400 let _ = syscall::cap_grant(VFS_RING_BASE + i, proc_slot, 11 + i, RIGHTS_ALL);
401 });
402 [
403 (VFS_CLIENT_NOTIF_SLOT, 28u64),
404 (VFS_NOTIF_SLOT, 29),
405 (VFS_PROC_SLOT, 15),
406 ]
407 .iter()
408 .for_each(|&(src, dst)| {
409 let _ = syscall::cap_grant(src, proc_slot, dst, RIGHTS_ALL);
410 });
411
412 let death_badge = NOTIFY_BIT_SHELL_DEATH_BASE << conn_id as u64;
413 syscall::proc_bind_death_notif(proc_slot, OWN_NOTIF_SLOT, death_badge);
414
415 let r = syscall::untyped_retype(SLOT_UNTYPED, TAG_SCHED_CONTEXT, 0, sched_slot, 1);
416 if r < 0 {
417 cleanup_shell_caps(conn_id);
418 return false;
419 }
420
421 let r = syscall::sched_configure(sched_slot, 10_000_000, 10_000_000, 100);
422 if r < 0 {
423 cleanup_shell_caps(conn_id);
424 return false;
425 }
426
427 let r = syscall::sched_attach(sched_slot, child_pid as u64);
428 if r < 0 {
429 cleanup_shell_caps(conn_id);
430 return false;
431 }
432
433 let r = syscall::proc_start(proc_slot, 2);
434 if r < 0 {
435 cleanup_shell_caps(conn_id);
436 return false;
437 }
438
439 syscall::notify_signal(VFS_CLIENT_NOTIF_SLOT, 1);
440
441 session.shell_running = true;
442 session.input_tx = Some(input_tx);
443 session.netsock_tx = Some(netsock_tx);
444 session.netsock_rx = Some(netsock_rx);
445
446 true
447}
448
449fn cleanup_shell_caps(conn_id: usize) {
450 let ir_vaddr = input_ring_vaddr(conn_id);
451 let ns_vaddr = netsock_ring_vaddr(conn_id);
452 let input_mem = slot_input_ring_mem(conn_id);
453 let netsock_mem = slot_netsock_mem(conn_id);
454
455 (0..INPUT_RING_PAGES).for_each(|i| {
456 syscall::frame_unmap(input_mem + i, ir_vaddr + i * 4096);
457 });
458 (0..INPUT_RING_PAGES).for_each(|i| {
459 syscall::cap_revoke(input_mem + i);
460 });
461
462 (0..NETSOCK_RING_PAGES).for_each(|i| {
463 syscall::frame_unmap(netsock_mem + i, ns_vaddr + i * 4096);
464 });
465 (0..NETSOCK_RING_PAGES).for_each(|i| {
466 syscall::cap_revoke(netsock_mem + i);
467 });
468
469 syscall::cap_revoke(slot_shell_notif(conn_id));
470 syscall::cap_revoke(slot_sched(conn_id));
471 syscall::cap_revoke(slot_output_ep(conn_id));
472}
473
474fn drain_shell_output(conn_id: usize, ssh: &mut SshSession, tx: &PacketRingWriter) -> bool {
475 let ep_slot = slot_output_ep(conn_id);
476 let cid = conn_id as u8;
477 let mut ipc_raw = [0u8; MAX_IPC_BYTES];
478 let mut had_data = false;
479
480 core::iter::from_fn(|| {
481 let (status, msg) = syscall::nb_recv(ep_slot);
482 match status >= 0 {
483 true => {
484 let count = unpack_bytes(&msg, &mut ipc_raw);
485 const _: () = assert!(MAX_IPC_BYTES * 2 <= 80);
486 if count > 0 {
487 had_data = true;
488 let mut translated = [0u8; MAX_IPC_BYTES * 2];
489 let tlen = (0..count).fold(0usize, |pos, i| match ipc_raw[i] {
490 b'\n' => {
491 translated[pos] = b'\r';
492 translated[pos + 1] = b'\n';
493 pos + 2
494 }
495 b => {
496 translated[pos] = b;
497 pos + 1
498 }
499 });
500 send_ssh_bytes(ssh, tx, cid, &translated[..tlen]);
501 }
502 Some(())
503 }
504 false => None,
505 }
506 })
507 .take(64)
508 .count();
509 had_data
510}
511
512fn drain_shell_netsock(
513 conn_id: usize,
514 netsock_rx: &PacketRingReader,
515 netstack_tx: &PacketRingWriter,
516) -> bool {
517 let cid = conn_id as u8;
518 let mut buf = [0u8; 128];
519 let mut relayed = false;
520
521 core::iter::from_fn(|| match netsock_rx.try_pop(&mut buf) {
522 Some(n) if n > 0 => {
523 buf[0] = cid;
524 match netstack_tx.try_push(&buf[..n]) {
525 true => {
526 relayed = true;
527 Some(())
528 }
529 false => None,
530 }
531 }
532 _ => None,
533 })
534 .take(16)
535 .count();
536
537 relayed
538}
539
540fn relay_netstack_to_shell(netsock_tx: &PacketRingWriter, data: &[u8], shell_notif: u64) {
541 let mut buf = [0u8; 128];
542 let copy_len = data.len().min(buf.len() - 1);
543 buf[0] = 0;
544 buf[1..1 + copy_len].copy_from_slice(&data[..copy_len]);
545 if netsock_tx.try_push(&buf[..1 + copy_len]) {
546 syscall::notify_signal(shell_notif, 4);
547 }
548}
549
550fn teardown_session(conn_id: usize, session: &mut Session, tx: &PacketRingWriter) {
551 if session.shell_running {
552 syscall::proc_destroy(slot_proc(conn_id));
553 cleanup_shell_caps(conn_id);
554 session.shell_running = false;
555 session.input_tx = None;
556 session.netsock_tx = None;
557 session.netsock_rx = None;
558 }
559 if let Some(ssh) = session.ssh.as_mut() {
560 ssh.close();
561 }
562 let ctrl = [conn_id as u8, MSG_DISCONNECT];
563 let _ = tx.try_push(&ctrl);
564 drop(session.ssh.take());
565}
566
567#[unsafe(no_mangle)]
568pub extern "C" fn lancer_main() -> ! {
569 lancer_user::show!(rsh, "starting");
570
571 if !(0..SHELL_RING_FRAME_COUNT).all(|i| syscall::frame_map(SHELL_RING_BASE_SLOT + i, RING_BASE + i * 4096, 1) >= 0) {
572 lancer_user::show!(rsh, error, "ring frame_map failed");
573 syscall::exit();
574 }
575
576 let rx_base = RING_BASE as *mut u8;
577 let tx_base = (RING_BASE + RING_HALF as u64) as *mut u8;
578
579 let rx = unsafe { PacketRingReader::attach(rx_base, RING_HALF) };
580 let tx = unsafe { PacketRingWriter::init(tx_base, RING_HALF, 128) };
581
582 let host_key = match sunset::SignKey::generate(sunset::KeyType::Ed25519, None) {
583 Ok(k) => k,
584 Err(_) => {
585 lancer_user::show!(rsh, error, "keygen failed");
586 syscall::exit();
587 }
588 };
589
590 let shell_module_idx = find_module_by_name(b"shell");
591 match shell_module_idx {
592 Some(_) => lancer_user::show!(rsh, "shell module found"),
593 None => lancer_user::show!(rsh, warn, "shell module not found"),
594 }
595
596 lancer_user::show!(rsh, "ready");
597
598 let mut sessions: [Session; MAX_SESSIONS] = [
599 Session::new(),
600 Session::new(),
601 Session::new(),
602 Session::new(),
603 ];
604 let mut pop_buf = [0u8; 128];
605 let mut plaintext_buf = [0u8; 256];
606
607 syscall::ntfn_bind(OWN_NOTIF_SLOT);
608
609 loop {
610 let active_shell = (0..MAX_SESSIONS).find(|&i| sessions[i].shell_running);
611 let mut any_output = false;
612
613 let bits = match active_shell {
614 Some(cid) => {
615 let (status, msg) = syscall::recv(slot_output_ep(cid));
616 match status == syscall::BOUND_NOTIFICATION_BADGE {
617 true => msg[1],
618 false => {
619 if status >= 0 {
620 let Session { ref mut ssh, .. } = sessions[cid];
621 if let Some(ssh) = ssh.as_mut() {
622 let mut ipc_raw = [0u8; MAX_IPC_BYTES];
623 let count = unpack_bytes(&msg, &mut ipc_raw);
624 if count > 0 {
625 let mut translated = [0u8; MAX_IPC_BYTES * 2];
626 let tlen = (0..count).fold(0usize, |pos, i| match ipc_raw[i] {
627 b'\n' => {
628 translated[pos] = b'\r';
629 translated[pos + 1] = b'\n';
630 pos + 2
631 }
632 b => {
633 translated[pos] = b;
634 pos + 1
635 }
636 });
637 send_ssh_bytes(ssh, &tx, cid as u8, &translated[..tlen]);
638 any_output = true;
639 }
640 }
641 }
642 let (_, poll_bits) = syscall::notify_poll(OWN_NOTIF_SLOT);
643 poll_bits
644 }
645 }
646 }
647 None => {
648 let (_, wait_bits) = syscall::notify_wait(OWN_NOTIF_SLOT);
649 wait_bits
650 }
651 };
652
653 let now_ms = syscall::clock_monotonic_ms().max(0);
654
655 (0..MAX_SESSIONS).for_each(|i| {
656 let death_bit = NOTIFY_BIT_SHELL_DEATH_BASE << i as u64;
657 if bits & death_bit == 0 || !sessions[i].shell_running {
658 return;
659 }
660
661 if let Some(ssh) = sessions[i].ssh.as_mut() {
662 drain_shell_output(i, ssh, &tx);
663 }
664 syscall::proc_destroy(slot_proc(i));
665 cleanup_shell_caps(i);
666 sessions[i].shell_running = false;
667 sessions[i].input_tx = None;
668 sessions[i].netsock_tx = None;
669 sessions[i].netsock_rx = None;
670 any_output = true;
671
672 if let Some(ssh) = sessions[i].ssh.as_mut() {
673 ssh.close();
674 }
675 let ctrl = [i as u8, MSG_DISCONNECT];
676 let _ = tx.try_push(&ctrl);
677 drop(sessions[i].ssh.take());
678 });
679
680 if bits & NOTIFY_BIT_SHELL_NETSOCK != 0 {
681 (0..MAX_SESSIONS).for_each(|i| {
682 if let Some(ref netsock_rx) = sessions[i].netsock_rx {
683 let relayed = drain_shell_netsock(i, netsock_rx, &tx);
684 if relayed {
685 syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
686 any_output = true;
687 }
688 }
689 });
690 }
691
692 (0..MAX_SESSIONS).for_each(|i| {
693 if let Some(ssh) = sessions[i].ssh.as_mut() {
694 any_output |= drive_ssh(ssh, &tx, i as u8);
695 }
696 });
697
698 let mut sessions_to_drop: u8 = 0;
699 let mut need_spawn: u8 = 0;
700
701 core::iter::from_fn(|| match rx.try_pop(&mut pop_buf) {
702 Some(n) if n >= 2 => {
703 let conn_id = pop_buf[0] as usize;
704 let msg_type = pop_buf[1];
705 if conn_id >= MAX_SESSIONS {
706 return Some(());
707 }
708
709 match msg_type {
710 MSG_CONNECT => {
711 if sessions[conn_id].shell_running {
712 syscall::proc_destroy(slot_proc(conn_id));
713 cleanup_shell_caps(conn_id);
714 sessions[conn_id].shell_running = false;
715 sessions[conn_id].input_tx = None;
716 sessions[conn_id].netsock_tx = None;
717 sessions[conn_id].netsock_rx = None;
718 }
719 drop(sessions[conn_id].ssh.take());
720 let inbuf: &'static mut [u8; 35000] = unsafe {
721 &mut *((&raw mut SSH_INBUFS) as *mut [u8; 35000]).add(conn_id)
722 };
723 let outbuf: &'static mut [u8; 35000] = unsafe {
724 &mut *((&raw mut SSH_OUTBUFS) as *mut [u8; 35000]).add(conn_id)
725 };
726 inbuf.fill(0);
727 outbuf.fill(0);
728 sessions[conn_id].ssh =
729 Some(SshSession::new(inbuf, outbuf, host_key.clone()));
730 sessions[conn_id].last_activity_ms = now_ms;
731 if let Some(ssh) = sessions[conn_id].ssh.as_mut() {
732 any_output |= drive_ssh(ssh, &tx, conn_id as u8);
733 syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
734 }
735 }
736 MSG_DISCONNECT => {
737 teardown_session(conn_id, &mut sessions[conn_id], &tx);
738 }
739 MSG_DATA if n > 2 => {
740 let Session {
741 ref mut ssh,
742 ref input_tx,
743 ref shell_running,
744 ref mut last_activity_ms,
745 ..
746 } = sessions[conn_id];
747 *last_activity_ms = now_ms;
748 if let Some(ssh) = ssh.as_mut() {
749 let data = &pop_buf[2..n];
750 let mut off = 0usize;
751 core::iter::from_fn(|| match off < data.len() {
752 true => {
753 let accepted = ssh.feed_input(&data[off..]);
754 match accepted {
755 0 => {
756 any_output |= drive_ssh(ssh, &tx, conn_id as u8);
757 if ssh.shell_ready && *shell_running {
758 let pt_n = ssh.read_plaintext(&mut plaintext_buf);
759 if pt_n > 0
760 && let Some(it) = input_tx
761 && it.try_push(&plaintext_buf[..pt_n])
762 {
763 syscall::notify_signal(
764 slot_shell_notif(conn_id),
765 2,
766 );
767 }
768 }
769 any_output |= flush_ssh_output(ssh, &tx, conn_id as u8);
770 let retry = ssh.feed_input(&data[off..]);
771 match retry {
772 0 => None,
773 n2 => {
774 off += n2;
775 Some(())
776 }
777 }
778 }
779 n2 => {
780 off += n2;
781 Some(())
782 }
783 }
784 }
785 false => None,
786 })
787 .take(32)
788 .count();
789
790 any_output |= drive_ssh(ssh, &tx, conn_id as u8);
791 if ssh.shell_ready && *shell_running {
792 let pt_n = ssh.read_plaintext(&mut plaintext_buf);
793 if pt_n > 0
794 && let Some(it) = input_tx
795 && it.try_push(&plaintext_buf[..pt_n])
796 {
797 syscall::notify_signal(slot_shell_notif(conn_id), 2);
798 }
799 }
800 }
801 }
802 MSG_DNS_RESULT
803 | MSG_DNS_RESULT6
804 | MSG_PING_RESULT
805 | MSG_PING_DONE
806 | MSG_UDP_RECV
807 | MSG_UDP_CLOSE
808 | MSG_UDP_BOUND
809 | MSG_IFCONFIG_RESULT
810 | MSG_ARP_ENTRY
811 | MSG_ARP_DONE
812 | MSG_NETSTAT_ENTRY
813 | MSG_NETSTAT_ENTRY_UDP
814 | MSG_NETSTAT_DONE
815 if sessions[conn_id].shell_running =>
816 {
817 if let Some(ref netsock_tx) = sessions[conn_id].netsock_tx {
818 relay_netstack_to_shell(
819 netsock_tx,
820 &pop_buf[1..n],
821 slot_shell_notif(conn_id),
822 );
823 }
824 any_output = true;
825 }
826 _ => {}
827 }
828 Some(())
829 }
830 _ => None,
831 })
832 .take(64)
833 .count();
834
835 (0..MAX_SESSIONS).for_each(|i| {
836 let Session {
837 ref mut ssh,
838 ref shell_running,
839 ref input_tx,
840 ..
841 } = sessions[i];
842 if let Some(ssh) = ssh.as_mut() {
843 any_output |= drive_ssh(ssh, &tx, i as u8);
844 if ssh.shell_ready && !*shell_running {
845 need_spawn |= 1 << i;
846 }
847 if ssh.shell_ready && *shell_running {
848 let pt_n = ssh.read_plaintext(&mut plaintext_buf);
849 if pt_n > 0
850 && let Some(it) = input_tx
851 && it.try_push(&plaintext_buf[..pt_n])
852 {
853 syscall::notify_signal(slot_shell_notif(i), 2);
854 }
855 }
856 any_output |= flush_ssh_output(ssh, &tx, i as u8);
857 }
858 });
859
860 (0..MAX_SESSIONS).for_each(|i| {
861 if need_spawn & (1 << i) == 0 {
862 return;
863 }
864 match shell_module_idx {
865 Some(idx) => {
866 let ok = spawn_shell(i, idx, &mut sessions[i]);
867 if !ok {
868 if let Some(ssh) = sessions[i].ssh.as_mut() {
869 send_ssh_bytes(ssh, &tx, i as u8, b"Failed to start shell\r\n");
870 flush_ssh_output(ssh, &tx, i as u8);
871 ssh.close();
872 }
873 sessions_to_drop |= 1 << i;
874 }
875 }
876 None => {
877 if let Some(ssh) = sessions[i].ssh.as_mut() {
878 send_ssh_bytes(ssh, &tx, i as u8, b"Shell module not found\r\n");
879 flush_ssh_output(ssh, &tx, i as u8);
880 ssh.close();
881 }
882 sessions_to_drop |= 1 << i;
883 }
884 }
885 });
886
887 (0..MAX_SESSIONS).for_each(|i| {
888 if sessions[i].shell_running
889 && let Some(ssh) = sessions[i].ssh.as_mut()
890 {
891 any_output |= drain_shell_output(i, ssh, &tx);
892 }
893 });
894
895 (0..MAX_SESSIONS).for_each(|i| {
896 if sessions_to_drop & (1 << i) != 0 {
897 teardown_session(i, &mut sessions[i], &tx);
898 }
899 });
900
901 (0..MAX_SESSIONS).for_each(|i| {
902 if sessions[i].ssh.is_none() {
903 return;
904 }
905 if now_ms - sessions[i].last_activity_ms > IDLE_TIMEOUT_MS {
906 teardown_session(i, &mut sessions[i], &tx);
907 }
908 });
909
910 if any_output {
911 syscall::notify_signal(NETSTACK_NOTIF_SLOT, 4);
912 }
913 }
914}