refactor: changes

src/instance.rs

@@ -33,6 +33,7 @@ const SSH_GUEST_DIR: &str = "/root/.vibebox";
 const DEFAULT_SSH_USER: &str = "vibebox";
 const SSH_CONNECT_RETRIES: usize = 20;
 const SSH_CONNECT_DELAY_MS: u64 = 500;
+const SSH_SETUP_SCRIPT: &str = include_str!("ssh.sh");
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct InstanceConfig {
@@ -61,7 +62,12 @@ pub fn run_with_ssh(
     let instance_dir = ensure_instance_dir(&project_root)?;
     let (ssh_key, _ssh_pub) = ensure_ssh_keypair(&instance_dir)?;
 
-    let config = load_or_create_instance_config(&instance_dir)?;
+    let mut config = load_or_create_instance_config(&instance_dir)?;
+    // Clear cached IP to avoid reusing a stale address on startup.
+    if config.vm_ipv4.is_some() {
+        config.vm_ipv4 = None;
+        write_instance_config(&instance_dir.join(INSTANCE_TOML), &config)?;
+    }
     let config = Arc::new(Mutex::new(config));
 
     let extra_shares = vec![DirectoryShare::new(
@@ -244,40 +250,15 @@ fn build_ssh_login_actions(
 
     let key_path = format!("{guest_dir}/{key_name}.pub");
 
-    let mask_cache = format!(
-        "if [ -d /root/{project_name}/.vibebox ]; then mount -t tmpfs tmpfs /root/{project_name}/.vibebox; fi"
-    );
+    let setup_script = SSH_SETUP_SCRIPT
+        .replace("__SSH_USER__", &ssh_user)
+        .replace("__SUDO_PASSWORD__", &sudo_password)
+        .replace("__PROJECT_NAME__", project_name)
+        .replace("__KEY_PATH__", &key_path);
+    let setup = vm::script_command_from_content("ssh_setup", &setup_script)
+        .expect("ssh setup script contained invalid marker");
 
-    let setup_lines = [
-        "if ! command -v sshd >/dev/null 2>&1; then apt-get update && apt-get install -y openssh-server sudo; fi",
-        "systemctl enable ssh >/dev/null 2>&1 || true",
-        &format!("id -u {ssh_user} >/dev/null 2>&1 || useradd -m -s /bin/bash {ssh_user}"),
-        &format!("echo \"{ssh_user}:{sudo_password}\" | chpasswd"),
-        &format!("usermod -aG sudo {ssh_user}"),
-        &format!("install -d -m 700 /home/{ssh_user}/.ssh"),
-        &format!("install -m 600 {key_path} /home/{ssh_user}/.ssh/authorized_keys"),
-        &format!("chown -R {ssh_user}:{ssh_user} /home/{ssh_user}/.ssh"),
-        "mkdir -p /etc/ssh/sshd_config.d",
-        "cat >/etc/ssh/sshd_config.d/vibebox.conf <<'VIBEBOX_SSHD'",
-        "PasswordAuthentication no",
-        "KbdInteractiveAuthentication no",
-        "ChallengeResponseAuthentication no",
-        "PubkeyAuthentication yes",
-        "PermitRootLogin no",
-        &format!("AllowUsers {ssh_user}"),
-        "VIBEBOX_SSHD",
-        "systemctl restart ssh",
-        "echo VIBEBOX_SSH_READY",
-    ];
-    let setup = setup_lines.join("\n");
-
-    let ip_probe = "while true; do ip=$(ip -4 -o addr show scope global | awk '{print $4}' | cut -d/ -f1 | head -n 1); if [ -n \"$ip\" ]; then echo VIBEBOX_IPV4=$ip; break; fi; sleep 1; done";
-
-    vec![
-        LoginAction::Send(mask_cache),
-        LoginAction::Send(setup),
-        LoginAction::Send(ip_probe.to_string()),
-    ]
+    vec![LoginAction::Send(setup)]
 }
 
 fn spawn_ssh_io(

src/ssh.sh

@@ -0,0 +1,46 @@
+#!/bin/sh
+
+SSH_USER="__SSH_USER__"
+SUDO_PASSWORD="__SUDO_PASSWORD__"
+PROJECT_NAME="__PROJECT_NAME__"
+KEY_PATH="__KEY_PATH__"
+
+if [ -d /root/${PROJECT_NAME}/.vibebox ]; then
+  mount -t tmpfs tmpfs /root/${PROJECT_NAME}/.vibebox
+fi
+
+if ! command -v sshd >/dev/null 2>&1; then
+  apt-get update && apt-get install -y openssh-server sudo
+fi
+
+systemctl enable ssh >/dev/null 2>&1 || true
+id -u ${SSH_USER} >/dev/null 2>&1 || useradd -m -s /bin/bash ${SSH_USER}
+echo "${SSH_USER}:${SUDO_PASSWORD}" | chpasswd
+usermod -aG sudo ${SSH_USER}
+install -d -m 700 /home/${SSH_USER}/.ssh
+install -m 600 ${KEY_PATH} /home/${SSH_USER}/.ssh/authorized_keys
+chown -R ${SSH_USER}:${SSH_USER} /home/${SSH_USER}/.ssh
+rm -f /home/${SSH_USER}/.bash_logout
+mkdir -p /etc/ssh/sshd_config.d
+cat >/etc/ssh/sshd_config.d/vibebox.conf <<'VIBEBOX_SSHD'
+PasswordAuthentication no
+KbdInteractiveAuthentication no
+ChallengeResponseAuthentication no
+PubkeyAuthentication yes
+PermitRootLogin no
+AllowUsers __SSH_USER__
+VIBEBOX_SSHD
+systemctl restart ssh
+echo VIBEBOX_SSH_READY
+
+echo "=== generated network file ==="
+sed -n '1,200p' /run/systemd/network/10-netplan-all-en.network || true
+
+while true; do
+  ip=$(ip -4 -o addr show scope global | awk '{print $4}' | cut -d/ -f1 | head -n 1)
+  if [ -n "$ip" ]; then
+    echo VIBEBOX_IPV4=$ip
+    break
+  fi
+  sleep 1
+done

src/vm.rs

@@ -415,7 +415,7 @@ fn script_command_from_path(
     script_command_from_content(&label, &script)
 }
 
-fn script_command_from_content(
+pub(crate) fn script_command_from_content(
     label: &str,
     script: &str,
 ) -> Result<String, Box<dyn std::error::Error>> {
@@ -725,7 +725,12 @@ where
         Error,
     }
 
-    fn poll_with_wakeup<'a>(main_fd: RawFd, wakeup_fd: RawFd, buf: &'a mut [u8]) -> PollResult<'a> {
+    fn poll_with_wakeup<'a>(
+        main_fd: RawFd,
+        wakeup_fd: RawFd,
+        buf: &'a mut [u8],
+        timeout_ms: i32,
+    ) -> PollResult<'a> {
         let mut fds = [
             libc::pollfd {
                 fd: main_fd,
@@ -739,8 +744,12 @@ where
             },
         ];
 
-        let ret = unsafe { libc::poll(fds.as_mut_ptr(), 2, -1) };
-        if ret <= 0 || fds[1].revents & libc::POLLIN != 0 {
+        let ret = unsafe { libc::poll(fds.as_mut_ptr(), 2, timeout_ms) };
+        if ret == 0 {
+            PollResult::Spurious
+        } else if ret < 0 {
+            PollResult::Error
+        } else if fds[1].revents & libc::POLLIN != 0 {
             PollResult::Shutdown
         } else if fds[0].revents & libc::POLLIN != 0 {
             let n = unsafe { libc::read(main_fd, buf.as_mut_ptr() as *mut _, buf.len()) };
@@ -786,7 +795,12 @@ where
                     continue;
                 }
 
-                match poll_with_wakeup(libc::STDIN_FILENO, wakeup_read.as_raw_fd(), &mut buf) {
+                match poll_with_wakeup(
+                    libc::STDIN_FILENO,
+                    wakeup_read.as_raw_fd(),
+                    &mut buf,
+                    -1,
+                ) {
                     PollResult::Shutdown | PollResult::Error => break,
                     PollResult::Spurious => continue,
                     PollResult::Ready(bytes) => {
@@ -839,7 +853,12 @@ where
                     let mut guard = raw_guard.lock().unwrap();
                     *guard = None;
                 }
-                match poll_with_wakeup(vm_output_fd.as_raw_fd(), wakeup_read.as_raw_fd(), &mut buf)
+                match poll_with_wakeup(
+                    vm_output_fd.as_raw_fd(),
+                    wakeup_read.as_raw_fd(),
+                    &mut buf,
+                    100,
+                )
                 {
                     PollResult::Shutdown | PollResult::Error => break,
                     PollResult::Spurious => continue,