Skip to main content

max / alloy

6.6 KB · 180 lines History Blame Raw
1 #!/usr/bin/env python3
2 """Talk to qemu's QMP socket: run a command, or wait for an event.
3
4 `vm.py` drives the human monitor, which is the right interface for sendkey and
5 screendump. This is the other one, and it exists for the thing HMP cannot do:
6 say when something happened inside the guest.
7
8 A scripted first-boot test has to know when the guest rebooted. HMP has no
9 answer — `info status` reports running either side of a reset, and a script
10 left watching the clock or the qcow2's size is guessing. QMP emits `RESET` as
11 an event, so "the machine rebooted" becomes a thing to wait for rather than a
12 thing to infer. When it does not arrive, the timeout is itself the finding.
13
14 `set_link` is here for the same reason it is not in vm.py: taking the guest's
15 route away from outside is the only way to do it that does not depend on the
16 guest cooperating. Doing it from inside means `ip link set enp0s2 down` over
17 the ssh session that command arrived on, which kills the session in the middle
18 of the thing being measured — the hand run on 2026-08-25 needed a detached
19 `setsid` script that brought the interface back up at the end, because there is
20 no other way home. From out here the link is a property of the device.
21
22 Use:
23
24 qmp.py cmd query-status
25 qmp.py cmd set_link '{"name": "n0", "up": false}'
26 qmp.py wait RESET 900
27 qmp.py wait RESET 900 '{"guest": true}'
28 qmp.py cmd quit
29
30 `wait` exits 0 when the event arrives and 1 when it does not, so a scenario
31 script reads the verdict off the exit status.
32 """
33
34 import json
35 import os
36 import socket
37 import sys
38 import time
39
40 SCRATCH = os.environ.get(
41 "VM_STATE", os.path.join(os.path.dirname(os.path.abspath(__file__)), "state")
42 )
43 SOCK = os.path.join(SCRATCH, "qmp.sock")
44
45
46 class QMP:
47 def __init__(self, path=SOCK, connect_timeout=30.0):
48 # qemu creates the socket as it starts, so a scenario that launches
49 # qemu and connects immediately is racing it. Retrying here rather than
50 # sleeping at every call site.
51 end = time.time() + connect_timeout
52 while True:
53 try:
54 self.s = socket.socket(socket.AF_UNIX)
55 self.s.connect(path)
56 break
57 except (FileNotFoundError, ConnectionRefusedError):
58 if time.time() >= end:
59 raise
60 time.sleep(0.2)
61 # A raw buffer rather than socket.makefile. A buffered reader that
62 # raises socket.timeout can leave a partial line in a buffer nobody
63 # can get at again, which turns one slow read into a stream that never
64 # parses. Owning the buffer keeps a timeout a timeout.
65 self.buf = b""
66 self.events = []
67 self._read(timeout=30.0) # the greeting
68 self.execute("qmp_capabilities")
69
70 def _line(self, timeout):
71 """One newline-terminated message from the socket."""
72 while b"\n" not in self.buf:
73 self.s.settimeout(timeout)
74 chunk = self.s.recv(65536)
75 if not chunk:
76 raise EOFError("qemu closed the QMP socket")
77 self.buf += chunk
78 line, self.buf = self.buf.split(b"\n", 1)
79 return line
80
81 def _read(self, timeout=None):
82 """One JSON reply, with events set aside rather than returned.
83
84 Events interleave with replies, so a reader that returned the first
85 message it saw would hand a caller someone else's RESET."""
86 while True:
87 msg = json.loads(self._line(timeout))
88 if "event" in msg:
89 self.events.append(msg)
90 continue
91 return msg
92
93 def execute(self, command, arguments=None, timeout=30.0):
94 payload = {"execute": command}
95 if arguments:
96 payload["arguments"] = arguments
97 self.s.sendall((json.dumps(payload) + "\n").encode())
98 # `quit` is answered and then the socket goes away, sometimes in that
99 # order and sometimes not, so losing the reply to it is not a failure.
100 try:
101 reply = self._read(timeout)
102 except EOFError:
103 if command == "quit":
104 return {}
105 raise
106 if "error" in reply:
107 raise RuntimeError("%s: %s" % (command, reply["error"].get("desc", reply)))
108 return reply.get("return", {})
109
110 def wait_event(self, name, timeout, data=None):
111 """The next `name` event, or None if the timeout runs out first.
112
113 `data` narrows it to events whose payload carries those keys and
114 values, which is how RESET becomes a question about the guest rather
115 than about the machine: a reset the host asked for and a reboot the
116 guest performed are the same event with a different `reason`, and only
117 one of them is anything being measured.
118
119 Events already buffered count: one that arrived while a command was in
120 flight is the same event, and dropping it would make this depend on
121 when it was called."""
122 end = time.time() + timeout
123
124 def matches(e):
125 if e.get("event") != name:
126 return False
127 payload = e.get("data", {})
128 return all(payload.get(k) == v for k, v in (data or {}).items())
129
130 while True:
131 for i, e in enumerate(self.events):
132 if matches(e):
133 return self.events.pop(i)
134 remaining = end - time.time()
135 if remaining <= 0:
136 return None
137 try:
138 # Nothing but events is expected here, and _read files those
139 # away itself; a reply arriving would be someone else's.
140 self._read(min(remaining, 1.0))
141 except socket.timeout:
142 continue
143 except EOFError:
144 return None
145
146 def close(self):
147 try:
148 self.s.close()
149 except OSError:
150 pass
151
152
153 def main():
154 if len(sys.argv) < 2:
155 raise SystemExit(__doc__)
156 verb = sys.argv[1]
157 q = QMP()
158 if verb == "cmd":
159 command = sys.argv[2]
160 arguments = json.loads(sys.argv[3]) if len(sys.argv) > 3 else None
161 print(json.dumps(q.execute(command, arguments)))
162 return
163 if verb == "wait":
164 name = sys.argv[2]
165 seconds = float(sys.argv[3]) if len(sys.argv) > 3 else 300.0
166 data = json.loads(sys.argv[4]) if len(sys.argv) > 4 else None
167 event = q.wait_event(name, seconds, data)
168 if event is None:
169 print("no %s%s in %gs"
170 % (name, " matching %s" % json.dumps(data) if data else "", seconds),
171 file=sys.stderr)
172 raise SystemExit(1)
173 print(json.dumps(event))
174 return
175 raise SystemExit(__doc__)
176
177
178 if __name__ == "__main__":
179 main()
180