@samitouri / QOSAMI-WSL-1
Code Commits Issues Pulls Sessions CI/CD
master
master
cpp 157 lines 5.17 KB
Raw
1 // Copyright (C) Microsoft Corporation. All rights reserved.
2
3 #include <netinet/in.h>
4 #include <arpa/inet.h>
5 #include "DnsTunnelingChannel.h"
6 #include "util.h"
7 #include "RuntimeErrorWithSourceLocation.h"
8 #include "Syscall.h"
9 #include "message.h"
10
11 DnsTunnelingChannel::DnsTunnelingChannel(int channelFd, DnsTunnelingCallback&& reportDnsResponse) :
12 m_channel(wil::unique_fd{channelFd}, "DnsTunneling"), m_reportDnsResponse(std::move(reportDnsResponse))
13 {
14 // Create a pipe to be used for signalling the receive loop to stop
15 m_shutdownReceiveWorkerPipe = wil::unique_pipe::create(0);
16
17 // Start loop waiting for incoming messages from Windows side
18 m_receiveWorkerThread = std::thread([this]() { ReceiveLoop(); });
19 }
20
21 DnsTunnelingChannel::~DnsTunnelingChannel()
22 {
23 Stop();
24 }
25
26 void DnsTunnelingChannel::SendDnsMessage(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept
27 try
28 {
29 wsl::shared::MessageWriter<LX_GNS_DNS_TUNNELING_MESSAGE> message(LxGnsMessageDnsTunneling);
30 message->DnsClientIdentifier = dnsClientIdentifier;
31 message.WriteSpan(dnsBuffer);
32
33 m_channel.SendMessage<LX_GNS_DNS_TUNNELING_MESSAGE>(message.Span());
34 }
35 CATCH_LOG()
36
37 void DnsTunnelingChannel::ReceiveLoop() noexcept
38 {
39 UtilSetThreadName("DnsTunneling");
40
41 // Returns false if the write pipe was closed, signaling that loop should exit
42 // Returns true if there is data to be received on the channel fd
43 auto wait_for_channel_fd = [this]() -> bool {
44 struct pollfd poll_fds[2];
45 poll_fds[0] = {.fd = m_channel.Socket(), .events = POLLIN, .revents = 0};
46 poll_fds[1] = {.fd = m_shutdownReceiveWorkerPipe.read().get(), .events = POLLIN, .revents = 0};
47
48 unsigned int retryCount = 0;
49 const unsigned int maxRetryCount = 3;
50
51 for (;;)
52 {
53 int return_value = SyscallInterruptable(poll, poll_fds, ARRAY_SIZE(poll_fds), -1);
54 if (return_value < 0)
55 {
56 GNS_LOG_ERROR("poll failed");
57 retryCount++;
58
59 if (retryCount < maxRetryCount)
60 {
61 continue;
62 }
63 else
64 {
65 return false;
66 }
67 }
68 else if (return_value == 0)
69 {
70 GNS_LOG_ERROR("poll returned 0 (timeout)");
71 return false;
72 }
73 else if (poll_fds[1].revents)
74 {
75 return false;
76 }
77 else if (poll_fds[0].revents & POLLIN)
78 {
79 return true;
80 }
81 }
82 };
83
84 std::vector<gsl::byte> receiveBuffer;
85
86 for (;;)
87 {
88 try
89 {
90 if (!wait_for_channel_fd())
91 {
92 break;
93 }
94
95 GNS_LOG_INFO("processing next message from Windows");
96
97 // Read next message. wsl::shared::socket::RecvMessage() first reads the message header, then uses it to determine the
98 // total size of the message and read the rest of the message, resizing the buffer if needed.
99 auto [message, span] = m_channel.ReceiveMessageOrClosed<MESSAGE_HEADER>();
100 if (message == nullptr)
101 {
102 GNS_LOG_ERROR("failed to read message");
103 return;
104 }
105
106 // Get the message type from the message header
107 switch (message->MessageType)
108 {
109 case LxGnsMessageDnsTunneling:
110 {
111 // Cast message to a LX_GNS_DNS_TUNNELING_MESSAGE struct
112 auto* dnsMessage = gslhelpers::try_get_struct<LX_GNS_DNS_TUNNELING_MESSAGE>(span);
113 if (!dnsMessage)
114 {
115 GNS_LOG_ERROR("failed to convert message to LX_GNS_DNS_TUNNELING_MESSAGE");
116 return;
117 }
118
119 // Extract DNS buffer from message
120 auto dnsBuffer = span.subspan(offsetof(LX_GNS_DNS_TUNNELING_MESSAGE, Buffer));
121
122 GNS_LOG_INFO(
123 "received DNS message DNS buffer size: {}, Protocol {}, DNS client id: {}",
124 dnsBuffer.size(),
125 dnsMessage->DnsClientIdentifier.Protocol == IPPROTO_UDP ? "UDP" : "TCP",
126 dnsMessage->DnsClientIdentifier.DnsClientId);
127
128 // Invoke callback to notify about the new DNS response
129 m_reportDnsResponse(dnsBuffer, dnsMessage->DnsClientIdentifier);
130
131 break;
132 }
133
134 default:
135 {
136 throw RuntimeErrorWithSourceLocation(std::format("Unexpected LX_MESSAGE_TYPE : {}", static_cast<int>(message->MessageType)));
137 }
138 }
139 }
140 CATCH_LOG()
141 }
142 }
143
144 void DnsTunnelingChannel::Stop() noexcept
145 try
146 {
147 GNS_LOG_INFO("stopping DNS server");
148
149 // Stop receive loop by closing the write fd of the pipe
150 m_shutdownReceiveWorkerPipe.write().reset();
151
152 if (m_receiveWorkerThread.joinable())
153 {
154 m_receiveWorkerThread.join();
155 }
156 }
157 CATCH_LOG()