@samitouri / QOSAMI-WSL
Code Commits Issues Pulls Sessions CI/CD
master
master
hpp 156 lines 5.37 KB
Raw
1 /*++
2
3 Copyright (c) Microsoft. All rights reserved.
4
5 Module Name:
6
7 relay.hpp
8
9 Abstract:
10
11 This file contains function declarations for the relay worker thread routines.
12
13 --*/
14
15 #pragma once
16
17 #include <winsock2.h>
18 #include "ConsoleState.h"
19 #include "HandleIO.h"
20
21 namespace wsl::windows::common::relay {
22
23 using namespace wsl::windows::common::io;
24
25 std::thread CreateThread(_In_ HANDLE InputHandle, _In_ HANDLE OutputHandle, _In_opt_ HANDLE ExitHandle = nullptr, _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE);
26
27 std::thread CreateThread(_In_ wil::unique_handle&& InputHandle, _In_ HANDLE OutputHandle, _In_opt_ HANDLE ExitHandle = nullptr, _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE);
28
29 std::thread CreateThread(_In_ HANDLE InputHandle, _In_ wil::unique_handle&& OutputHandle, _In_opt_ HANDLE ExitHandle = nullptr, _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE);
30
31 std::thread CreateThread(
32 _In_ wil::unique_handle&& InputHandle,
33 _In_ wil::unique_handle&& OutputHandle,
34 _In_opt_ HANDLE ExitHandle = nullptr,
35 _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE);
36
37 DWORD
38 InterruptableRead(_In_ HANDLE InputHandle, _In_ gsl::span<gsl::byte> Buffer, _In_ const std::vector<HANDLE>& ExitHandles, _In_opt_ LPOVERLAPPED Overlapped = nullptr);
39
40 void InterruptableRelay(_In_ HANDLE InputHandle, _In_opt_ HANDLE OutputHandle, _In_opt_ HANDLE ExitHandle = nullptr, _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE);
41
42 bool InterruptableWait(_In_ HANDLE WaitObject, _In_ const std::vector<HANDLE>& ExitHandles = {});
43
44 DWORD
45 InterruptableWrite(_In_ HANDLE OutputHandle, _In_ gsl::span<const gsl::byte> Buffer, _In_ const std::vector<HANDLE>& ExitHandles, _In_ LPOVERLAPPED Overlapped);
46
47 bool StandardInputRelay(HANDLE ConsoleHandle, HANDLE OutputHandle, std::function<void()>&& UpdateTerminalSize, HANDLE ExitEvent);
48
49 enum class RelayFlags
50 {
51 None = 0,
52 LeftIsSocket = 1,
53 RightIsSocket = 2
54 };
55
56 DEFINE_ENUM_FLAG_OPERATORS(RelayFlags);
57
58 void BidirectionalRelay(_In_ HANDLE LeftHandle, _In_ HANDLE RightHandle, _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE, _In_ RelayFlags Flags = RelayFlags::None);
59
60 void SocketRelay(_In_ SOCKET LeftSocket, _In_ SOCKET RightSocket, _In_ size_t BufferSize = LX_RELAY_BUFFER_SIZE);
61
62 class ScopedMultiRelay
63 {
64 public:
65 using TWriteMethod = std::function<void(size_t, const gsl::span<gsl::byte>& buffer)>;
66 ScopedMultiRelay(const std::vector<HANDLE>& Inputs, const TWriteMethod& Write, size_t BufferSize = LX_RELAY_BUFFER_SIZE);
67
68 ~ScopedMultiRelay();
69
70 ScopedMultiRelay(ScopedMultiRelay&& other) = default;
71 ScopedMultiRelay(const ScopedMultiRelay&) = delete;
72
73 ScopedMultiRelay& operator=(const ScopedMultiRelay&) = delete;
74 ScopedMultiRelay& operator=(ScopedMultiRelay&&) = delete;
75
76 // Blocks until the relaying is complete.
77 // This is useful for situations where the relay should make sure that all
78 // the content has been flushed before exiting.
79 void Sync();
80
81 private:
82 void Run(const std::vector<HANDLE>& Inputs, const TWriteMethod& Write, size_t BufferSize = LX_RELAY_BUFFER_SIZE) const;
83
84 std::thread m_thread;
85 wil::unique_event m_exitEvent{wil::EventOptions::ManualReset};
86 };
87
88 // Helper class to relay the output of a handle to another.
89 // Note: The relay can take ownership of the handles if desired.
90 // Doing that will cause the handle to be released when the relaying is complete.
91
92 class ScopedRelay
93 {
94 public:
95 template <typename TInput, typename TOutput>
96 ScopedRelay(
97 TInput&& Input, TOutput&& Output, size_t BufferSize = LX_RELAY_BUFFER_SIZE, std::function<void()>&& OnDestroy = []() {}) :
98 m_onDestroy(std::move(OnDestroy))
99 {
100 m_thread = std::thread{[this, Input = std::move(Input), Output = std::move(Output), BufferSize = BufferSize]() {
101 try
102 {
103 Run(GetUnderlyingHandle(Input), GetUnderlyingHandle(Output), BufferSize);
104 }
105 CATCH_LOG();
106 }};
107 }
108
109 ~ScopedRelay();
110
111 ScopedRelay(ScopedRelay&& other) = default;
112 ScopedRelay(const ScopedRelay&) = delete;
113
114 ScopedRelay& operator=(const ScopedRelay&) = delete;
115 ScopedRelay& operator=(ScopedRelay&&) = delete;
116
117 // Blocks until the relaying is complete.
118 // This is useful for situations where the relay should make sure that all
119 // the content has been flushed before exiting.
120 void Sync();
121
122 private:
123 template <typename THandle>
124 static HANDLE GetUnderlyingHandle(THandle& handle)
125 {
126 if constexpr (std::is_same_v<std::remove_cv_t<THandle>, HANDLE>)
127 {
128 return handle;
129 }
130 else if constexpr (std::is_same_v<std::remove_cv_t<THandle>, wil::unique_handle>)
131 {
132 return handle.get();
133 }
134 else if constexpr (std::is_same_v<std::remove_cv_t<THandle>, wil::unique_socket>)
135 {
136 return reinterpret_cast<HANDLE>(handle.get());
137 }
138 else if constexpr (std::is_same_v<std::remove_cv_t<THandle>, SOCKET>)
139 {
140 return reinterpret_cast<HANDLE>(handle);
141 }
142 else
143 {
144 // If this assert fails, an invalid type was passed to ScopedRelay
145 static_assert(sizeof(THandle) != sizeof(THandle));
146 }
147 }
148
149 void Run(_In_ HANDLE Input, _In_ HANDLE Output, size_t BufferSize) const;
150
151 std::thread m_thread;
152 wil::unique_event m_exitEvent{wil::EventOptions::ManualReset};
153 std::function<void()> m_onDestroy;
154 };
155
156 } // namespace wsl::windows::common::relay