@samitouri / QOSAMI-WSL
Code Commits Issues Pulls Sessions CI/CD
master
master
h 538 lines 15.3 KB
Raw
1 // Copyright (C) Microsoft Corporation. All rights reserved.
2
3 #pragma once
4
5 #include <concurrent_queue.h>
6 #include <deque>
7 #include <list>
8
9 #define LX_RELAY_BUFFER_SIZE 0x1000
10
11 namespace wsl::windows::common::io {
12
13 // Returns the current file pointer for disk handles, and a zero offset for every other handle type.
14 inline LARGE_INTEGER InitializeFileOffset(HANDLE File)
15 {
16 LARGE_INTEGER Offset{};
17 if (GetFileType(File) == FILE_TYPE_DISK)
18 {
19 LOG_IF_WIN32_BOOL_FALSE(SetFilePointerEx(File, {}, &Offset, FILE_CURRENT));
20 }
21
22 return Offset;
23 }
24
25 enum class IOHandleStatus
26 {
27 Standby,
28 Pending,
29 Completed,
30 // A persistent handle with no work to do. It stays registered with the MultiHandleWait loop but is not
31 // waited on; it returns to Standby once more data is pushed into it.
32 Idle
33 };
34
35 struct HandleWrapper
36 {
37 NON_COPYABLE(HandleWrapper)
38
39 HandleWrapper() = default;
40 HandleWrapper(HandleWrapper&& other) noexcept;
41 HandleWrapper& operator=(HandleWrapper&& other) noexcept;
42 HandleWrapper(wil::unique_handle&& handle, std::function<void()>&& OnClose = []() {});
43 HandleWrapper(wil::unique_socket&& handle, std::function<void()>&& OnClose = []() {});
44 HandleWrapper(wil::shared_handle handle, std::function<void()>&& OnClose = []() {});
45 HandleWrapper(wil::shared_socket handle, std::function<void()>&& OnClose = []() {});
46 HandleWrapper(wil::unique_event&& handle, std::function<void()>&& OnClose = []() {});
47 HandleWrapper(SOCKET handle, std::function<void()>&& OnClose = []() {});
48 HandleWrapper(HANDLE handle, std::function<void()>&& OnClose = []() {});
49 HandleWrapper(wil::unique_hfile&& handle, std::function<void()>&& OnClose = []() {});
50 ~HandleWrapper();
51
52 HANDLE Get() const;
53 bool IsValid() const;
54 void Reset();
55
56 private:
57 HANDLE Handle{};
58 std::variant<wil::unique_handle, wil::unique_socket, wil::shared_handle, wil::shared_socket> OwnedHandle;
59 std::function<void()> OnClose;
60 };
61
62 // A buffer that may either own its underlying storage (constructed from a size, allocating an
63 // internal std::vector<char>) or borrow it from a caller-provided gsl::span<gsl::byte>.
64 class BufferWrapper
65 {
66 public:
67 DEFAULT_MOVABLE(BufferWrapper);
68 NON_COPYABLE(BufferWrapper);
69
70 explicit BufferWrapper(size_t size);
71 explicit BufferWrapper(gsl::span<gsl::byte> span);
72
73 bool Owned() const noexcept;
74 void Resize(size_t size);
75 void Append(gsl::span<char> Span);
76 void Consume(size_t bytes) noexcept;
77 gsl::span<gsl::byte> Span() noexcept;
78 size_t Size() const noexcept;
79
80 private:
81 std::optional<std::vector<char>> m_owned;
82 gsl::span<gsl::byte> m_unowned;
83 };
84
85 class OverlappedIOHandle
86 {
87 public:
88 NON_COPYABLE(OverlappedIOHandle)
89 NON_MOVABLE(OverlappedIOHandle)
90
91 OverlappedIOHandle() = default;
92 virtual ~OverlappedIOHandle() = default;
93 virtual void Schedule() = 0;
94 virtual void Collect() = 0;
95 virtual HANDLE GetHandle() const = 0;
96 IOHandleStatus GetState() const;
97
98 protected:
99 IOHandleStatus State = IOHandleStatus::Standby;
100 };
101
102 class EventHandle : public OverlappedIOHandle
103 {
104 public:
105 NON_COPYABLE(EventHandle)
106 NON_MOVABLE(EventHandle)
107
108 EventHandle(HandleWrapper&& Handle, std::function<void()>&& OnSignalled = []() {});
109 void Schedule() override;
110 void Collect() override;
111 HANDLE GetHandle() const override;
112
113 private:
114 HandleWrapper Handle;
115 std::function<void()> OnSignalled;
116 };
117
118 class ReadHandle : public OverlappedIOHandle
119 {
120 public:
121 NON_COPYABLE(ReadHandle);
122 NON_MOVABLE(ReadHandle);
123
124 ReadHandle(HandleWrapper&& MovedHandle, std::function<void(const gsl::span<char>& Buffer)>&& OnRead);
125 virtual ~ReadHandle();
126
127 void Schedule() override;
128 void Collect() override;
129 HANDLE GetHandle() const override;
130
131 protected:
132 HandleWrapper Handle;
133 OVERLAPPED Overlapped{};
134
135 private:
136 std::function<void(const gsl::span<char>& Buffer)> OnRead;
137 wil::unique_event Event{wil::EventOptions::ManualReset};
138 BufferWrapper Buffer{LX_RELAY_BUFFER_SIZE};
139 LARGE_INTEGER Offset{};
140 };
141
142 // A ReadHandle for a server named pipe. It waits for a client to connect (ConnectNamedPipe) before
143 // reading, so it can be scheduled directly with a freshly created server pipe. The connect reuses the
144 // base read handle's overlapped and event, so the base destructor cancels a pending connect correctly.
145 class ReadNamedPipe : public ReadHandle
146 {
147 public:
148 NON_COPYABLE(ReadNamedPipe);
149 NON_MOVABLE(ReadNamedPipe);
150
151 ReadNamedPipe(HandleWrapper&& Pipe, std::function<void(const gsl::span<char>& Buffer)>&& OnRead);
152
153 void Schedule() override;
154 void Collect() override;
155
156 private:
157 bool m_connected = false;
158 };
159
160 class AcceptHandle : public OverlappedIOHandle
161 {
162 public:
163 NON_COPYABLE(AcceptHandle)
164 NON_MOVABLE(AcceptHandle)
165
166 AcceptHandle(HandleWrapper&& ListenSocket, bool AcceptOnce, std::function<void(wil::unique_socket&&)>&& OnAccepted);
167 ~AcceptHandle();
168
169 void Schedule() override;
170 void Collect() override;
171 HANDLE GetHandle() const override;
172
173 private:
174 void CreateAcceptSocket();
175 void OnComplete();
176
177 HandleWrapper ListenSocket;
178 wil::unique_socket AcceptedSocket;
179 int AddressFamily{};
180 int SocketType{};
181 int Protocol{};
182 bool AcceptOnce{};
183 wil::unique_event Event{wil::EventOptions::ManualReset};
184 OVERLAPPED Overlapped{};
185 std::function<void(wil::unique_socket&&)> OnAccepted;
186 char AcceptBuffer[2 * sizeof(SOCKADDR_STORAGE)];
187 };
188
189 class LineBasedReadHandle : public ReadHandle
190 {
191 public:
192 NON_COPYABLE(LineBasedReadHandle);
193 NON_MOVABLE(LineBasedReadHandle);
194
195 LineBasedReadHandle(HandleWrapper&& Handle, std::function<void(const gsl::span<char>& Buffer)>&& OnLine, bool Crlf);
196 ~LineBasedReadHandle();
197
198 private:
199 void OnRead(const gsl::span<char>& Buffer);
200
201 std::function<void(const gsl::span<char>& Buffer)> OnLine;
202 std::string PendingBuffer;
203 bool Crlf{};
204 };
205
206 class HTTPChunkBasedReadHandle : public ReadHandle
207 {
208 public:
209 NON_COPYABLE(HTTPChunkBasedReadHandle);
210 NON_MOVABLE(HTTPChunkBasedReadHandle);
211
212 HTTPChunkBasedReadHandle(HandleWrapper&& Handler, std::function<void(const gsl::span<char>& Buffer)>&& OnChunk);
213 ~HTTPChunkBasedReadHandle();
214
215 void OnRead(const gsl::span<char>& Line);
216
217 private:
218 std::function<void(const gsl::span<char>& Buffer)> OnChunk;
219 std::string PendingBuffer;
220 uint64_t PendingChunkSize = 0;
221 bool ExpectHeader = true;
222 };
223
224 class ReadSocketMessageHandle : public OverlappedIOHandle
225 {
226 public:
227 NON_COPYABLE(ReadSocketMessageHandle);
228 NON_MOVABLE(ReadSocketMessageHandle);
229
230 ReadSocketMessageHandle(
231 HandleWrapper&& Socket,
232 std::vector<gsl::byte>& Buffer,
233 std::vector<gsl::byte>& PendingBytes,
234 std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage);
235 ~ReadSocketMessageHandle();
236
237 void Schedule() override;
238 void Collect() override;
239 HANDLE GetHandle() const override;
240
241 private:
242 void ScheduleRecv();
243 void ProcessRecvResult(DWORD BytesRead);
244 bool ProcessChunk();
245
246 HandleWrapper Socket;
247 std::vector<gsl::byte>& Buffer;
248 std::vector<gsl::byte>& PendingBytes;
249 std::function<void(const gsl::span<gsl::byte>& Message)> OnMessage;
250 wil::unique_event Event{wil::EventOptions::ManualReset};
251 OVERLAPPED Overlapped{};
252 bool ReadingHeader = true;
253 size_t BytesRemaining = sizeof(MESSAGE_HEADER);
254 size_t CurrentOffset = 0;
255 };
256
257 class ReadConsoleHandle : public OverlappedIOHandle
258 {
259 public:
260 NON_COPYABLE(ReadConsoleHandle);
261 NON_MOVABLE(ReadConsoleHandle);
262
263 ReadConsoleHandle(
264 HandleWrapper&& Console,
265 std::function<void(const gsl::span<char>& Buffer)>&& OnRead,
266 std::function<void()>&& UpdateTerminalSize = []() {},
267 std::vector<char> DetachSequence = {},
268 std::function<void()>&& OnDetach = []() {});
269
270 void Schedule() override;
271 void Collect() override;
272 HANDLE GetHandle() const override;
273
274 private:
275 HandleWrapper Console;
276 std::function<void(const gsl::span<char>& Buffer)> OnRead;
277 std::function<void()> UpdateTerminalSize;
278 std::vector<char> DetachSequence;
279 std::function<void()> OnDetach;
280 std::deque<char> CurrentSequence;
281 };
282
283 class WriteHandle : public OverlappedIOHandle
284 {
285 public:
286 NON_COPYABLE(WriteHandle);
287 NON_MOVABLE(WriteHandle);
288
289 WriteHandle(HandleWrapper&& Handle, const std::vector<char>& Source = {}, bool CompleteOnDrained = true);
290 WriteHandle(HandleWrapper&& Handle, gsl::span<gsl::byte> Source);
291 ~WriteHandle();
292 void Schedule() override;
293 void Collect() override;
294 HANDLE GetHandle() const override;
295 void Push(const gsl::span<char>& Buffer);
296
297 // Controls whether the writer completes (and is removed from the loop) or becomes Idle once its buffer drains.
298 void SetCompleteOnDrained(bool CompleteOnDrained);
299
300 // Returns the number of bytes that have been queued for writing but not yet written to the handle.
301 size_t PendingBytes() const;
302
303 private:
304 // Returns the state to adopt once the active buffer drains: Completed for one-shot writers, or Idle/Standby
305 // for reusable writers depending on whether more data is queued.
306 IOHandleStatus DrainedState() const;
307
308 HandleWrapper Handle;
309 wil::unique_event Event{wil::EventOptions::ManualReset};
310 OVERLAPPED Overlapped{};
311 BufferWrapper Buffer;
312 LARGE_INTEGER Offset{};
313 bool CompleteOnDrained = true;
314 std::vector<char> Pending;
315 };
316
317 // A persistent writer for a named pipe that transparently handles the server-side connection lifecycle. Data
318 // pushed via Push() is written to the pipe by the IO loop. 'Connected' indicates the pipe already has a connected
319 // peer.
320 class WriteNamedPipe : public OverlappedIOHandle
321 {
322 public:
323 NON_COPYABLE(WriteNamedPipe);
324 NON_MOVABLE(WriteNamedPipe);
325
326 WriteNamedPipe(HandleWrapper&& Pipe, bool Reconnect, bool Connected);
327 ~WriteNamedPipe();
328 void Schedule() override;
329 void Collect() override;
330 HANDLE GetHandle() const override;
331 void Push(const gsl::span<char>& Buffer);
332
333 // Returns the number of bytes that have been queued for writing but not yet written to the pipe.
334 size_t PendingBytes() const;
335
336 private:
337 // Drops the current client and arms a fresh connection so the next Schedule() reconnects before writing.
338 void Reconnect();
339
340 HandleWrapper Pipe;
341 std::optional<WriteHandle> Write;
342 wil::unique_event ConnectEvent{wil::EventOptions::ManualReset};
343 OVERLAPPED ConnectOverlapped{};
344 bool ReconnectOnFailure = false;
345 bool NeedConnect = false;
346 bool Connecting = false;
347 };
348
349 template <typename TRead = ReadHandle>
350 class RelayHandle : public OverlappedIOHandle
351 {
352 public:
353 NON_COPYABLE(RelayHandle);
354 NON_MOVABLE(RelayHandle);
355
356 template <typename... TArgs>
357 RelayHandle(HandleWrapper&& Input, HandleWrapper&& Output, TArgs&&... InputArgs) :
358 Read(
359 std::move(Input), [this](const gsl::span<char>& Buffer) { return OnRead(Buffer); }, std::forward<TArgs>(InputArgs)...),
360 Write(std::move(Output), {}, false)
361 {
362 }
363
364 void Schedule() override
365 {
366 WI_ASSERT(State == IOHandleStatus::Standby);
367
368 // If the Buffer is empty, then we're reading.
369 if (PendingBuffer.empty())
370 {
371 if (Read.GetState() == IOHandleStatus::Completed)
372 {
373 // If all reading is complete, flush any pending writes before transitioning to Completed.
374 Write.SetCompleteOnDrained(true);
375
376 if (Write.PendingBytes() > 0)
377 {
378 Write.Schedule();
379 State = Write.GetState();
380 }
381 else
382 {
383 State = IOHandleStatus::Completed;
384 }
385
386 return;
387 }
388
389 Read.Schedule();
390
391 // If the read is pending, update to 'Pending'
392 if (Read.GetState() == IOHandleStatus::Pending)
393 {
394 State = IOHandleStatus::Pending;
395 }
396 }
397 else
398 {
399 Write.Push(PendingBuffer);
400 PendingBuffer.clear();
401
402 Write.Schedule();
403
404 if (Write.GetState() == IOHandleStatus::Pending)
405 {
406 // The write is pending, update to 'Pending'
407 State = IOHandleStatus::Pending;
408 }
409 }
410 }
411
412 void Collect() override
413 {
414 WI_ASSERT(State == IOHandleStatus::Pending);
415
416 // Transition back to standby
417 State = IOHandleStatus::Standby;
418
419 if (Read.GetState() == IOHandleStatus::Pending)
420 {
421 Read.Collect();
422 }
423 else
424 {
425 WI_ASSERT(Write.GetState() == IOHandleStatus::Pending);
426 Write.Collect();
427 }
428 }
429
430 HANDLE GetHandle() const override
431 {
432 if (Read.GetState() == IOHandleStatus::Pending)
433 {
434 return Read.GetHandle();
435 }
436 else
437 {
438 WI_ASSERT(Write.GetState() == IOHandleStatus::Pending);
439 return Write.GetHandle();
440 }
441 }
442
443 private:
444 void OnRead(const gsl::span<char>& Content)
445 {
446 PendingBuffer.insert(PendingBuffer.end(), Content.begin(), Content.end());
447 }
448
449 TRead Read;
450 WriteHandle Write;
451 std::vector<char> PendingBuffer;
452 };
453
454 class DockerIORelayHandle : public OverlappedIOHandle
455 {
456 public:
457 NON_COPYABLE(DockerIORelayHandle);
458 NON_MOVABLE(DockerIORelayHandle);
459
460 enum class Format
461 {
462 Raw,
463 HttpChunked
464 };
465
466 DockerIORelayHandle(HandleWrapper&& Input, HandleWrapper&& Stdout, HandleWrapper&& Stderr, Format ReadFormat);
467 void Schedule() override;
468 void Collect() override;
469 HANDLE GetHandle() const override;
470
471 #pragma pack(push, 1)
472 struct MultiplexedHeader
473 {
474 uint8_t Fd;
475 char Zeroes[3];
476 uint32_t Length;
477 };
478 #pragma pack(pop)
479
480 static_assert(sizeof(MultiplexedHeader) == 8);
481
482 private:
483 void OnRead(const gsl::span<char>& Buffer);
484 void ProcessNextHeader();
485
486 std::unique_ptr<OverlappedIOHandle> Read;
487 WriteHandle WriteStdout;
488 WriteHandle WriteStderr;
489 std::vector<char> PendingBuffer;
490 WriteHandle* ActiveHandle = nullptr;
491 size_t RemainingBytes = 0;
492 };
493 class MultiHandleWait
494 {
495 public:
496 NON_COPYABLE(MultiHandleWait);
497
498 using OnError = std::function<void()>;
499
500 enum Flags
501 {
502 None = 0,
503 CancelOnCompleted = 1,
504 IgnoreErrors = 2,
505 NeedNotComplete = 4,
506 };
507
508 MultiHandleWait() = default;
509 MultiHandleWait(MultiHandleWait&&) noexcept;
510 MultiHandleWait& operator=(MultiHandleWait&&) noexcept;
511
512 void AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags = Flags::None, OnError&& onError = []() { throw; });
513 bool Run(std::optional<std::chrono::milliseconds> Timeout);
514 void Cancel();
515
516 private:
517 struct Entry
518 {
519 Flags HandleFlags{};
520 std::unique_ptr<OverlappedIOHandle> Handle;
521 MultiHandleWait* Self;
522 OnError ErrorCallback;
523 };
524
525 static void NTAPI WaitCallback(PVOID Context, BOOLEAN TimerOrWaitFired);
526
527 concurrency::concurrent_queue<Entry*> m_signaledHandles;
528 wil::unique_event m_handleSignaledEvent{wil::EventOptions::ManualReset};
529
530 // N.B. A std::list is used (rather than a vector) so handles can be added from a callback while Run() is
531 // iterating m_handles without invalidating the loop's iterator.
532 std::list<std::unique_ptr<Entry>> m_handles;
533 bool m_cancel = false;
534 };
535
536 DEFINE_ENUM_FLAG_OPERATORS(MultiHandleWait::Flags);
537
538 } // namespace wsl::windows::common::io