Skip to content

Commit adf3601

Browse files
chruffinschruffins
committed
Make QMP handshake cancellation race-safe
1 parent 12091b5 commit adf3601

4 files changed

Lines changed: 234 additions & 49 deletions

File tree

‎lib/hypervisor/qemu/pool.go‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,11 +68,6 @@ func GetOrCreateForTypeContext(ctx context.Context, socketPath string, hyperviso
6868
}
6969

7070
clientPool.Lock()
71-
if existing := clientPool.clients[socketPath]; existing != nil {
72-
clientPool.Unlock()
73-
closeClientAsync(client)
74-
return existing, nil
75-
}
7671
if creationLock.generation != generation {
7772
clientPool.Unlock()
7873
closeClientAsync(client)

‎lib/hypervisor/qemu/qemu_test.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package qemu
22

33
import (
44
"context"
5+
"encoding/json"
56
"errors"
67
"net"
78
"sync"
@@ -51,6 +52,51 @@ func awaitTestValue[T any](t *testing.T, ch <-chan T, timeoutMessage string) T {
5152
}
5253
}
5354

55+
func TestSocketMonitorHandshakeAndCommand(t *testing.T) {
56+
listener, err := net.Listen("unix", t.TempDir()+"/qemu.sock")
57+
require.NoError(t, err)
58+
defer listener.Close()
59+
60+
serverErr := make(chan error, 1)
61+
go func() {
62+
conn, err := listener.Accept()
63+
if err != nil {
64+
serverErr <- err
65+
return
66+
}
67+
defer conn.Close()
68+
if _, err := conn.Write([]byte("{\"QMP\":{}}\n")); err != nil {
69+
serverErr <- err
70+
return
71+
}
72+
decoder := json.NewDecoder(conn)
73+
var command map[string]any
74+
if err := decoder.Decode(&command); err != nil {
75+
serverErr <- err
76+
return
77+
}
78+
if _, err := conn.Write([]byte("{\"return\":{}}\n")); err != nil {
79+
serverErr <- err
80+
return
81+
}
82+
if err := decoder.Decode(&command); err != nil {
83+
serverErr <- err
84+
return
85+
}
86+
_, err = conn.Write([]byte("{\"return\":{\"status\":\"running\"}}\n"))
87+
serverErr <- err
88+
}()
89+
90+
monitor, err := newSocketMonitor(context.Background(), listener.Addr().String())
91+
require.NoError(t, err)
92+
t.Cleanup(func() { _ = monitor.Disconnect() })
93+
response, err := monitor.Run([]byte(`{"execute":"query-status"}` + "\n"))
94+
require.NoError(t, err)
95+
assert.JSONEq(t, `{"return":{"status":"running"}}`, string(response))
96+
require.NoError(t, monitor.Disconnect())
97+
require.NoError(t, <-serverErr)
98+
}
99+
54100
func TestNewClientContextBoundsQMPHandshake(t *testing.T) {
55101
listener, err := net.Listen("unix", t.TempDir()+"/qemu.sock")
56102
require.NoError(t, err)

‎lib/hypervisor/qemu/qmp.go‎

Lines changed: 5 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ const (
2424
type Client struct {
2525
domain *qemu.Domain
2626
raw *raw.Monitor
27-
mon *qmp.SocketMonitor
2827
closeOnce sync.Once
2928
closeErr error
3029
}
@@ -36,61 +35,23 @@ func NewClient(socketPath string) (*Client, error) {
3635
return NewClientContext(ctx, socketPath)
3736
}
3837

39-
// NewClientContext creates a QEMU client while bounding socket connection and
40-
// QMP handshake time by the caller's context.
38+
// NewClientContext creates a QEMU client using the caller's context for dial and handshake.
4139
func NewClientContext(ctx context.Context, socketPath string) (*Client, error) {
4240
if _, ok := ctx.Deadline(); !ok {
4341
var cancel context.CancelFunc
4442
ctx, cancel = context.WithTimeout(ctx, qmpConnectTimeout)
4543
defer cancel()
4644
}
47-
if err := ctx.Err(); err != nil {
48-
return nil, err
49-
}
50-
connectTimeout := qmpConnectTimeout
51-
if deadline, ok := ctx.Deadline(); ok {
52-
remaining := time.Until(deadline)
53-
if remaining <= 0 {
54-
if err := ctx.Err(); err != nil {
55-
return nil, err
56-
}
57-
return nil, context.DeadlineExceeded
58-
}
59-
if remaining < connectTimeout {
60-
connectTimeout = remaining
61-
}
62-
}
63-
mon, err := qmp.NewSocketMonitor("unix", socketPath, connectTimeout)
45+
mon, err := newSocketMonitor(ctx, socketPath)
6446
if err != nil {
65-
return nil, fmt.Errorf("create socket monitor: %w", err)
66-
}
67-
68-
connectCh := make(chan error, 1)
69-
go func() { connectCh <- mon.Connect() }()
70-
select {
71-
case err := <-connectCh:
72-
if err != nil {
73-
_ = mon.Disconnect()
74-
return nil, fmt.Errorf("connect to qmp: %w", err)
75-
}
76-
case <-ctx.Done():
77-
_ = mon.Disconnect()
78-
<-connectCh
79-
_ = mon.Disconnect()
80-
return nil, fmt.Errorf("connect to qmp: %w", ctx.Err())
47+
return nil, err
8148
}
82-
8349
domain, err := qemu.NewDomain(mon, "vm")
8450
if err != nil {
85-
mon.Disconnect()
51+
_ = mon.Disconnect()
8652
return nil, fmt.Errorf("create domain: %w", err)
8753
}
88-
89-
return &Client{
90-
domain: domain,
91-
raw: raw.NewMonitor(mon),
92-
mon: mon,
93-
}, nil
54+
return &Client{domain: domain, raw: raw.NewMonitor(mon)}, nil
9455
}
9556

9657
// Close disconnects from the QMP socket.

‎lib/hypervisor/qemu/qmp_monitor.go‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
package qemu
2+
3+
import (
4+
"bufio"
5+
"context"
6+
"encoding/json"
7+
"errors"
8+
"fmt"
9+
"io"
10+
"net"
11+
"sync"
12+
"sync/atomic"
13+
14+
"github.com/digitalocean/go-qemu/qmp"
15+
)
16+
17+
type socketMonitor struct {
18+
conn net.Conn
19+
connectDone chan struct{}
20+
disconnect chan struct{}
21+
disconnect1 sync.Once
22+
responses chan monitorResponse
23+
events chan qmp.Event
24+
listenDone chan struct{}
25+
listeners atomic.Int32
26+
commandMu sync.Mutex
27+
closeErr error
28+
}
29+
30+
type monitorResponse struct {
31+
data []byte
32+
err error
33+
}
34+
35+
func newSocketMonitor(ctx context.Context, socketPath string) (*socketMonitor, error) {
36+
conn, err := (&net.Dialer{Timeout: qmpConnectTimeout}).DialContext(ctx, "unix", socketPath)
37+
if err != nil {
38+
return nil, fmt.Errorf("create socket monitor: %w", err)
39+
}
40+
m := &socketMonitor{
41+
conn: conn,
42+
connectDone: make(chan struct{}),
43+
disconnect: make(chan struct{}),
44+
}
45+
connected := make(chan error, 1)
46+
go func() { connected <- m.Connect() }()
47+
select {
48+
case err := <-connected:
49+
if err != nil {
50+
_ = m.Disconnect()
51+
return nil, fmt.Errorf("connect to qmp: %w", err)
52+
}
53+
case <-ctx.Done():
54+
_ = m.Disconnect()
55+
<-connected
56+
return nil, fmt.Errorf("connect to qmp: %w", ctx.Err())
57+
}
58+
return m, nil
59+
}
60+
61+
func (m *socketMonitor) Connect() error {
62+
defer close(m.connectDone)
63+
decoder := json.NewDecoder(m.conn)
64+
if err := decoder.Decode(new(any)); err != nil {
65+
return err
66+
}
67+
if err := json.NewEncoder(m.conn).Encode(qmp.Command{Execute: "qmp_capabilities"}); err != nil {
68+
return err
69+
}
70+
var response json.RawMessage
71+
if err := decoder.Decode(&response); err != nil {
72+
return err
73+
}
74+
if err := qmpResponseError(response); err != nil {
75+
return err
76+
}
77+
78+
m.responses = make(chan monitorResponse, 1)
79+
m.events = make(chan qmp.Event)
80+
m.listenDone = make(chan struct{})
81+
go m.listen()
82+
return nil
83+
}
84+
85+
func (m *socketMonitor) Disconnect() error {
86+
m.disconnect1.Do(func() {
87+
close(m.disconnect)
88+
m.closeErr = m.conn.Close()
89+
})
90+
<-m.connectDone
91+
if m.listenDone != nil {
92+
<-m.listenDone
93+
}
94+
return m.closeErr
95+
}
96+
97+
func (m *socketMonitor) Run(command []byte) ([]byte, error) {
98+
m.commandMu.Lock()
99+
defer m.commandMu.Unlock()
100+
responses := m.responses
101+
if responses == nil {
102+
return nil, errors.New("QMP monitor is not connected")
103+
}
104+
if _, err := m.conn.Write(command); err != nil {
105+
return nil, err
106+
}
107+
select {
108+
case response, ok := <-responses:
109+
if !ok {
110+
return nil, io.EOF
111+
}
112+
if response.err != nil {
113+
return nil, response.err
114+
}
115+
if err := qmpResponseError(response.data); err != nil {
116+
return nil, err
117+
}
118+
return response.data, nil
119+
case <-m.disconnect:
120+
return nil, net.ErrClosed
121+
}
122+
}
123+
124+
func (m *socketMonitor) Events(context.Context) (<-chan qmp.Event, error) {
125+
events := m.events
126+
if events == nil {
127+
return nil, errors.New("QMP monitor is not connected")
128+
}
129+
m.listeners.Add(1)
130+
return events, nil
131+
}
132+
133+
func qmpResponseError(data []byte) error {
134+
var response struct {
135+
Error *struct {
136+
Desc string `json:"desc"`
137+
} `json:"error"`
138+
}
139+
if err := json.Unmarshal(data, &response); err != nil {
140+
return err
141+
}
142+
if response.Error != nil && response.Error.Desc != "" {
143+
return errors.New(response.Error.Desc)
144+
}
145+
return nil
146+
}
147+
148+
func (m *socketMonitor) listen() {
149+
defer close(m.listenDone)
150+
defer close(m.events)
151+
defer close(m.responses)
152+
scanner := bufio.NewScanner(m.conn)
153+
for scanner.Scan() {
154+
data := append([]byte(nil), scanner.Bytes()...)
155+
var event qmp.Event
156+
if json.Unmarshal(data, &event) != nil {
157+
continue
158+
}
159+
if event.Event != "" {
160+
if m.listeners.Load() > 0 {
161+
select {
162+
case m.events <- event:
163+
case <-m.disconnect:
164+
return
165+
}
166+
}
167+
continue
168+
}
169+
select {
170+
case m.responses <- monitorResponse{data: data}:
171+
case <-m.disconnect:
172+
return
173+
}
174+
}
175+
if err := scanner.Err(); err != nil {
176+
select {
177+
case m.responses <- monitorResponse{err: err}:
178+
case <-m.disconnect:
179+
}
180+
}
181+
}
182+
183+
var _ qmp.Monitor = (*socketMonitor)(nil)

0 commit comments

Comments
 (0)