fix(ipi): validate initialization and atom-count framing

doc/third-party/ipi.md

@@ -41,3 +41,17 @@ The option **`port`** should be the same as that in input.xml:
 The option **`graph_file`** provides the file name of the frozen model. The model can have either double or single float precision interface.
 
 The `dp_ipi` gets the atom names from an [XYZ file](https://en.wikipedia.org/wiki/XYZ_file_format) provided by **`coord_file`** (meanwhile ignores all coordinates in it) and translates the names to atom types by rules provided by **`atom_type`**.
+
+The client accepts the standard i-PI `INIT` message (replica ID, payload length,
+and payload). Its model and atom types remain configured by the JSON input;
+initialization payloads are read in chunks with no additional size cap beyond
+the signed 32-bit protocol length limit. Negative payload lengths are rejected
+with an error.
+
+Socket reads assemble each requested protocol field completely, including when
+its bytes arrive in separate packets. Interrupted reads are retried. If the
+server closes the connection before a field is complete, the client rejects the
+incomplete read and exits with an error. Bytes received before EOF have already
+been read from the socket, and verbose mode may have printed earlier complete
+INIT chunks. To request a successful client shutdown, send the standard `EXIT`
+message rather than closing the socket without it.

source/ipi/driver.cc

@@ -1,4 +1,5 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
+#include <algorithm>
 #include <cstdint>
 #include <exception>
 #include <fstream>
@@ -162,11 +163,33 @@ int main(int argc, char* argv[]) {
       }
     } else if (header_str == "INIT") {
       assert(4 == sizeof(int32_t));
+      int32_t replica_id;
+      readbuffer_(&socket, (char*)(&replica_id), sizeof(int32_t));
       readbuffer_(&socket, (char*)(&cbuf), sizeof(int32_t));
-      readbuffer_(&socket, initbuffer, cbuf);
+      if (cbuf < 0) {
+        std::cerr << "dp_ipi: INIT payload length must be nonnegative."
+                  << std::endl;
+        return 1;
+      }
+      if (b_verb) {
+        std::cout << "Init sys from wrapper, using ";
+      }
+      // This driver is configured by its input file. Consume the optional
+      // initialization string without imposing a fixed payload-size limit.
+      while (cbuf > 0) {
+        const int chunk =
+            std::min(cbuf, static_cast<int32_t>(sizeof(initbuffer)));
+        readbuffer_(&socket, initbuffer, chunk);
+        if (b_verb) {
+          std::cout.write(initbuffer, chunk);
+        }
+        cbuf -= chunk;
+      }
       if (b_verb) {
-        std::cout << "Init sys from wrapper, using " << initbuffer << std::endl;
+        std::cout << std::endl;
       }
+      isinit = true;
+      hasdata = false;
     } else if (header_str == "POSDATA") {
       assert(8 == sizeof(double));
 
@@ -179,6 +202,12 @@ int main(int argc, char* argv[]) {
 
       // get number of atoms
       readbuffer_(&socket, (char*)(&cbuf), sizeof(int32_t));
+      if (cbuf <= 0 || static_cast<size_t>(cbuf) != dtype.size()) {
+        std::cerr << "dp_ipi: POSDATA atom count " << cbuf
+                  << " does not match configured atom count " << dtype.size()
+                  << "." << std::endl;
+        return 1;
+      }
       if (natoms < 0) {
         natoms = cbuf;
         if (b_verb) {
@@ -225,6 +254,8 @@ int main(int argc, char* argv[]) {
       writebuffer_(&socket, (char*)(&cbuf), sizeof(int32_t));
       writebuffer_(&socket, msg_nothing, 7);
       hasdata = false;
+    } else if (header_str == "EXIT") {
+      break;
     } else {
       std::cerr << "unexpected header " << std::endl;
       return 1;

source/ipi/include/sockets.h

@@ -39,7 +39,9 @@ void open_socket_(int* psockfd, int* inet, int* port, const char* host);
 */
 void writebuffer_(int* psockfd, char* data, int len);
 
-/* Reads from a socket.
+/* Reads exactly len bytes from a socket. Interrupted reads are retried.
+   Exits on premature EOF, another read error, or a negative length.
+   A zero-length request succeeds without reading.
    Args:
    psockfd: The id of the socket that will be read from.
    data: The storage array for data read from the socket.

source/ipi/src/sockets.c

@@ -110,6 +110,34 @@ int deepmd_write_all(int sockfd,
   return 0;
 }
 
+int deepmd_read_all(int sockfd,
+                    char* data,
+                    size_t len,
+                    deepmd_socket_read_fn read_fn) {
+  size_t received = 0;
+  if (read_fn == NULL || (data == NULL && len != 0)) {
+    errno = EINVAL;
+    return -1;
+  }
+  while (received < len) {
+    ssize_t count = read_fn(sockfd, data + received, len - received);
+    if (count > 0) {
+      if ((size_t)count > len - received) {
+        errno = EIO;
+        return -1;
+      }
+      received += (size_t)count;
+    } else if (count == 0) {
+      // A partial frame must never be consumed as a complete protocol field.
+      errno = ECONNRESET;
+      return -1;
+    } else if (errno != EINTR) {
+      return -1;
+    }
+  }
+  return 0;
+}
+
 void error(const char* msg)
 // Prints an error message and then exits.
 {
@@ -213,17 +241,11 @@ void readbuffer_(int* psockfd, char* data, int len)
 */
 
 {
-  int n, nr;
-  int sockfd = *psockfd;
-
-  n = nr = read(sockfd, data, len);
-
-  while (nr > 0 && n < len) {
-    nr = read(sockfd, &data[n], len - n);
-    n += nr;
+  if (len < 0) {
+    errno = EINVAL;
+    error("Error reading from socket: invalid buffer length");
   }
-
-  if (n == 0) {
+  if (deepmd_read_all(*psockfd, data, (size_t)len, read) < 0) {
     error("Error reading from socket: server has quit or connection broke");
   }
 }

source/ipi/src/sockets_internal.h

@@ -19,3 +19,10 @@ int deepmd_write_all(int sockfd,
                      const char* data,
                      size_t len,
                      deepmd_socket_write_fn write_fn);
+
+// Receive exactly len bytes; EOF before completion is an error.
+typedef ssize_t (*deepmd_socket_read_fn)(int, void*, size_t);
+int deepmd_read_all(int sockfd,
+                    char* data,
+                    size_t len,
+                    deepmd_socket_read_fn read_fn);

source/ipi/tests/test_driver_exit.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import json
+import os
+import socket
+import subprocess
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+
+@unittest.skipIf(
+    os.environ.get("ENABLE_PYTORCH", "1") != "1",
+    reason="Skip test because PyTorch support is not enabled.",
+)
+class TestDPIPIExit(unittest.TestCase):
+    def test_exit_header_stops_driver_cleanly(self) -> None:
+        tests_path = Path(__file__).parent.parent.parent / "tests"
+        model_file = tests_path / "infer" / "deeppot_sea.pth"
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir)
+            coord_file = tmp_path / "coord.xyz"
+            config_file = tmp_path / "config.json"
+            socket_name = f"dpipi_exit_{os.getpid()}_{tmp_path.name}"
+            socket_path = Path("/tmp") / f"ipi_{socket_name}"
+            coord_file.write_text(
+                "1\nEXIT protocol regression test\nO 0.0 0.0 0.0\n",
+                encoding="utf-8",
+            )
+            config_file.write_text(
+                json.dumps(
+                    {
+                        "verbose": False,
+                        "use_unix": True,
+                        "port": 31415,
+                        "host": socket_name,
+                        "graph_file": str(model_file),
+                        "coord_file": str(coord_file),
+                        "atom_type": {"O": 0},
+                    }
+                ),
+                encoding="utf-8",
+            )
+
+            server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            process = None
+            try:
+                server.bind(str(socket_path))
+                server.listen(1)
+                server.settimeout(30)
+                process = subprocess.Popen(
+                    ["dp_ipi", str(config_file)],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                )
+                connection, _ = server.accept()
+                with connection:
+                    connection.sendall(b"EXIT        ")
+                stdout, stderr = process.communicate(timeout=60)
+                self.assertEqual(
+                    process.returncode,
+                    0,
+                    msg=f"stdout:\n{stdout}\nstderr:\n{stderr}",
+                )
+            finally:
+                if process is not None and process.poll() is None:
+                    process.kill()
+                    process.wait()
+                server.close()
+                socket_path.unlink(missing_ok=True)

source/ipi/tests/test_driver_initialization.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Validate framing of the standard i-PI initialization message."""
+
+import json
+import os
+import socket
+import struct
+import subprocess
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+
+@unittest.skipIf(
+    os.environ.get("ENABLE_PYTORCH", "1") != "1",
+    reason="Skip test because PyTorch support is not enabled.",
+)
+class TestDPIPIInitialization(unittest.TestCase):
+    def exchange(
+        self, payload: bytes, length: int | None = None, *, verbose: bool = False
+    ) -> None:
+        """Check INIT framing, optionally advertising an invalid payload length.
+
+        A complete payload must preserve the following STATUS header. A short
+        payload is followed by EOF and must fail before its bytes are logged.
+        """
+        model = Path(__file__).resolve().parents[2] / "tests/infer/deeppot_sea.pth"
+        with tempfile.TemporaryDirectory() as temporary:
+            root = Path(temporary)
+            coordinate = root / "coord.xyz"
+            coordinate.write_text("1\nINIT protocol test\nO 0 0 0\n")
+            name = f"dpipi_init_{os.getpid()}_{root.name}"
+            address = Path("/tmp") / f"ipi_{name}"
+            config = root / "config.json"
+            config.write_text(
+                json.dumps(
+                    {
+                        "verbose": verbose,
+                        "use_unix": True,
+                        "port": 31415,
+                        "host": name,
+                        "graph_file": str(model),
+                        "coord_file": str(coordinate),
+                        "atom_type": {"O": 0},
+                    }
+                )
+            )
+            process = None
+            server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            try:
+                server.bind(str(address))
+                server.listen(1)
+                server.settimeout(30)
+                process = subprocess.Popen(
+                    ["dp_ipi", str(config)],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                )
+                connection, _ = server.accept()
+                with connection:
+                    connection.settimeout(10)
+                    size = len(payload) if length is None else length
+                    connection.sendall(b"INIT        " + struct.pack("=ii", 0, size))
+                    connection.sendall(payload)
+                    if size == len(payload):
+                        connection.sendall(b"STATUS      ")
+                        response = b""
+                        while len(response) < 12:
+                            part = connection.recv(12 - len(response))
+                            self.assertTrue(part, "driver closed before STATUS reply")
+                            response += part
+                        self.assertEqual(response, b"READY       ")
+                        connection.sendall(b"EXIT        ")
+                    elif size > len(payload):
+                        connection.shutdown(socket.SHUT_WR)
+                stdout, stderr = process.communicate(timeout=30)
+                self.assertEqual(
+                    process.returncode,
+                    0 if size == len(payload) else (1 if size < 0 else 255),
+                    msg=f"stdout:\n{stdout}\nstderr:\n{stderr}",
+                )
+                if size > len(payload):
+                    self.assertIn("Error reading from socket", stderr)
+                    self.assertNotIn(payload.decode(), stdout)
+                if size < 0:
+                    self.assertIn(
+                        "dp_ipi: INIT payload length must be nonnegative.", stderr
+                    )
+            finally:
+                if process is not None and process.poll() is None:
+                    process.kill()
+                    process.wait()
+                server.close()
+                address.unlink(missing_ok=True)
+
+    def test_empty_initialization_preserves_next_header(self) -> None:
+        self.exchange(b"")
+
+    def test_initialization_consumes_replica_id_and_payload(self) -> None:
+        self.exchange(b'{"temperature": 300}')
+
+    def test_initialization_accepts_payload_larger_than_read_buffer(self) -> None:
+        self.exchange(b"x" * 4097)
+
+    def test_negative_initialization_length_is_rejected(self) -> None:
+        self.exchange(b"", length=-1)
+
+    def test_truncated_verbose_initialization_is_rejected(self) -> None:
+        self.exchange(b"partial-init-marker", length=64, verbose=True)

source/ipi/tests/test_driver_input_validation.py

@@ -50,10 +50,12 @@ def test_unknown_atom_name_is_rejected_before_model_loading(self) -> None:
             )
 
         self.assertEqual(result.returncode, 1, msg=result.stderr)
-        self.assertEqual(
-            result.stderr,
-            "dp_ipi: Unknown atom name 'Xx' in coordinate file: "
-            "no matching entry in atom_type.\n",
+        self.assertTrue(
+            result.stderr.endswith(
+                "dp_ipi: Unknown atom name 'Xx' in coordinate file: "
+                "no matching entry in atom_type.\n"
+            ),
+            msg=result.stderr,
         )
 
     def test_missing_atom_type_is_rejected_before_model_loading(self) -> None: