@samitouri / QOSAMI-WSL-2
Code Commits Issues Pulls Sessions CI/CD
master
master
h 146 lines 5.68 KB
Raw
1 // Copyright (C) Microsoft Corporation. All rights reserved.
2
3 #pragma once
4
5 #include <chrono>
6 #include <list>
7 #include <map>
8 #include "common.h"
9 #include "lxinitshared.h"
10
11 using DnsTunnelingCallback = std::function<void(const gsl::span<gsl::byte>, const LX_GNS_DNS_CLIENT_IDENTIFIER&)>;
12
13 // Number of bytes used to store the length of DNS over TCP requests
14 constexpr int c_byteCountTcpRequestLength = 2;
15
16 class DnsServer
17 {
18 public:
19 DnsServer(DnsTunnelingCallback&& tunnelDnsRequest);
20 ~DnsServer() noexcept;
21
22 DnsServer(const DnsServer&) = delete;
23 DnsServer(DnsServer&&) = delete;
24 DnsServer& operator=(const DnsServer&) = delete;
25 DnsServer& operator=(DnsServer&&) = delete;
26
27 // Start DNS server.
28 //
29 // Arguments:
30 // ipAddress - IP address to start server on.
31 void Start(const std::string& ipAddress) noexcept;
32
33 // Process DNS response received from Windows.
34 //
35 // Arguments:
36 // dnsBuffer - buffer containing DNS response.
37 // dnsClientIdentifier - struct containing protocol (TCP/UDP) and unique id of the Linux DNS client making the request.
38 void HandleDnsResponse(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept;
39
40 void Stop() noexcept;
41
42 private:
43 struct TcpConnectionContext
44 {
45 // Connection fd
46 wil::unique_fd m_tcpConnection;
47
48 // Offset in m_currentDnsRequest indicating how much of the current DNS request on
49 // the TCP connection has been read.
50 size_t m_currentRequestOffset = 0;
51
52 // Buffer containing the current DNS request received on the TCP connection.
53 std::vector<gsl::byte> m_currentDnsRequest;
54
55 // Unique connection id. The connection fd would be a candidate for this, but the fd might be reused, so we need a different id.
56 uint32_t m_connectionId{};
57
58 TcpConnectionContext(uint32_t id, wil::unique_fd&& tcpConnection) :
59 m_tcpConnection(std::move(tcpConnection)), m_connectionId(id)
60 {
61 // Resize to fit the bytes that represent the request length
62 m_currentDnsRequest.resize(c_byteCountTcpRequestLength);
63 }
64
65 ~TcpConnectionContext() noexcept = default;
66
67 TcpConnectionContext(const TcpConnectionContext&) = delete;
68 TcpConnectionContext& operator=(const TcpConnectionContext&) = delete;
69 TcpConnectionContext(TcpConnectionContext&&) = delete;
70 TcpConnectionContext& operator=(TcpConnectionContext&&) = delete;
71 };
72
73 using UdpRequestExpirationQueue = std::list<std::pair<std::chrono::steady_clock::time_point, uint32_t>>;
74
75 struct UdpRequestContext
76 {
77 sockaddr_in m_remoteAddress;
78 UdpRequestExpirationQueue::iterator m_expiration;
79 };
80
81 void StartUdpDnsServer(const std::string& ipAddress) noexcept;
82
83 void StartTcpDnsServer(const std::string& ipAddress) noexcept;
84
85 // Main server loop, processing epoll notifications.
86 void ServerLoop() noexcept;
87
88 // Accept new incoming TCP connection.
89 void HandleNewTcpConnection() noexcept;
90
91 // Handle new data received on an existing TCP connection.
92 void HandleNewTcpData(TcpConnectionContext* context) noexcept;
93
94 // Read the next DNS request from the UDP socket.
95 void HandleUdpDnsRequest() noexcept;
96
97 int ExpireUdpRequestsAndGetTimeout() noexcept;
98
99 void HandleUdpDnsResponse(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept;
100
101 void HandleTcpDnsResponse(const gsl::span<gsl::byte> dnsBuffer, const LX_GNS_DNS_CLIENT_IDENTIFIER& dnsClientIdentifier) noexcept;
102
103 // File descriptor used to interact with epoll. Declared before the UDP and TCP sockets and the shutdown pipe so it will be closed after them.
104 // Note: Closing a socket fd automatically leads to unregistering it from epoll - EPOLL_CTL_DEL is not necessary for that fd.
105 wil::unique_fd m_epollFd;
106
107 std::mutex m_udpLock;
108
109 // _Guarded_by_(m_udpLock)
110 wil::unique_fd m_udpSocket;
111
112 // Unique id that is incremented for each DNS request over UDP. In case the value reaches MAX_UINT and is reset to 0,
113 // it's assumed previous requests with id's 0, 1, ... finished in the meantime and the id can be reused.
114 // _Guarded_by_(m_udpLock)
115 uint32_t m_currentUdpRequestId = 0;
116
117 // Mapping id of an UDP DNS request to the sockaddr_in struct storing the IP and port used by the Linux DNS client that made
118 // the DNS request. Note: Since we only configure an IPv4 DNS server in Linux, we expect all Linux DNS clients to use IPv4
119 // addresses. _Guarded_by_(m_udpLock)
120 std::map<uint32_t, UdpRequestContext> m_udpRequests;
121
122 // UDP requests ordered by expiration time. _Guarded_by_(m_udpLock)
123 UdpRequestExpirationQueue m_udpRequestExpirations;
124
125 wil::unique_fd m_tcpListenSocket;
126
127 std::mutex m_tcpLock;
128
129 // Unique id that is incremented for each TCP connection. In case the value reaches MAX_UINT and is reset to 0,
130 // it's assumed previous connections with id's 0, 1, ... were closed in the meantime and the id can be reused.
131 // _Guarded_by_(m_udpLock)
132 uint32_t m_currentTcpConnectionId = 0;
133
134 // Mapping TCP connection unique id to connection context
135 // _Guarded_by_(m_udpLock)
136 std::map<uint32_t, std::unique_ptr<TcpConnectionContext>> m_tcpConnectionContexts;
137
138 // Pipe used to stop m_serverThread.
139 wil::unique_pipe m_shutdownServerLoopPipe;
140
141 // Thread running the server loop.
142 std::thread m_serverThread;
143
144 // Callback used for tunneling a DNS request to Windows to be resolved.
145 DnsTunnelingCallback m_tunnelDnsRequest;
146 };