| 1 | // Copyright (C) Microsoft Corporation. All rights reserved. |
| 2 | |
| 3 | #include "precomp.h" |
| 4 | #include "HandleIO.h" |
| 5 | #pragma hdrstop |
| 6 | |
| 7 | using wsl::windows::common::io::AcceptHandle; |
| 8 | using wsl::windows::common::io::BufferWrapper; |
| 9 | using wsl::windows::common::io::DockerIORelayHandle; |
| 10 | using wsl::windows::common::io::EventHandle; |
| 11 | using wsl::windows::common::io::HandleWrapper; |
| 12 | using wsl::windows::common::io::HTTPChunkBasedReadHandle; |
| 13 | using wsl::windows::common::io::InitializeFileOffset; |
| 14 | using wsl::windows::common::io::IOHandleStatus; |
| 15 | using wsl::windows::common::io::LineBasedReadHandle; |
| 16 | using wsl::windows::common::io::MultiHandleWait; |
| 17 | using wsl::windows::common::io::OverlappedIOHandle; |
| 18 | using wsl::windows::common::io::ReadHandle; |
| 19 | using wsl::windows::common::io::ReadNamedPipe; |
| 20 | using wsl::windows::common::io::ReadSocketMessageHandle; |
| 21 | using wsl::windows::common::io::WriteHandle; |
| 22 | using wsl::windows::common::io::WriteNamedPipe; |
| 23 | |
| 24 | namespace { |
| 25 | |
| 26 | DWORD CancelPendingIo(auto Handle, OVERLAPPED& Overlapped) |
| 27 | { |
| 28 | DWORD bytesTransferred{}; |
| 29 | if (CancelIoEx((HANDLE)Handle, &Overlapped) || GetLastError() == ERROR_NOT_FOUND) |
| 30 | { |
| 31 | if constexpr (std::is_same_v<decltype(Handle), SOCKET>) |
| 32 | { |
| 33 | DWORD flagsReturned{}; |
| 34 | if (!WSAGetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true, &flagsReturned)) |
| 35 | { |
| 36 | auto error = WSAGetLastError(); |
| 37 | LOG_LAST_ERROR_IF(error != WSAECONNABORTED && error != WSA_OPERATION_ABORTED && error != WSAECONNRESET); |
| 38 | } |
| 39 | } |
| 40 | else |
| 41 | { |
| 42 | static_assert(std::is_same_v<decltype(Handle), HANDLE>); |
| 43 | if (!GetOverlappedResult(Handle, &Overlapped, &bytesTransferred, true)) |
| 44 | { |
| 45 | auto error = GetLastError(); |
| 46 | LOG_LAST_ERROR_IF(error != ERROR_CONNECTION_ABORTED && error != ERROR_OPERATION_ABORTED); |
| 47 | } |
| 48 | } |
| 49 | } |
| 50 | else |
| 51 | { |
| 52 | LOG_LAST_ERROR_MSG("Unexpected error while cancelling IO on handle: 0x%p", (void*)Handle); |
| 53 | } |
| 54 | |
| 55 | return bytesTransferred; |
| 56 | } |
| 57 | |
| 58 | inline void UnregisterWait(HANDLE waitHandle) noexcept |
| 59 | { |
| 60 | // INVALID_HANDLE_VALUE makes UnregisterWaitEx block until any in-flight wait callback returns. |
| 61 | LOG_LAST_ERROR_IF(!UnregisterWaitEx(waitHandle, INVALID_HANDLE_VALUE)); |
| 62 | } |
| 63 | |
| 64 | using unique_registered_wait = wil::unique_any_handle_null<decltype(&UnregisterWait), &UnregisterWait>; |
| 65 | |
| 66 | #define TTY_ALT_NUMPAD_VK_MENU (0x12) |
| 67 | #define TTY_ESCAPE_CHARACTER (L'\x1b') |
| 68 | #define TTY_INPUT_EVENT_BUFFER_SIZE (16) |
| 69 | #define TTY_UTF8_TRANSLATION_BUFFER_SIZE (4 * TTY_INPUT_EVENT_BUFFER_SIZE) |
| 70 | |
| 71 | BOOL IsActionableKey(_In_ PKEY_EVENT_RECORD KeyEvent) |
| 72 | { |
| 73 | // |
| 74 | // This is a bit complicated to discern. |
| 75 | // |
| 76 | // 1. Our first check is that we only want structures that |
| 77 | // represent at least one key press. If we have 0, then we don't |
| 78 | // need to bother. If we have >1, we'll send the key through |
| 79 | // that many times into the pipe. |
| 80 | // 2. Our second check is where it gets confusing. |
| 81 | // a. Characters that are non-null get an automatic pass. Copy |
| 82 | // them through to the pipe. |
| 83 | // b. Null characters need further scrutiny. We generally do not |
| 84 | // pass nulls through EXCEPT if they're sourced from the |
| 85 | // virtual terminal engine (or another application living |
| 86 | // above our layer). If they're sourced by a non-keyboard |
| 87 | // source, they'll have no scan code (since they didn't come |
| 88 | // from a keyboard). But that rule has an exception too: |
| 89 | // "Enhanced keys" from above the standard range of scan |
| 90 | // codes will return 0 also with a special flag set that says |
| 91 | // they're an enhanced key. That means the desired behavior |
| 92 | // is: |
| 93 | // Scan Code = 0, ENHANCED_KEY = 0 |
| 94 | // -> This came from the VT engine or another app |
| 95 | // above our layer. |
| 96 | // Scan Code = 0, ENHANCED_KEY = 1 |
| 97 | // -> This came from the keyboard, but is a special |
| 98 | // key like 'Volume Up' that wasn't generally a |
| 99 | // part of historic (pre-1990s) keyboards. |
| 100 | // Scan Code = <anything else> |
| 101 | // -> This came from a keyboard directly. |
| 102 | // |
| 103 | |
| 104 | if ((KeyEvent->wRepeatCount == 0) || ((KeyEvent->uChar.UnicodeChar == UNICODE_NULL) && |
| 105 | ((KeyEvent->wVirtualScanCode != 0) || (WI_IsFlagSet(KeyEvent->dwControlKeyState, ENHANCED_KEY))))) |
| 106 | { |
| 107 | return FALSE; |
| 108 | } |
| 109 | |
| 110 | return TRUE; |
| 111 | } |
| 112 | |
| 113 | BOOL GetNextCharacter(_In_ INPUT_RECORD* InputRecord, _Out_ PWCHAR NextCharacter) |
| 114 | { |
| 115 | BOOL IsNextCharacterValid = FALSE; |
| 116 | if (InputRecord->EventType == KEY_EVENT) |
| 117 | { |
| 118 | const auto KeyEvent = &InputRecord->Event.KeyEvent; |
| 119 | if ((IsActionableKey(KeyEvent) != FALSE) && ((KeyEvent->bKeyDown != FALSE) || (KeyEvent->wVirtualKeyCode == TTY_ALT_NUMPAD_VK_MENU))) |
| 120 | { |
| 121 | *NextCharacter = KeyEvent->uChar.UnicodeChar; |
| 122 | IsNextCharacterValid = TRUE; |
| 123 | } |
| 124 | } |
| 125 | |
| 126 | return IsNextCharacterValid; |
| 127 | } |
| 128 | |
| 129 | } // namespace |
| 130 | |
| 131 | // HandleWrapper |
| 132 | |
| 133 | HandleWrapper::HandleWrapper(HandleWrapper&& other) noexcept : |
| 134 | Handle(std::exchange(other.Handle, nullptr)), OwnedHandle(std::move(other.OwnedHandle)), OnClose(std::move(other.OnClose)) |
| 135 | { |
| 136 | other.OnClose = nullptr; |
| 137 | } |
| 138 | |
| 139 | HandleWrapper& HandleWrapper::operator=(HandleWrapper&& other) noexcept |
| 140 | { |
| 141 | if (this != &other) |
| 142 | { |
| 143 | Reset(); |
| 144 | Handle = std::exchange(other.Handle, nullptr); |
| 145 | OwnedHandle = std::move(other.OwnedHandle); |
| 146 | OnClose = std::move(other.OnClose); |
| 147 | other.OnClose = nullptr; |
| 148 | } |
| 149 | |
| 150 | return *this; |
| 151 | } |
| 152 | |
| 153 | HandleWrapper::HandleWrapper(wil::unique_handle&& handle, std::function<void()>&& OnClose) : |
| 154 | Handle(handle.get()), OwnedHandle(std::move(handle)), OnClose(std::move(OnClose)) |
| 155 | { |
| 156 | } |
| 157 | |
| 158 | HandleWrapper::HandleWrapper(wil::unique_socket&& handle, std::function<void()>&& OnClose) : |
| 159 | Handle((HANDLE)handle.get()), OwnedHandle(wil::unique_socket{handle.release()}), OnClose(std::move(OnClose)) |
| 160 | { |
| 161 | } |
| 162 | |
| 163 | HandleWrapper::HandleWrapper(wil::shared_handle handle, std::function<void()>&& OnClose) : |
| 164 | Handle(handle.get()), OwnedHandle(std::move(handle)), OnClose(std::move(OnClose)) |
| 165 | { |
| 166 | } |
| 167 | |
| 168 | HandleWrapper::HandleWrapper(wil::shared_socket handle, std::function<void()>&& OnClose) : |
| 169 | Handle(reinterpret_cast<HANDLE>(handle.get())), OwnedHandle(std::move(handle)), OnClose(std::move(OnClose)) |
| 170 | { |
| 171 | } |
| 172 | |
| 173 | HandleWrapper::HandleWrapper(wil::unique_event&& handle, std::function<void()>&& OnClose) : |
| 174 | Handle(handle.get()), OwnedHandle(wil::unique_handle{handle.release()}), OnClose(std::move(OnClose)) |
| 175 | { |
| 176 | } |
| 177 | |
| 178 | HandleWrapper::HandleWrapper(SOCKET handle, std::function<void()>&& OnClose) : |
| 179 | Handle(reinterpret_cast<HANDLE>(handle)), OnClose(std::move(OnClose)) |
| 180 | { |
| 181 | } |
| 182 | |
| 183 | HandleWrapper::HandleWrapper(HANDLE handle, std::function<void()>&& OnClose) : Handle(handle), OnClose(std::move(OnClose)) |
| 184 | { |
| 185 | } |
| 186 | |
| 187 | HandleWrapper::HandleWrapper(wil::unique_hfile&& handle, std::function<void()>&& OnClose) : |
| 188 | Handle(handle.get()), OwnedHandle(wil::unique_handle{handle.release()}), OnClose(std::move(OnClose)) |
| 189 | { |
| 190 | } |
| 191 | |
| 192 | HandleWrapper::~HandleWrapper() |
| 193 | { |
| 194 | Reset(); |
| 195 | } |
| 196 | |
| 197 | HANDLE HandleWrapper::Get() const |
| 198 | { |
| 199 | return Handle; |
| 200 | } |
| 201 | |
| 202 | bool HandleWrapper::IsValid() const |
| 203 | { |
| 204 | return Handle != nullptr && Handle != INVALID_HANDLE_VALUE; |
| 205 | } |
| 206 | |
| 207 | void HandleWrapper::Reset() |
| 208 | { |
| 209 | if (OnClose != nullptr) |
| 210 | { |
| 211 | OnClose(); |
| 212 | OnClose = nullptr; |
| 213 | } |
| 214 | |
| 215 | OwnedHandle = {}; |
| 216 | Handle = nullptr; |
| 217 | } |
| 218 | |
| 219 | // BufferWrapper |
| 220 | |
| 221 | BufferWrapper::BufferWrapper(size_t size) : m_owned(std::in_place, size) |
| 222 | { |
| 223 | } |
| 224 | |
| 225 | BufferWrapper::BufferWrapper(gsl::span<gsl::byte> span) : m_unowned(span) |
| 226 | { |
| 227 | } |
| 228 | |
| 229 | bool BufferWrapper::Owned() const noexcept |
| 230 | { |
| 231 | return m_owned.has_value(); |
| 232 | } |
| 233 | |
| 234 | void BufferWrapper::Resize(size_t size) |
| 235 | { |
| 236 | THROW_HR_IF_MSG(E_UNEXPECTED, !Owned(), "BufferWrapper::Resize called on a non-owned buffer"); |
| 237 | m_owned->resize(size); |
| 238 | } |
| 239 | |
| 240 | void BufferWrapper::Append(gsl::span<char> Span) |
| 241 | { |
| 242 | THROW_HR_IF_MSG(E_UNEXPECTED, !Owned(), "BufferWrapper::Append called on a non-owned buffer"); |
| 243 | |
| 244 | m_owned->insert(m_owned->end(), Span.begin(), Span.end()); |
| 245 | } |
| 246 | |
| 247 | void BufferWrapper::Consume(size_t bytes) noexcept |
| 248 | { |
| 249 | WI_ASSERT(bytes <= Size()); |
| 250 | if (Owned()) |
| 251 | { |
| 252 | m_owned->erase(m_owned->begin(), m_owned->begin() + bytes); |
| 253 | } |
| 254 | else |
| 255 | { |
| 256 | m_unowned = m_unowned.subspan(bytes); |
| 257 | } |
| 258 | } |
| 259 | |
| 260 | gsl::span<gsl::byte> BufferWrapper::Span() noexcept |
| 261 | { |
| 262 | return Owned() ? gsl::make_span(reinterpret_cast<gsl::byte*>(m_owned->data()), m_owned->size()) : m_unowned; |
| 263 | } |
| 264 | |
| 265 | size_t BufferWrapper::Size() const noexcept |
| 266 | { |
| 267 | return Owned() ? m_owned->size() : m_unowned.size(); |
| 268 | } |
| 269 | |
| 270 | // OverlappedIOHandle |
| 271 | |
| 272 | IOHandleStatus OverlappedIOHandle::GetState() const |
| 273 | { |
| 274 | return State; |
| 275 | } |
| 276 | |
| 277 | // EventHandle |
| 278 | |
| 279 | EventHandle::EventHandle(HandleWrapper&& Handle, std::function<void()>&& OnSignalled) : |
| 280 | Handle(std::move(Handle)), OnSignalled(std::move(OnSignalled)) |
| 281 | { |
| 282 | } |
| 283 | |
| 284 | void EventHandle::Schedule() |
| 285 | { |
| 286 | State = IOHandleStatus::Pending; |
| 287 | } |
| 288 | |
| 289 | void EventHandle::Collect() |
| 290 | { |
| 291 | State = IOHandleStatus::Completed; |
| 292 | OnSignalled(); |
| 293 | } |
| 294 | |
| 295 | HANDLE EventHandle::GetHandle() const |
| 296 | { |
| 297 | return Handle.Get(); |
| 298 | } |
| 299 | |
| 300 | // ReadHandle |
| 301 | |
| 302 | ReadHandle::ReadHandle(HandleWrapper&& MovedHandle, std::function<void(const gsl::span<char>& Buffer)>&& OnRead) : |
| 303 | Handle(std::move(MovedHandle)), OnRead(OnRead), Offset(InitializeFileOffset(Handle.Get())) |
| 304 | { |
| 305 | Overlapped.hEvent = Event.get(); |
| 306 | } |
| 307 | |
| 308 | ReadHandle::~ReadHandle() |
| 309 | { |
| 310 | if (State == IOHandleStatus::Pending) |
| 311 | { |
| 312 | CancelPendingIo(Handle.Get(), Overlapped); |
| 313 | } |
| 314 | } |
| 315 | |
| 316 | void ReadHandle::Schedule() |
| 317 | { |
| 318 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 319 | |
| 320 | Event.ResetEvent(); |
| 321 | |
| 322 | // Schedule the read. |
| 323 | DWORD bytesRead{}; |
| 324 | Overlapped.Offset = Offset.LowPart; |
| 325 | Overlapped.OffsetHigh = Offset.HighPart; |
| 326 | auto* bufferData = reinterpret_cast<char*>(Buffer.Span().data()); |
| 327 | if (ReadFile(Handle.Get(), bufferData, static_cast<DWORD>(Buffer.Size()), &bytesRead, &Overlapped)) |
| 328 | { |
| 329 | Offset.QuadPart += bytesRead; |
| 330 | |
| 331 | // Signal the read. |
| 332 | OnRead(gsl::make_span<char>(bufferData, static_cast<size_t>(bytesRead))); |
| 333 | |
| 334 | // ReadFile completed immediately, process the result right away. |
| 335 | if (bytesRead == 0) |
| 336 | { |
| 337 | State = IOHandleStatus::Completed; |
| 338 | return; // Handle is completely read, don't try again. |
| 339 | } |
| 340 | |
| 341 | // Read was done synchronously, remain in 'standby' state. |
| 342 | } |
| 343 | else |
| 344 | { |
| 345 | auto error = GetLastError(); |
| 346 | if (error == ERROR_HANDLE_EOF || error == ERROR_BROKEN_PIPE) |
| 347 | { |
| 348 | // Signal an empty read for EOF. |
| 349 | OnRead({}); |
| 350 | |
| 351 | State = IOHandleStatus::Completed; |
| 352 | return; |
| 353 | } |
| 354 | |
| 355 | THROW_LAST_ERROR_IF_MSG(error != ERROR_IO_PENDING, "Handle: 0x%p", (void*)Handle.Get()); |
| 356 | |
| 357 | // The read is pending, update to 'Pending' |
| 358 | State = IOHandleStatus::Pending; |
| 359 | } |
| 360 | } |
| 361 | |
| 362 | void ReadHandle::Collect() |
| 363 | { |
| 364 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 365 | |
| 366 | // Transition back to standby |
| 367 | State = IOHandleStatus::Standby; |
| 368 | |
| 369 | // Complete the read. |
| 370 | DWORD bytesRead{}; |
| 371 | if (!GetOverlappedResult(Handle.Get(), &Overlapped, &bytesRead, false)) |
| 372 | { |
| 373 | auto error = GetLastError(); |
| 374 | THROW_WIN32_IF(error, error != ERROR_HANDLE_EOF && error != ERROR_BROKEN_PIPE); |
| 375 | |
| 376 | // We received ERROR_HANDLE_EOF or ERROR_BROKEN_PIPE. Validate that this was indeed a zero byte read. |
| 377 | WI_ASSERT(bytesRead == 0); |
| 378 | } |
| 379 | |
| 380 | Offset.QuadPart += bytesRead; |
| 381 | |
| 382 | // Signal the read. |
| 383 | OnRead(gsl::make_span<char>(reinterpret_cast<char*>(Buffer.Span().data()), static_cast<size_t>(bytesRead))); |
| 384 | |
| 385 | // Transition to Complete if this was a zero byte read. |
| 386 | if (bytesRead == 0) |
| 387 | { |
| 388 | State = IOHandleStatus::Completed; |
| 389 | } |
| 390 | } |
| 391 | |
| 392 | HANDLE ReadHandle::GetHandle() const |
| 393 | { |
| 394 | return Event.get(); |
| 395 | } |
| 396 | |
| 397 | // ReadNamedPipe |
| 398 | |
| 399 | ReadNamedPipe::ReadNamedPipe(HandleWrapper&& Pipe, std::function<void(const gsl::span<char>& Buffer)>&& OnRead) : |
| 400 | ReadHandle(std::move(Pipe), std::move(OnRead)) |
| 401 | { |
| 402 | } |
| 403 | |
| 404 | void ReadNamedPipe::Schedule() |
| 405 | { |
| 406 | if (!m_connected) |
| 407 | { |
| 408 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 409 | |
| 410 | if (!ConnectNamedPipe(Handle.Get(), &Overlapped)) |
| 411 | { |
| 412 | const auto error = GetLastError(); |
| 413 | if (error == ERROR_IO_PENDING) |
| 414 | { |
| 415 | State = IOHandleStatus::Pending; |
| 416 | return; |
| 417 | } |
| 418 | |
| 419 | THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != ERROR_PIPE_CONNECTED, "Handle: 0x%p", (void*)Handle.Get()); |
| 420 | } |
| 421 | |
| 422 | m_connected = true; |
| 423 | } |
| 424 | |
| 425 | ReadHandle::Schedule(); |
| 426 | } |
| 427 | |
| 428 | void ReadNamedPipe::Collect() |
| 429 | { |
| 430 | if (!m_connected) |
| 431 | { |
| 432 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 433 | |
| 434 | DWORD bytes{}; |
| 435 | if (!GetOverlappedResult(Handle.Get(), &Overlapped, &bytes, FALSE)) |
| 436 | { |
| 437 | const auto error = GetLastError(); |
| 438 | THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != ERROR_PIPE_CONNECTED, "Handle: 0x%p", (void*)Handle.Get()); |
| 439 | } |
| 440 | |
| 441 | m_connected = true; |
| 442 | |
| 443 | // Transition back to standby so the IO loop schedules the first read. |
| 444 | State = IOHandleStatus::Standby; |
| 445 | return; |
| 446 | } |
| 447 | |
| 448 | ReadHandle::Collect(); |
| 449 | } |
| 450 | |
| 451 | // AcceptHandle |
| 452 | |
| 453 | AcceptHandle::AcceptHandle(HandleWrapper&& ListenSocket, bool AcceptOnce, std::function<void(wil::unique_socket&&)>&& OnAccepted) : |
| 454 | ListenSocket(std::move(ListenSocket)), AcceptOnce(AcceptOnce), OnAccepted(std::move(OnAccepted)) |
| 455 | { |
| 456 | Overlapped.hEvent = Event.get(); |
| 457 | |
| 458 | // Query the listen socket so accepted sockets can be created with a matching address family, type, and protocol. |
| 459 | WSAPROTOCOL_INFOW protocolInfo{}; |
| 460 | int length = sizeof(protocolInfo); |
| 461 | THROW_LAST_ERROR_IF( |
| 462 | getsockopt(reinterpret_cast<SOCKET>(this->ListenSocket.Get()), SOL_SOCKET, SO_PROTOCOL_INFOW, reinterpret_cast<char*>(&protocolInfo), &length) == |
| 463 | SOCKET_ERROR); |
| 464 | |
| 465 | AddressFamily = protocolInfo.iAddressFamily; |
| 466 | SocketType = protocolInfo.iSocketType; |
| 467 | Protocol = protocolInfo.iProtocol; |
| 468 | } |
| 469 | |
| 470 | AcceptHandle::~AcceptHandle() |
| 471 | { |
| 472 | if (State == IOHandleStatus::Pending) |
| 473 | { |
| 474 | CancelPendingIo(reinterpret_cast<SOCKET>(ListenSocket.Get()), Overlapped); |
| 475 | } |
| 476 | } |
| 477 | |
| 478 | void AcceptHandle::CreateAcceptSocket() |
| 479 | { |
| 480 | AcceptedSocket.reset(WSASocketW(AddressFamily, SocketType, Protocol, nullptr, 0, WSA_FLAG_OVERLAPPED)); |
| 481 | THROW_LAST_ERROR_IF(!AcceptedSocket); |
| 482 | |
| 483 | if (AddressFamily == AF_HYPERV) |
| 484 | { |
| 485 | ULONG enable = 1; |
| 486 | THROW_LAST_ERROR_IF( |
| 487 | setsockopt(AcceptedSocket.get(), HV_PROTOCOL_RAW, HVSOCKET_CONNECTED_SUSPEND, reinterpret_cast<char*>(&enable), sizeof(enable)) == |
| 488 | SOCKET_ERROR); |
| 489 | } |
| 490 | } |
| 491 | |
| 492 | void AcceptHandle::OnComplete() |
| 493 | { |
| 494 | wsl::windows::common::socket::SetAcceptContext(AcceptedSocket.get(), reinterpret_cast<SOCKET>(ListenSocket.Get())); |
| 495 | |
| 496 | OnAccepted(std::move(AcceptedSocket)); |
| 497 | |
| 498 | if (AcceptOnce) |
| 499 | { |
| 500 | State = IOHandleStatus::Completed; |
| 501 | } |
| 502 | else |
| 503 | { |
| 504 | State = IOHandleStatus::Standby; |
| 505 | } |
| 506 | } |
| 507 | |
| 508 | void AcceptHandle::Schedule() |
| 509 | { |
| 510 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 511 | |
| 512 | CreateAcceptSocket(); |
| 513 | |
| 514 | Event.ResetEvent(); |
| 515 | |
| 516 | // Schedule the accept. |
| 517 | DWORD bytesReturned{}; |
| 518 | if (AcceptEx((SOCKET)ListenSocket.Get(), AcceptedSocket.get(), &AcceptBuffer, 0, sizeof(SOCKADDR_STORAGE), sizeof(SOCKADDR_STORAGE), &bytesReturned, &Overlapped)) |
| 519 | { |
| 520 | // Accept completed immediately. |
| 521 | OnComplete(); |
| 522 | } |
| 523 | else |
| 524 | { |
| 525 | auto error = WSAGetLastError(); |
| 526 | THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != ERROR_IO_PENDING, "Handle: 0x%p", (void*)ListenSocket.Get()); |
| 527 | |
| 528 | State = IOHandleStatus::Pending; |
| 529 | } |
| 530 | } |
| 531 | |
| 532 | void AcceptHandle::Collect() |
| 533 | { |
| 534 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 535 | |
| 536 | DWORD bytesReceived{}; |
| 537 | DWORD flagsReturned{}; |
| 538 | |
| 539 | THROW_IF_WIN32_BOOL_FALSE(WSAGetOverlappedResult((SOCKET)ListenSocket.Get(), &Overlapped, &bytesReceived, false, &flagsReturned)); |
| 540 | |
| 541 | OnComplete(); |
| 542 | } |
| 543 | |
| 544 | HANDLE AcceptHandle::GetHandle() const |
| 545 | { |
| 546 | return Event.get(); |
| 547 | } |
| 548 | |
| 549 | // LineBasedReadHandle |
| 550 | |
| 551 | LineBasedReadHandle::LineBasedReadHandle(HandleWrapper&& Handle, std::function<void(const gsl::span<char>& Line)>&& OnLine, bool Crlf) : |
| 552 | ReadHandle(std::move(Handle), [this](const gsl::span<char>& Buffer) { OnRead(Buffer); }), OnLine(OnLine), Crlf(Crlf) |
| 553 | { |
| 554 | } |
| 555 | |
| 556 | LineBasedReadHandle::~LineBasedReadHandle() |
| 557 | { |
| 558 | // N.B. PendingBuffer can contain remaining data is an exception was thrown during parsing. |
| 559 | } |
| 560 | |
| 561 | void LineBasedReadHandle::OnRead(const gsl::span<char>& Buffer) |
| 562 | { |
| 563 | // If we reach of the end, signal a line with the remaining buffer. |
| 564 | if (Buffer.empty() && !PendingBuffer.empty()) |
| 565 | { |
| 566 | OnLine(PendingBuffer); |
| 567 | PendingBuffer.clear(); |
| 568 | return; |
| 569 | } |
| 570 | |
| 571 | auto begin = Buffer.begin(); |
| 572 | auto end = std::ranges::find(Buffer, Crlf ? '\r' : '\n'); |
| 573 | while (end != Buffer.end()) |
| 574 | { |
| 575 | if (Crlf) |
| 576 | { |
| 577 | end++; // Move to the following '\n' |
| 578 | |
| 579 | if (end == Buffer.end() || *end != '\n') // Incomplete CRLF sequence. Append to buffer and continue. |
| 580 | { |
| 581 | PendingBuffer.insert(PendingBuffer.end(), begin, end); |
| 582 | begin = end; |
| 583 | end = std::ranges::find(end, Buffer.end(), '\r'); |
| 584 | continue; |
| 585 | } |
| 586 | } |
| 587 | |
| 588 | // Discard the '\r' in CRLF mode. |
| 589 | PendingBuffer.insert(PendingBuffer.end(), begin, Crlf ? end - 1 : end); |
| 590 | |
| 591 | if (!PendingBuffer.empty()) |
| 592 | { |
| 593 | OnLine(PendingBuffer); |
| 594 | PendingBuffer.clear(); |
| 595 | } |
| 596 | |
| 597 | begin = end + 1; |
| 598 | end = std::ranges::find(begin, Buffer.end(), Crlf ? '\r' : '\n'); |
| 599 | } |
| 600 | |
| 601 | PendingBuffer.insert(PendingBuffer.end(), begin, end); |
| 602 | } |
| 603 | |
| 604 | // HTTPChunkBasedReadHandle |
| 605 | |
| 606 | HTTPChunkBasedReadHandle::HTTPChunkBasedReadHandle(HandleWrapper&& MovedHandle, std::function<void(const gsl::span<char>& Line)>&& OnChunk) : |
| 607 | ReadHandle(std::move(MovedHandle), [this](const gsl::span<char>& Buffer) { OnRead(Buffer); }), OnChunk(std::move(OnChunk)) |
| 608 | { |
| 609 | } |
| 610 | |
| 611 | HTTPChunkBasedReadHandle::~HTTPChunkBasedReadHandle() |
| 612 | { |
| 613 | // N.B. PendingBuffer can contain remaining data is an exception was thrown during parsing. |
| 614 | LOG_HR_IF(E_UNEXPECTED, !PendingBuffer.empty() || PendingChunkSize != 0 || ExpectHeader); |
| 615 | } |
| 616 | |
| 617 | void HTTPChunkBasedReadHandle::OnRead(const gsl::span<char>& Input) |
| 618 | { |
| 619 | // See: https://httpwg.org/specs/rfc9112.html#field.transfer-encoding |
| 620 | |
| 621 | if (Input.empty()) |
| 622 | { |
| 623 | // N.B. The body can be terminated by a zero-length chunk. |
| 624 | THROW_HR_IF(E_INVALIDARG, PendingChunkSize != 0 || ExpectHeader); |
| 625 | } |
| 626 | |
| 627 | auto buffer = Input; |
| 628 | |
| 629 | auto advance = [&](size_t count) { |
| 630 | WI_ASSERT(buffer.size() >= count); |
| 631 | buffer = buffer.subspan(count); |
| 632 | }; |
| 633 | |
| 634 | while (!buffer.empty()) |
| 635 | { |
| 636 | if (PendingChunkSize == 0) |
| 637 | { |
| 638 | // Consume CRLF's between chunks. |
| 639 | if (PendingBuffer.empty() && (buffer.front() == '\r' || buffer.front() == '\n')) |
| 640 | { |
| 641 | advance(1); |
| 642 | continue; |
| 643 | } |
| 644 | |
| 645 | ExpectHeader = true; |
| 646 | |
| 647 | auto end = std::ranges::find(buffer, '\n'); |
| 648 | PendingBuffer.insert(PendingBuffer.end(), buffer.begin(), end); |
| 649 | if (end == buffer.end()) |
| 650 | { |
| 651 | // Incomplete size header, buffer until next read. |
| 652 | break; |
| 653 | } |
| 654 | // Advance beyond the LF |
| 655 | advance(end - buffer.begin() + 1); |
| 656 | |
| 657 | THROW_HR_IF_MSG( |
| 658 | E_INVALIDARG, |
| 659 | PendingBuffer.size() < 2 || PendingBuffer.back() != '\r', |
| 660 | "Malformed chunk header: %hs", |
| 661 | PendingBuffer.c_str()); |
| 662 | PendingBuffer.erase(PendingBuffer.end() - 1, PendingBuffer.end()); // Remove CR. |
| 663 | |
| 664 | #ifdef WSLC_HTTP_DEBUG |
| 665 | |
| 666 | WSL_LOG("HTTPChunkHeader", TraceLoggingValue(PendingBuffer.c_str(), "Size")); |
| 667 | |
| 668 | #endif |
| 669 | |
| 670 | try |
| 671 | { |
| 672 | size_t parsed{}; |
| 673 | PendingChunkSize = std::stoul(PendingBuffer.c_str(), &parsed, 16); |
| 674 | THROW_HR_IF(E_INVALIDARG, parsed != PendingBuffer.size()); |
| 675 | } |
| 676 | catch (...) |
| 677 | { |
| 678 | THROW_HR_MSG(E_INVALIDARG, "Failed to parse chunk size: %hs", PendingBuffer.c_str()); |
| 679 | } |
| 680 | |
| 681 | ExpectHeader = false; |
| 682 | PendingBuffer.clear(); |
| 683 | } |
| 684 | else |
| 685 | { |
| 686 | // Consume the chunk. |
| 687 | auto consumedBytes = std::min(PendingChunkSize, buffer.size()); |
| 688 | PendingBuffer.append(buffer.data(), consumedBytes); |
| 689 | advance(consumedBytes); |
| 690 | |
| 691 | WI_ASSERT(PendingChunkSize >= consumedBytes); |
| 692 | PendingChunkSize -= consumedBytes; |
| 693 | |
| 694 | if (PendingChunkSize == 0) |
| 695 | { |
| 696 | |
| 697 | #ifdef WSLC_HTTP_DEBUG |
| 698 | |
| 699 | WSL_LOG("HTTPChunk", TraceLoggingValue(PendingBuffer.c_str(), "Content")); |
| 700 | |
| 701 | #endif |
| 702 | OnChunk(PendingBuffer); |
| 703 | PendingBuffer.clear(); |
| 704 | } |
| 705 | } |
| 706 | } |
| 707 | } |
| 708 | |
| 709 | // ReadSocketMessageHandle |
| 710 | |
| 711 | ReadSocketMessageHandle::ReadSocketMessageHandle( |
| 712 | HandleWrapper&& MovedSocket, |
| 713 | std::vector<gsl::byte>& Buffer, |
| 714 | std::vector<gsl::byte>& PendingBytes, |
| 715 | std::function<void(const gsl::span<gsl::byte>& Message)>&& OnMessage) : |
| 716 | Socket(std::move(MovedSocket)), Buffer(Buffer), PendingBytes(PendingBytes), OnMessage(std::move(OnMessage)) |
| 717 | { |
| 718 | Overlapped.hEvent = Event.get(); |
| 719 | |
| 720 | if (Buffer.size() < sizeof(MESSAGE_HEADER)) |
| 721 | { |
| 722 | Buffer.resize(sizeof(MESSAGE_HEADER)); |
| 723 | } |
| 724 | |
| 725 | if (PendingBytes.empty()) |
| 726 | { |
| 727 | return; |
| 728 | } |
| 729 | |
| 730 | // If bytes from a previously cancelled transaction are passed, process them now. |
| 731 | if (Buffer.size() < PendingBytes.size()) |
| 732 | { |
| 733 | Buffer.resize(PendingBytes.size()); |
| 734 | } |
| 735 | |
| 736 | std::copy(PendingBytes.begin(), PendingBytes.end(), Buffer.begin()); |
| 737 | CurrentOffset = PendingBytes.size(); |
| 738 | PendingBytes.clear(); |
| 739 | |
| 740 | if (CurrentOffset < sizeof(MESSAGE_HEADER)) |
| 741 | { |
| 742 | BytesRemaining = sizeof(MESSAGE_HEADER) - CurrentOffset; |
| 743 | } |
| 744 | else |
| 745 | { |
| 746 | BytesRemaining = 0; |
| 747 | } |
| 748 | } |
| 749 | |
| 750 | ReadSocketMessageHandle::~ReadSocketMessageHandle() |
| 751 | { |
| 752 | if (State != IOHandleStatus::Completed) |
| 753 | { |
| 754 | auto pendingSize = CurrentOffset; |
| 755 | |
| 756 | if (State == IOHandleStatus::Pending) |
| 757 | { |
| 758 | // Cancel the pending receive and move any bytes already buffered for the in-flight message into PendingBytes |
| 759 | const auto socket = reinterpret_cast<SOCKET>(Socket.Get()); |
| 760 | pendingSize += CancelPendingIo(socket, Overlapped); |
| 761 | } |
| 762 | |
| 763 | if (pendingSize > 0) |
| 764 | { |
| 765 | WI_ASSERT(pendingSize <= Buffer.size()); |
| 766 | PendingBytes.assign(Buffer.begin(), Buffer.begin() + pendingSize); |
| 767 | |
| 768 | WSL_LOG( |
| 769 | "CanceledMessageRead", TraceLoggingValue(pendingSize, "TotalBytes"), TraceLoggingValue(Socket.Get(), "Socket")); |
| 770 | } |
| 771 | } |
| 772 | } |
| 773 | |
| 774 | void ReadSocketMessageHandle::ScheduleRecv() |
| 775 | { |
| 776 | Event.ResetEvent(); |
| 777 | |
| 778 | auto target = gsl::make_span(Buffer).subspan(CurrentOffset, BytesRemaining); |
| 779 | WSABUF wsaBuf = {gsl::narrow_cast<ULONG>(target.size()), reinterpret_cast<CHAR*>(target.data())}; |
| 780 | DWORD bytesRead{}; |
| 781 | DWORD flags = 0; |
| 782 | if (WSARecv(reinterpret_cast<SOCKET>(Socket.Get()), &wsaBuf, 1, &bytesRead, &flags, &Overlapped, nullptr) == 0) |
| 783 | { |
| 784 | ProcessRecvResult(bytesRead); |
| 785 | } |
| 786 | else |
| 787 | { |
| 788 | auto error = WSAGetLastError(); |
| 789 | if (error == WSAECONNABORTED || error == WSAECONNRESET) |
| 790 | { |
| 791 | ProcessRecvResult(0); |
| 792 | return; |
| 793 | } |
| 794 | |
| 795 | THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != WSA_IO_PENDING, "Socket: 0x%p", (void*)Socket.Get()); |
| 796 | |
| 797 | State = IOHandleStatus::Pending; |
| 798 | } |
| 799 | } |
| 800 | |
| 801 | void ReadSocketMessageHandle::ProcessRecvResult(DWORD BytesRead) |
| 802 | { |
| 803 | if (BytesRead == 0) |
| 804 | { |
| 805 | // If the socket was closed before any bytes of the next message were read, signal a clean end-of-stream. |
| 806 | // If some bytes were already buffered, the peer closed mid-message which is a protocol error. |
| 807 | THROW_HR_IF_MSG( |
| 808 | E_UNEXPECTED, |
| 809 | CurrentOffset > 0, |
| 810 | "Socket closed before a complete message could be read. ReadingHeader: %d, CurrentOffset: %zu, BytesRemaining: %zu", |
| 811 | ReadingHeader, |
| 812 | CurrentOffset, |
| 813 | BytesRemaining); |
| 814 | |
| 815 | OnMessage({}); |
| 816 | State = IOHandleStatus::Completed; |
| 817 | return; |
| 818 | } |
| 819 | |
| 820 | CurrentOffset += BytesRead; |
| 821 | BytesRemaining -= BytesRead; |
| 822 | |
| 823 | if (BytesRemaining > 0) |
| 824 | { |
| 825 | return; |
| 826 | } |
| 827 | |
| 828 | ProcessChunk(); |
| 829 | } |
| 830 | |
| 831 | bool ReadSocketMessageHandle::ProcessChunk() |
| 832 | { |
| 833 | const auto messageSize = gslhelpers::get_struct<MESSAGE_HEADER>(gsl::make_span(Buffer.data(), sizeof(MESSAGE_HEADER)))->MessageSize; |
| 834 | |
| 835 | if (ReadingHeader) |
| 836 | { |
| 837 | THROW_HR_IF_MSG(E_UNEXPECTED, messageSize < sizeof(MESSAGE_HEADER), "Unexpected message size: %u", messageSize); |
| 838 | THROW_HR_IF_MSG(E_UNEXPECTED, messageSize > 16 * 1024 * 1024, "Message size too large: %u", messageSize); |
| 839 | |
| 840 | if (Buffer.size() < messageSize) |
| 841 | { |
| 842 | Buffer.resize(messageSize); |
| 843 | } |
| 844 | |
| 845 | ReadingHeader = false; |
| 846 | if (CurrentOffset < messageSize) |
| 847 | { |
| 848 | BytesRemaining = messageSize - CurrentOffset; |
| 849 | } |
| 850 | |
| 851 | if (BytesRemaining > 0) |
| 852 | { |
| 853 | return true; |
| 854 | } |
| 855 | } |
| 856 | |
| 857 | OnMessage(gsl::make_span(Buffer.data(), messageSize)); |
| 858 | State = IOHandleStatus::Completed; |
| 859 | return false; |
| 860 | } |
| 861 | |
| 862 | void ReadSocketMessageHandle::Schedule() |
| 863 | { |
| 864 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 865 | |
| 866 | // Process previously received bytes, if any. |
| 867 | if (BytesRemaining == 0 && !ProcessChunk()) |
| 868 | { |
| 869 | return; // Message has been fully received, no need to schedule a receive. |
| 870 | } |
| 871 | |
| 872 | ScheduleRecv(); |
| 873 | } |
| 874 | |
| 875 | void ReadSocketMessageHandle::Collect() |
| 876 | { |
| 877 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 878 | |
| 879 | State = IOHandleStatus::Standby; |
| 880 | |
| 881 | DWORD bytesRead{}; |
| 882 | DWORD flags{}; |
| 883 | if (!WSAGetOverlappedResult(reinterpret_cast<SOCKET>(Socket.Get()), &Overlapped, &bytesRead, FALSE, &flags)) |
| 884 | { |
| 885 | long error = WSAGetLastError(); |
| 886 | THROW_WIN32_IF(error, error != WSAECONNABORTED && error != WSAECONNRESET); |
| 887 | |
| 888 | WI_ASSERT(bytesRead == 0); |
| 889 | } |
| 890 | |
| 891 | ProcessRecvResult(bytesRead); |
| 892 | } |
| 893 | |
| 894 | HANDLE ReadSocketMessageHandle::GetHandle() const |
| 895 | { |
| 896 | return Event.get(); |
| 897 | } |
| 898 | |
| 899 | // ReadConsoleHandle |
| 900 | |
| 901 | wsl::windows::common::io::ReadConsoleHandle::ReadConsoleHandle( |
| 902 | HandleWrapper&& Console, |
| 903 | std::function<void(const gsl::span<char>& Buffer)>&& OnRead, |
| 904 | std::function<void()>&& UpdateTerminalSize, |
| 905 | std::vector<char> DetachSequence, |
| 906 | std::function<void()>&& OnDetach) : |
| 907 | Console(std::move(Console)), |
| 908 | OnRead(std::move(OnRead)), |
| 909 | UpdateTerminalSize(std::move(UpdateTerminalSize)), |
| 910 | DetachSequence(std::move(DetachSequence)), |
| 911 | OnDetach(std::move(OnDetach)) |
| 912 | { |
| 913 | } |
| 914 | |
| 915 | void wsl::windows::common::io::ReadConsoleHandle::Schedule() |
| 916 | { |
| 917 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 918 | |
| 919 | // |
| 920 | // Use the console handle as the signal event. |
| 921 | // N.B. This behavior is documented here: https://learn.microsoft.com/en-us/windows/console/readconsoleinput |
| 922 | // |
| 923 | |
| 924 | State = IOHandleStatus::Pending; |
| 925 | } |
| 926 | |
| 927 | HANDLE wsl::windows::common::io::ReadConsoleHandle::GetHandle() const |
| 928 | { |
| 929 | return Console.Get(); |
| 930 | } |
| 931 | |
| 932 | void wsl::windows::common::io::ReadConsoleHandle::Collect() |
| 933 | { |
| 934 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 935 | |
| 936 | // |
| 937 | // Re-arm by default; a detected detach sequence overrides this to Completed below. |
| 938 | // |
| 939 | |
| 940 | State = IOHandleStatus::Standby; |
| 941 | |
| 942 | // |
| 943 | // N.B. ReadConsoleInputEx has no associated import library. |
| 944 | // |
| 945 | |
| 946 | static LxssDynamicFunction<decltype(ReadConsoleInputExW)> readConsoleInput(L"Kernel32.dll", "ReadConsoleInputExW"); |
| 947 | |
| 948 | INPUT_RECORD InputRecordBuffer[TTY_INPUT_EVENT_BUFFER_SIZE]; |
| 949 | INPUT_RECORD* InputRecordPeek = &(InputRecordBuffer[1]); |
| 950 | KEY_EVENT_RECORD* KeyEvent; |
| 951 | DWORD RecordsRead; |
| 952 | |
| 953 | // |
| 954 | // The console handle stays signaled while input is available, so drain all currently available |
| 955 | // input here and return to Standby once none remains (the handle is waited on again by the IO loop). |
| 956 | // |
| 957 | |
| 958 | for (;;) |
| 959 | { |
| 960 | // Detach if the escape sequence was detected. |
| 961 | // N.B. This needs to be done at the beginning of the loop so the escape sequence is also delivered. |
| 962 | if (!CurrentSequence.empty() && std::ranges::equal(CurrentSequence, DetachSequence)) |
| 963 | { |
| 964 | OnDetach(); |
| 965 | State = IOHandleStatus::Completed; |
| 966 | return; |
| 967 | } |
| 968 | |
| 969 | // |
| 970 | // Because some input events generated by the console are encoded with |
| 971 | // more than one input event, we have to be smart about reading the |
| 972 | // events. |
| 973 | // |
| 974 | // First, we peek at the next input event. |
| 975 | // If it's an escape (wch == L'\x1b') event, then the characters that |
| 976 | // follow are part of an input sequence. We can't know for sure |
| 977 | // how long that sequence is, but we can assume it's all sent to |
| 978 | // the input queue at once, and it's less that 16 events. |
| 979 | // Furthermore, we can assume that if there's an Escape in those |
| 980 | // 16 events, that the escape marks the start of a new sequence. |
| 981 | // So, we'll peek at another 15 events looking for escapes. |
| 982 | // If we see an escape, then we'll read one less than that, |
| 983 | // such that the escape remains the next event in the input. |
| 984 | // From those read events, we'll aggregate chars into a single |
| 985 | // string to send to the subsystem. |
| 986 | // If it's not an escape, send the event through one at a time. |
| 987 | // |
| 988 | |
| 989 | // |
| 990 | // Read one input event without blocking. If none is available, all input has been drained. |
| 991 | // |
| 992 | |
| 993 | THROW_IF_WIN32_BOOL_FALSE(readConsoleInput(Console.Get(), InputRecordBuffer, 1, &RecordsRead, CONSOLE_READ_NOWAIT)); |
| 994 | if (RecordsRead == 0) |
| 995 | { |
| 996 | return; |
| 997 | } |
| 998 | |
| 999 | // |
| 1000 | // Don't read additional records if the first entry is a window size |
| 1001 | // event, or a repeated character. Handle those events on their own. |
| 1002 | // |
| 1003 | |
| 1004 | DWORD RecordsPeeked = 0; |
| 1005 | if ((InputRecordBuffer[0].EventType != WINDOW_BUFFER_SIZE_EVENT) && |
| 1006 | ((InputRecordBuffer[0].EventType != KEY_EVENT) || (InputRecordBuffer[0].Event.KeyEvent.wRepeatCount < 2))) |
| 1007 | { |
| 1008 | // |
| 1009 | // Read additional input records into the buffer if available. |
| 1010 | // |
| 1011 | |
| 1012 | THROW_IF_WIN32_BOOL_FALSE(PeekConsoleInputW(Console.Get(), InputRecordPeek, (RTL_NUMBER_OF(InputRecordBuffer) - 1), &RecordsPeeked)); |
| 1013 | } |
| 1014 | |
| 1015 | // |
| 1016 | // Iterate over peeked records [1, RecordsPeeked]. |
| 1017 | // |
| 1018 | |
| 1019 | DWORD AdditionalRecordsToRead = 0; |
| 1020 | WCHAR NextCharacter; |
| 1021 | for (DWORD RecordIndex = 1; RecordIndex <= RecordsPeeked; RecordIndex++) |
| 1022 | { |
| 1023 | if (GetNextCharacter(&InputRecordBuffer[RecordIndex], &NextCharacter) != FALSE) |
| 1024 | { |
| 1025 | KeyEvent = &InputRecordBuffer[RecordIndex].Event.KeyEvent; |
| 1026 | if (NextCharacter == TTY_ESCAPE_CHARACTER) |
| 1027 | { |
| 1028 | // |
| 1029 | // CurrentRecord is an escape event. We will start here |
| 1030 | // on the next input loop. |
| 1031 | // |
| 1032 | |
| 1033 | break; |
| 1034 | } |
| 1035 | else if (KeyEvent->wRepeatCount > 1) |
| 1036 | { |
| 1037 | // |
| 1038 | // Repeated keys are handled on their own. Start with this |
| 1039 | // key on the next input loop. |
| 1040 | // |
| 1041 | |
| 1042 | break; |
| 1043 | } |
| 1044 | else if (IS_HIGH_SURROGATE(NextCharacter) && (RecordIndex >= (RecordsPeeked - 1))) |
| 1045 | { |
| 1046 | // |
| 1047 | // If there is not enough room for the second character of |
| 1048 | // a surrogate pair, start with this character on the next |
| 1049 | // input loop. |
| 1050 | // |
| 1051 | // N.B. The test is for at least two remaining records |
| 1052 | // because typically a surrogate pair will be entered |
| 1053 | // via copy/paste, which will appear as an input |
| 1054 | // record with alt-down, alt-up and character. So to |
| 1055 | // include the next character of the surrogate pair it |
| 1056 | // is likely that the alt-up record will need to be |
| 1057 | // read first. |
| 1058 | // |
| 1059 | |
| 1060 | break; |
| 1061 | } |
| 1062 | } |
| 1063 | else if (InputRecordBuffer[RecordIndex].EventType == WINDOW_BUFFER_SIZE_EVENT) |
| 1064 | { |
| 1065 | // |
| 1066 | // A window size event is handled on its own. |
| 1067 | // |
| 1068 | |
| 1069 | break; |
| 1070 | } |
| 1071 | |
| 1072 | // |
| 1073 | // Process the additional input record. |
| 1074 | // |
| 1075 | |
| 1076 | AdditionalRecordsToRead += 1; |
| 1077 | } |
| 1078 | |
| 1079 | if (AdditionalRecordsToRead > 0) |
| 1080 | { |
| 1081 | THROW_IF_WIN32_BOOL_FALSE(readConsoleInput(Console.Get(), InputRecordPeek, AdditionalRecordsToRead, &RecordsRead, CONSOLE_READ_NOWAIT)); |
| 1082 | |
| 1083 | if (RecordsRead == 0) |
| 1084 | { |
| 1085 | // |
| 1086 | // This would be an unexpected case. We've already peeked to see |
| 1087 | // that there are AdditionalRecordsToRead # of records in the |
| 1088 | // input that need reading, yet we didn't get them when we read. |
| 1089 | // In this case, stop draining and wait to be signaled again. |
| 1090 | // |
| 1091 | |
| 1092 | return; |
| 1093 | } |
| 1094 | |
| 1095 | // |
| 1096 | // We already had one input record in the buffer before reading |
| 1097 | // additional, So account for that one too |
| 1098 | // |
| 1099 | |
| 1100 | RecordsRead += 1; |
| 1101 | } |
| 1102 | |
| 1103 | // |
| 1104 | // Process each input event. Keydowns will get aggregated into |
| 1105 | // Utf8String before getting injected into the subsystem. |
| 1106 | // |
| 1107 | |
| 1108 | WCHAR Utf16String[TTY_INPUT_EVENT_BUFFER_SIZE]; |
| 1109 | ULONG Utf16StringSize = 0; |
| 1110 | for (DWORD RecordIndex = 0; RecordIndex < RecordsRead; RecordIndex++) |
| 1111 | { |
| 1112 | INPUT_RECORD* CurrentInputRecord = &(InputRecordBuffer[RecordIndex]); |
| 1113 | switch (CurrentInputRecord->EventType) |
| 1114 | { |
| 1115 | case KEY_EVENT: |
| 1116 | |
| 1117 | KeyEvent = &CurrentInputRecord->Event.KeyEvent; |
| 1118 | |
| 1119 | if (KeyEvent->bKeyDown && IsActionableKey(KeyEvent) && !DetachSequence.empty()) |
| 1120 | { |
| 1121 | if (CurrentSequence.size() >= DetachSequence.size()) |
| 1122 | { |
| 1123 | CurrentSequence.pop_front(); |
| 1124 | } |
| 1125 | |
| 1126 | CurrentSequence.push_back(CurrentInputRecord->Event.KeyEvent.uChar.AsciiChar); |
| 1127 | } |
| 1128 | |
| 1129 | // |
| 1130 | // Filter out key up events unless they are from an <Alt> key. |
| 1131 | // Key up with an <Alt> key could contain a Unicode character |
| 1132 | // pasted from the clipboard and converted to an <Alt>+<Numpad> sequence. |
| 1133 | // |
| 1134 | |
| 1135 | if ((KeyEvent->bKeyDown == FALSE) && (KeyEvent->wVirtualKeyCode != TTY_ALT_NUMPAD_VK_MENU)) |
| 1136 | { |
| 1137 | break; |
| 1138 | } |
| 1139 | |
| 1140 | // |
| 1141 | // Filter out key presses that are not actionable, such as just |
| 1142 | // pressing <Ctrl>, <Alt>, <Shift> etc. These key presses return |
| 1143 | // the character of null but will have a valid scan code off the |
| 1144 | // keyboard. Certain other key sequences such as Ctrl+A, |
| 1145 | // Ctrl+<space>, and Ctrl+@ will also return the character null |
| 1146 | // but have no scan code. |
| 1147 | // <Alt> + <NumPad> sequences will show an <Alt> but will have |
| 1148 | // a scancode and character specified, so they should be actionable. |
| 1149 | // |
| 1150 | |
| 1151 | if (IsActionableKey(KeyEvent) == FALSE) |
| 1152 | { |
| 1153 | break; |
| 1154 | } |
| 1155 | |
| 1156 | Utf16String[Utf16StringSize] = KeyEvent->uChar.UnicodeChar; |
| 1157 | Utf16StringSize += 1; |
| 1158 | break; |
| 1159 | |
| 1160 | case WINDOW_BUFFER_SIZE_EVENT: |
| 1161 | |
| 1162 | // |
| 1163 | // Query the window size and send an update message via the |
| 1164 | // control channel. |
| 1165 | // |
| 1166 | |
| 1167 | UpdateTerminalSize(); |
| 1168 | break; |
| 1169 | } |
| 1170 | } |
| 1171 | |
| 1172 | CHAR Utf8String[TTY_UTF8_TRANSLATION_BUFFER_SIZE]; |
| 1173 | DWORD Utf8StringSize = 0; |
| 1174 | if (Utf16StringSize > 0) |
| 1175 | { |
| 1176 | // |
| 1177 | // Windows uses UTF-16LE encoding, Linux uses UTF-8 by default. |
| 1178 | // Convert each UTF-16LE character into the proper UTF-8 byte |
| 1179 | // sequence equivalent. |
| 1180 | // |
| 1181 | |
| 1182 | THROW_LAST_ERROR_IF( |
| 1183 | (Utf8StringSize = WideCharToMultiByte( |
| 1184 | CP_UTF8, 0, Utf16String, Utf16StringSize, Utf8String, sizeof(Utf8String), nullptr, nullptr)) == 0); |
| 1185 | } |
| 1186 | |
| 1187 | // |
| 1188 | // Deliver the translated input bytes. |
| 1189 | // |
| 1190 | |
| 1191 | const auto Utf8Span = gsl::make_span(Utf8String, static_cast<size_t>(Utf8StringSize)); |
| 1192 | if ((RecordsRead == 1) && (InputRecordBuffer[0].EventType == KEY_EVENT) && (InputRecordBuffer[0].Event.KeyEvent.wRepeatCount > 1)) |
| 1193 | { |
| 1194 | WI_ASSERT(Utf16StringSize == 1); |
| 1195 | |
| 1196 | // |
| 1197 | // Handle repeated characters. They aren't part of an input |
| 1198 | // sequence, so there's only one event that's generating characters. |
| 1199 | // |
| 1200 | |
| 1201 | for (WORD RepeatIndex = 0; RepeatIndex < InputRecordBuffer[0].Event.KeyEvent.wRepeatCount; RepeatIndex += 1) |
| 1202 | { |
| 1203 | OnRead(Utf8Span); |
| 1204 | } |
| 1205 | } |
| 1206 | else if (Utf8StringSize > 0) |
| 1207 | { |
| 1208 | OnRead(Utf8Span); |
| 1209 | } |
| 1210 | } |
| 1211 | } |
| 1212 | |
| 1213 | // WriteHandle |
| 1214 | |
| 1215 | WriteHandle::WriteHandle(HandleWrapper&& MovedHandle, const std::vector<char>& Source, bool CompleteOnDrained) : |
| 1216 | Handle(std::move(MovedHandle)), Buffer(Source.size()), Offset(InitializeFileOffset(Handle.Get())), CompleteOnDrained(CompleteOnDrained) |
| 1217 | { |
| 1218 | if (!Source.empty()) |
| 1219 | { |
| 1220 | std::memcpy(Buffer.Span().data(), Source.data(), Source.size()); |
| 1221 | } |
| 1222 | |
| 1223 | Overlapped.hEvent = Event.get(); |
| 1224 | |
| 1225 | if (!CompleteOnDrained && Buffer.Size() == 0) |
| 1226 | { |
| 1227 | State = IOHandleStatus::Idle; |
| 1228 | } |
| 1229 | } |
| 1230 | |
| 1231 | WriteHandle::WriteHandle(HandleWrapper&& MovedHandle, gsl::span<gsl::byte> Source) : |
| 1232 | Handle(std::move(MovedHandle)), Buffer(Source), Offset(InitializeFileOffset(Handle.Get())) |
| 1233 | { |
| 1234 | Overlapped.hEvent = Event.get(); |
| 1235 | } |
| 1236 | |
| 1237 | WriteHandle::~WriteHandle() |
| 1238 | { |
| 1239 | if (State == IOHandleStatus::Pending) |
| 1240 | { |
| 1241 | CancelPendingIo(Handle.Get(), Overlapped); |
| 1242 | } |
| 1243 | } |
| 1244 | |
| 1245 | void WriteHandle::SetCompleteOnDrained(bool Value) |
| 1246 | { |
| 1247 | CompleteOnDrained = Value; |
| 1248 | } |
| 1249 | |
| 1250 | IOHandleStatus WriteHandle::DrainedState() const |
| 1251 | { |
| 1252 | if (CompleteOnDrained) |
| 1253 | { |
| 1254 | return IOHandleStatus::Completed; |
| 1255 | } |
| 1256 | |
| 1257 | return Pending.empty() ? IOHandleStatus::Idle : IOHandleStatus::Standby; |
| 1258 | } |
| 1259 | |
| 1260 | void WriteHandle::Schedule() |
| 1261 | { |
| 1262 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 1263 | |
| 1264 | if (!Pending.empty()) |
| 1265 | { |
| 1266 | Buffer.Append(gsl::make_span(Pending)); |
| 1267 | Pending.clear(); |
| 1268 | } |
| 1269 | |
| 1270 | if (Buffer.Size() == 0) |
| 1271 | { |
| 1272 | State = DrainedState(); |
| 1273 | return; |
| 1274 | } |
| 1275 | |
| 1276 | Event.ResetEvent(); |
| 1277 | |
| 1278 | Overlapped.Offset = Offset.LowPart; |
| 1279 | Overlapped.OffsetHigh = Offset.HighPart; |
| 1280 | |
| 1281 | // Schedule the write. |
| 1282 | const auto buffer = Buffer.Span(); |
| 1283 | DWORD bytesWritten{}; |
| 1284 | if (WriteFile(Handle.Get(), buffer.data(), static_cast<DWORD>(buffer.size()), &bytesWritten, &Overlapped)) |
| 1285 | { |
| 1286 | Offset.QuadPart += bytesWritten; |
| 1287 | |
| 1288 | Buffer.Consume(bytesWritten); |
| 1289 | if (Buffer.Size() == 0) |
| 1290 | { |
| 1291 | State = DrainedState(); |
| 1292 | } |
| 1293 | } |
| 1294 | else |
| 1295 | { |
| 1296 | auto error = GetLastError(); |
| 1297 | THROW_LAST_ERROR_IF_MSG(error != ERROR_IO_PENDING, "Handle: 0x%p, size: %zu", (void*)Handle.Get(), buffer.size()); |
| 1298 | |
| 1299 | // The write is pending, update to 'Pending' |
| 1300 | State = IOHandleStatus::Pending; |
| 1301 | } |
| 1302 | } |
| 1303 | |
| 1304 | void WriteHandle::Collect() |
| 1305 | { |
| 1306 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 1307 | |
| 1308 | // Transition back to standby |
| 1309 | State = IOHandleStatus::Standby; |
| 1310 | |
| 1311 | // Complete the write. |
| 1312 | DWORD bytesWritten{}; |
| 1313 | THROW_IF_WIN32_BOOL_FALSE(GetOverlappedResult(Handle.Get(), &Overlapped, &bytesWritten, false)); |
| 1314 | Offset.QuadPart += bytesWritten; |
| 1315 | |
| 1316 | Buffer.Consume(bytesWritten); |
| 1317 | if (Buffer.Size() == 0) |
| 1318 | { |
| 1319 | State = DrainedState(); |
| 1320 | } |
| 1321 | } |
| 1322 | |
| 1323 | void WriteHandle::Push(const gsl::span<char>& Content) |
| 1324 | { |
| 1325 | WI_ASSERT(!Content.empty()); |
| 1326 | |
| 1327 | // Put any pending output to a different buffer, since the active buffer could be in the middle of a write. |
| 1328 | Pending.insert(Pending.end(), Content.begin(), Content.end()); |
| 1329 | |
| 1330 | if (State == IOHandleStatus::Idle) |
| 1331 | { |
| 1332 | State = IOHandleStatus::Standby; |
| 1333 | } |
| 1334 | } |
| 1335 | |
| 1336 | size_t WriteHandle::PendingBytes() const |
| 1337 | { |
| 1338 | return Pending.size() + Buffer.Size(); |
| 1339 | } |
| 1340 | |
| 1341 | HANDLE WriteHandle::GetHandle() const |
| 1342 | { |
| 1343 | return Event.get(); |
| 1344 | } |
| 1345 | |
| 1346 | WriteNamedPipe::WriteNamedPipe(HandleWrapper&& MovedPipe, bool Reconnect, bool Connected) : |
| 1347 | Pipe(std::move(MovedPipe)), ReconnectOnFailure(Reconnect), NeedConnect(!Connected) |
| 1348 | { |
| 1349 | ConnectOverlapped.hEvent = ConnectEvent.get(); |
| 1350 | |
| 1351 | Write.emplace(HandleWrapper{Pipe.Get()}, std::vector<char>{}, false); |
| 1352 | |
| 1353 | State = IOHandleStatus::Idle; |
| 1354 | } |
| 1355 | |
| 1356 | WriteNamedPipe::~WriteNamedPipe() |
| 1357 | { |
| 1358 | if (Connecting) |
| 1359 | { |
| 1360 | CancelPendingIo(Pipe.Get(), ConnectOverlapped); |
| 1361 | } |
| 1362 | } |
| 1363 | |
| 1364 | void WriteNamedPipe::Reconnect() |
| 1365 | { |
| 1366 | // Drop the disconnected client so a new one can connect, and retry the buffered data once reconnected. |
| 1367 | LOG_IF_WIN32_BOOL_FALSE(DisconnectNamedPipe(Pipe.Get())); |
| 1368 | |
| 1369 | NeedConnect = true; |
| 1370 | State = IOHandleStatus::Standby; |
| 1371 | } |
| 1372 | |
| 1373 | void WriteNamedPipe::Schedule() |
| 1374 | { |
| 1375 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 1376 | |
| 1377 | if (NeedConnect) |
| 1378 | { |
| 1379 | ConnectEvent.ResetEvent(); |
| 1380 | ConnectOverlapped.Offset = 0; |
| 1381 | ConnectOverlapped.OffsetHigh = 0; |
| 1382 | |
| 1383 | if (!ConnectNamedPipe(Pipe.Get(), &ConnectOverlapped)) |
| 1384 | { |
| 1385 | const auto error = GetLastError(); |
| 1386 | if (error == ERROR_IO_PENDING) |
| 1387 | { |
| 1388 | Connecting = true; |
| 1389 | State = IOHandleStatus::Pending; |
| 1390 | return; |
| 1391 | } |
| 1392 | |
| 1393 | THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != ERROR_PIPE_CONNECTED, "Handle: 0x%p", (void*)Pipe.Get()); |
| 1394 | } |
| 1395 | |
| 1396 | NeedConnect = false; |
| 1397 | } |
| 1398 | |
| 1399 | try |
| 1400 | { |
| 1401 | Write->Schedule(); |
| 1402 | State = Write->GetState(); |
| 1403 | } |
| 1404 | catch (...) |
| 1405 | { |
| 1406 | if (!ReconnectOnFailure) |
| 1407 | { |
| 1408 | throw; |
| 1409 | } |
| 1410 | |
| 1411 | LOG_CAUGHT_EXCEPTION(); |
| 1412 | Reconnect(); |
| 1413 | } |
| 1414 | } |
| 1415 | |
| 1416 | void WriteNamedPipe::Collect() |
| 1417 | { |
| 1418 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 1419 | |
| 1420 | // Complete a pending connection, then let the loop schedule the first write. |
| 1421 | if (Connecting) |
| 1422 | { |
| 1423 | Connecting = false; |
| 1424 | |
| 1425 | DWORD bytes{}; |
| 1426 | if (!GetOverlappedResult(Pipe.Get(), &ConnectOverlapped, &bytes, FALSE)) |
| 1427 | { |
| 1428 | const auto error = GetLastError(); |
| 1429 | THROW_HR_IF_MSG(HRESULT_FROM_WIN32(error), error != ERROR_PIPE_CONNECTED, "Handle: 0x%p", (void*)Pipe.Get()); |
| 1430 | } |
| 1431 | |
| 1432 | NeedConnect = false; |
| 1433 | State = IOHandleStatus::Standby; |
| 1434 | return; |
| 1435 | } |
| 1436 | |
| 1437 | try |
| 1438 | { |
| 1439 | Write->Collect(); |
| 1440 | State = Write->GetState(); |
| 1441 | } |
| 1442 | catch (...) |
| 1443 | { |
| 1444 | if (!ReconnectOnFailure) |
| 1445 | { |
| 1446 | throw; |
| 1447 | } |
| 1448 | |
| 1449 | LOG_CAUGHT_EXCEPTION(); |
| 1450 | Reconnect(); |
| 1451 | } |
| 1452 | } |
| 1453 | |
| 1454 | HANDLE WriteNamedPipe::GetHandle() const |
| 1455 | { |
| 1456 | return Connecting ? ConnectEvent.get() : Write->GetHandle(); |
| 1457 | } |
| 1458 | |
| 1459 | void WriteNamedPipe::Push(const gsl::span<char>& Content) |
| 1460 | { |
| 1461 | Write->Push(Content); |
| 1462 | |
| 1463 | if (State == IOHandleStatus::Idle) |
| 1464 | { |
| 1465 | State = IOHandleStatus::Standby; |
| 1466 | } |
| 1467 | } |
| 1468 | |
| 1469 | size_t WriteNamedPipe::PendingBytes() const |
| 1470 | { |
| 1471 | return Write ? Write->PendingBytes() : 0; |
| 1472 | } |
| 1473 | |
| 1474 | // DockerIORelayHandle |
| 1475 | |
| 1476 | DockerIORelayHandle::DockerIORelayHandle(HandleWrapper&& ReadHandle, HandleWrapper&& Stdout, HandleWrapper&& Stderr, Format ReadFormat) : |
| 1477 | WriteStdout(std::move(Stdout), {}, false), WriteStderr(std::move(Stderr), {}, false) |
| 1478 | { |
| 1479 | if (ReadFormat == Format::HttpChunked) |
| 1480 | { |
| 1481 | Read = std::make_unique<HTTPChunkBasedReadHandle>( |
| 1482 | std::move(ReadHandle), [this](const gsl::span<char>& Line) { this->OnRead(Line); }); |
| 1483 | } |
| 1484 | else |
| 1485 | { |
| 1486 | Read = |
| 1487 | std::make_unique<io::ReadHandle>(std::move(ReadHandle), [this](const gsl::span<char>& Buffer) { this->OnRead(Buffer); }); |
| 1488 | } |
| 1489 | } |
| 1490 | |
| 1491 | void DockerIORelayHandle::Schedule() |
| 1492 | { |
| 1493 | WI_ASSERT(State == IOHandleStatus::Standby); |
| 1494 | WI_ASSERT(Read->GetState() != IOHandleStatus::Pending); |
| 1495 | |
| 1496 | // If we have an active handle and a buffer, try to flush that first. |
| 1497 | if (ActiveHandle != nullptr && !PendingBuffer.empty()) |
| 1498 | { |
| 1499 | // Push the data to the selected handle. |
| 1500 | DWORD bytesToWrite = std::min(static_cast<DWORD>(RemainingBytes), static_cast<DWORD>(PendingBuffer.size())); |
| 1501 | |
| 1502 | ActiveHandle->Push(gsl::make_span(PendingBuffer.data(), bytesToWrite)); |
| 1503 | |
| 1504 | // Consume the written bytes. |
| 1505 | RemainingBytes -= bytesToWrite; |
| 1506 | PendingBuffer.erase(PendingBuffer.begin(), PendingBuffer.begin() + bytesToWrite); |
| 1507 | |
| 1508 | // Schedule the write. |
| 1509 | ActiveHandle->Schedule(); |
| 1510 | |
| 1511 | // If the write is pending, update to 'Pending' |
| 1512 | if (ActiveHandle->GetState() == IOHandleStatus::Pending) |
| 1513 | { |
| 1514 | State = IOHandleStatus::Pending; |
| 1515 | } |
| 1516 | else if (ActiveHandle->GetState() == IOHandleStatus::Completed || ActiveHandle->GetState() == IOHandleStatus::Idle) |
| 1517 | { |
| 1518 | if (RemainingBytes == 0) |
| 1519 | { |
| 1520 | // Switch back to reading if we've written all bytes for this chunk. |
| 1521 | ActiveHandle = nullptr; |
| 1522 | |
| 1523 | ProcessNextHeader(); |
| 1524 | } |
| 1525 | } |
| 1526 | } |
| 1527 | else |
| 1528 | { |
| 1529 | if (Read->GetState() == IOHandleStatus::Completed) |
| 1530 | { |
| 1531 | LOG_HR_IF(E_UNEXPECTED, ActiveHandle != nullptr); |
| 1532 | |
| 1533 | // No more data to read, we're done. |
| 1534 | State = IOHandleStatus::Completed; |
| 1535 | return; |
| 1536 | } |
| 1537 | |
| 1538 | // Schedule a read from the input. |
| 1539 | Read->Schedule(); |
| 1540 | if (Read->GetState() == IOHandleStatus::Pending) |
| 1541 | { |
| 1542 | State = IOHandleStatus::Pending; |
| 1543 | } |
| 1544 | } |
| 1545 | } |
| 1546 | |
| 1547 | void DockerIORelayHandle::Collect() |
| 1548 | { |
| 1549 | WI_ASSERT(State == IOHandleStatus::Pending); |
| 1550 | |
| 1551 | if (ActiveHandle != nullptr && ActiveHandle->GetState() == IOHandleStatus::Pending) |
| 1552 | { |
| 1553 | // Complete the write. |
| 1554 | ActiveHandle->Collect(); |
| 1555 | |
| 1556 | // If the write is completed, switch back to reading. |
| 1557 | if (RemainingBytes == 0) |
| 1558 | { |
| 1559 | if (ActiveHandle->GetState() == IOHandleStatus::Completed || ActiveHandle->GetState() == IOHandleStatus::Idle) |
| 1560 | { |
| 1561 | ActiveHandle = nullptr; |
| 1562 | |
| 1563 | ProcessNextHeader(); |
| 1564 | } |
| 1565 | } |
| 1566 | |
| 1567 | // Transition back to standby if there's still data to read. |
| 1568 | // Otherwise switch to Completed since everything is done. |
| 1569 | if (Read->GetState() == IOHandleStatus::Completed) |
| 1570 | { |
| 1571 | LOG_HR_IF(E_UNEXPECTED, RemainingBytes != 0); |
| 1572 | |
| 1573 | State = IOHandleStatus::Completed; |
| 1574 | } |
| 1575 | else |
| 1576 | { |
| 1577 | State = IOHandleStatus::Standby; |
| 1578 | } |
| 1579 | } |
| 1580 | else |
| 1581 | { |
| 1582 | WI_ASSERT(Read->GetState() == IOHandleStatus::Pending); |
| 1583 | |
| 1584 | // Complete the read. |
| 1585 | Read->Collect(); |
| 1586 | |
| 1587 | // Transition back to standby. |
| 1588 | State = IOHandleStatus::Standby; |
| 1589 | } |
| 1590 | } |
| 1591 | |
| 1592 | HANDLE DockerIORelayHandle::GetHandle() const |
| 1593 | { |
| 1594 | if (ActiveHandle != nullptr && ActiveHandle->GetState() == IOHandleStatus::Pending) |
| 1595 | { |
| 1596 | return ActiveHandle->GetHandle(); |
| 1597 | } |
| 1598 | else |
| 1599 | { |
| 1600 | return Read->GetHandle(); |
| 1601 | } |
| 1602 | } |
| 1603 | |
| 1604 | void DockerIORelayHandle::ProcessNextHeader() |
| 1605 | { |
| 1606 | if (PendingBuffer.size() < sizeof(MultiplexedHeader)) |
| 1607 | { |
| 1608 | // Not enough data for a header yet. |
| 1609 | return; |
| 1610 | } |
| 1611 | |
| 1612 | const auto* header = reinterpret_cast<const MultiplexedHeader*>(PendingBuffer.data()); |
| 1613 | RemainingBytes = ntohl(header->Length); |
| 1614 | |
| 1615 | if (header->Fd == 1) |
| 1616 | { |
| 1617 | ActiveHandle = &WriteStdout; |
| 1618 | } |
| 1619 | else if (header->Fd == 2) |
| 1620 | { |
| 1621 | ActiveHandle = &WriteStderr; |
| 1622 | } |
| 1623 | else |
| 1624 | { |
| 1625 | THROW_HR_MSG(E_INVALIDARG, "Invalid Docker IO multiplexed header fd: %u", header->Fd); |
| 1626 | } |
| 1627 | |
| 1628 | // Consume the header. |
| 1629 | PendingBuffer.erase(PendingBuffer.begin(), PendingBuffer.begin() + sizeof(MultiplexedHeader)); |
| 1630 | } |
| 1631 | |
| 1632 | void DockerIORelayHandle::OnRead(const gsl::span<char>& Buffer) |
| 1633 | { |
| 1634 | PendingBuffer.insert(PendingBuffer.end(), Buffer.begin(), Buffer.end()); |
| 1635 | |
| 1636 | if (ActiveHandle == nullptr) |
| 1637 | { |
| 1638 | // If no handle is active, expect a header. |
| 1639 | ProcessNextHeader(); |
| 1640 | } |
| 1641 | } |
| 1642 | |
| 1643 | // MultiHandleWait |
| 1644 | |
| 1645 | MultiHandleWait::MultiHandleWait(MultiHandleWait&& other) noexcept |
| 1646 | { |
| 1647 | *this = std::move(other); |
| 1648 | } |
| 1649 | |
| 1650 | MultiHandleWait& MultiHandleWait::operator=(MultiHandleWait&& other) noexcept |
| 1651 | { |
| 1652 | if (this != &other) |
| 1653 | { |
| 1654 | m_handles = std::move(other.m_handles); |
| 1655 | m_handleSignaledEvent = std::move(other.m_handleSignaledEvent); |
| 1656 | m_cancel = other.m_cancel; |
| 1657 | |
| 1658 | for (auto& entry : m_handles) |
| 1659 | { |
| 1660 | entry->Self = this; |
| 1661 | } |
| 1662 | |
| 1663 | // N.B. moving a MultiHandleWait() while running is not supported |
| 1664 | WI_ASSERT(m_signaledHandles.empty()); |
| 1665 | } |
| 1666 | |
| 1667 | return *this; |
| 1668 | } |
| 1669 | |
| 1670 | void MultiHandleWait::AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags, OnError&& onError) |
| 1671 | { |
| 1672 | auto entry = std::make_unique<Entry>(); |
| 1673 | entry->HandleFlags = flags; |
| 1674 | entry->Handle = std::move(handle); |
| 1675 | entry->Self = this; |
| 1676 | |
| 1677 | if (WI_IsFlagSet(flags, Flags::IgnoreErrors)) |
| 1678 | { |
| 1679 | entry->ErrorCallback = []() {}; |
| 1680 | } |
| 1681 | else |
| 1682 | { |
| 1683 | entry->ErrorCallback = std::move(onError); |
| 1684 | } |
| 1685 | m_handles.emplace_back(std::move(entry)); |
| 1686 | } |
| 1687 | |
| 1688 | void MultiHandleWait::Cancel() |
| 1689 | { |
| 1690 | m_cancel = true; |
| 1691 | } |
| 1692 | |
| 1693 | void NTAPI MultiHandleWait::WaitCallback(PVOID Context, BOOLEAN /*TimerOrWaitFired*/) |
| 1694 | { |
| 1695 | auto* entry = static_cast<Entry*>(Context); |
| 1696 | |
| 1697 | entry->Self->m_signaledHandles.push(entry); |
| 1698 | entry->Self->m_handleSignaledEvent.SetEvent(); |
| 1699 | } |
| 1700 | |
| 1701 | bool MultiHandleWait::Run(std::optional<std::chrono::milliseconds> Timeout) |
| 1702 | { |
| 1703 | m_cancel = false; // Run may be called multiple times. |
| 1704 | |
| 1705 | std::optional<std::chrono::steady_clock::time_point> deadline; |
| 1706 | if (Timeout.has_value()) |
| 1707 | { |
| 1708 | deadline = std::chrono::steady_clock::now() + Timeout.value(); |
| 1709 | } |
| 1710 | |
| 1711 | std::vector<unique_registered_wait> callbacks; |
| 1712 | |
| 1713 | while (!m_cancel) |
| 1714 | { |
| 1715 | // Cancel any pending callback. |
| 1716 | callbacks.clear(); |
| 1717 | |
| 1718 | Entry* signaledEntry = nullptr; |
| 1719 | while (m_signaledHandles.try_pop(signaledEntry)) |
| 1720 | { |
| 1721 | try |
| 1722 | { |
| 1723 | signaledEntry->Handle->Collect(); |
| 1724 | } |
| 1725 | catch (...) |
| 1726 | { |
| 1727 | signaledEntry->ErrorCallback(); // Might throw and cancel the IO. |
| 1728 | signaledEntry->Handle.reset(); |
| 1729 | continue; |
| 1730 | } |
| 1731 | } |
| 1732 | |
| 1733 | m_handleSignaledEvent.ResetEvent(); |
| 1734 | |
| 1735 | bool hasHandleToWaitFor = false; |
| 1736 | for (auto it = m_handles.begin(); it != m_handles.end();) |
| 1737 | { |
| 1738 | auto& entry = **it; |
| 1739 | |
| 1740 | while (entry.Handle && entry.Handle->GetState() == IOHandleStatus::Standby && !m_cancel) |
| 1741 | { |
| 1742 | try |
| 1743 | { |
| 1744 | entry.Handle->Schedule(); |
| 1745 | } |
| 1746 | catch (...) |
| 1747 | { |
| 1748 | entry.ErrorCallback(); // Might throw and cancel the IO. |
| 1749 | entry.Handle.reset(); |
| 1750 | break; |
| 1751 | } |
| 1752 | } |
| 1753 | |
| 1754 | if (!entry.Handle || entry.Handle->GetState() == IOHandleStatus::Completed) |
| 1755 | { |
| 1756 | if (entry.Handle && WI_IsFlagSet(entry.HandleFlags, Flags::CancelOnCompleted)) |
| 1757 | { |
| 1758 | m_cancel = true; |
| 1759 | } |
| 1760 | |
| 1761 | it = m_handles.erase(it); |
| 1762 | continue; |
| 1763 | } |
| 1764 | |
| 1765 | // N.B. An Idle handle cannot be waited for since it's not doing any IO. |
| 1766 | if (entry.Handle->GetState() == IOHandleStatus::Idle) |
| 1767 | { |
| 1768 | ++it; |
| 1769 | continue; |
| 1770 | } |
| 1771 | |
| 1772 | auto& callback = callbacks.emplace_back(); |
| 1773 | |
| 1774 | THROW_IF_WIN32_BOOL_FALSE(RegisterWaitForSingleObject( |
| 1775 | &callback, entry.Handle->GetHandle(), &WaitCallback, &entry, INFINITE, WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE)); |
| 1776 | |
| 1777 | if (WI_IsFlagClear(entry.HandleFlags, Flags::NeedNotComplete)) |
| 1778 | { |
| 1779 | hasHandleToWaitFor = true; |
| 1780 | } |
| 1781 | |
| 1782 | ++it; |
| 1783 | } |
| 1784 | |
| 1785 | if (m_handles.empty() || !hasHandleToWaitFor || m_cancel) |
| 1786 | { |
| 1787 | break; |
| 1788 | } |
| 1789 | |
| 1790 | DWORD waitTimeout = INFINITE; |
| 1791 | if (deadline.has_value()) |
| 1792 | { |
| 1793 | auto milliseconds = |
| 1794 | std::chrono::duration_cast<std::chrono::milliseconds>(deadline.value() - std::chrono::steady_clock::now()).count(); |
| 1795 | |
| 1796 | waitTimeout = static_cast<DWORD>(std::max<long long>(0, milliseconds)); |
| 1797 | } |
| 1798 | |
| 1799 | THROW_HR_IF_MSG( |
| 1800 | HRESULT_FROM_WIN32(ERROR_TIMEOUT), |
| 1801 | !m_handleSignaledEvent.wait(waitTimeout), |
| 1802 | "Timed out waiting for %llu handles. Timeout: %lu", |
| 1803 | m_handles.size(), |
| 1804 | waitTimeout); |
| 1805 | } |
| 1806 | |
| 1807 | return !m_cancel; |
| 1808 | } |