1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
|
// Copyright Epic Games, Inc. All Rights Reserved.
#pragma once
#include "httptrafficinspector.h"
#include "httptrafficrecorder.h"
#include <zencore/logging.h>
#include <zencore/thread.h>
#include <zentelemetry/stats.h>
ZEN_THIRD_PARTY_INCLUDES_START
#include <asio.hpp>
#if defined(ASIO_HAS_LOCAL_SOCKETS)
# include <asio/local/stream_protocol.hpp>
#endif
ZEN_THIRD_PARTY_INCLUDES_END
#include <atomic>
#include <chrono>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>
namespace zen {
struct ProxyMapping
{
std::string ListenAddress;
uint16_t ListenPort = 0;
std::string ListenUnixSocket;
std::string TargetHost;
uint16_t TargetPort = 0;
std::string TargetUnixSocket;
bool IsUnixListen() const { return !ListenUnixSocket.empty(); }
bool IsUnixTarget() const { return !TargetUnixSocket.empty(); }
std::string ListenDescription() const;
std::string TargetDescription() const;
};
class TcpProxyService;
class TcpProxySession : public std::enable_shared_from_this<TcpProxySession>
{
public:
TcpProxySession(asio::ip::tcp::socket ClientSocket, const ProxyMapping& Mapping, TcpProxyService& Owner);
#if defined(ASIO_HAS_LOCAL_SOCKETS)
TcpProxySession(asio::local::stream_protocol::socket ClientSocket, const ProxyMapping& Mapping, TcpProxyService& Owner);
#endif
void Start();
const std::string& GetClientLabel() const { return m_ClientLabel; }
std::chrono::steady_clock::time_point GetStartTime() const { return m_StartTime; }
uint64_t GetBytesFromClient() const { return m_BytesFromClient.load(std::memory_order_relaxed); }
uint64_t GetBytesToClient() const { return m_BytesToClient.load(std::memory_order_relaxed); }
uint64_t GetRequestCount() const { return m_RequestInspector ? m_RequestInspector->GetMessageCount() : 0; }
bool IsWebSocket() const { return m_RequestInspector && m_RequestInspector->IsUpgraded(); }
bool HasSessionId() const { return m_RequestInspector && m_RequestInspector->HasSessionId(); }
Oid GetSessionId() const { return m_RequestInspector ? m_RequestInspector->GetSessionId() : Oid::Zero; }
private:
LoggerRef Log();
void ConnectToTcpTarget();
#if defined(ASIO_HAS_LOCAL_SOCKETS)
void ConnectToUnixTarget();
#endif
void StartRelay();
void ReadFromClient();
void ReadFromUpstream();
template<typename Fn>
void DispatchClientSocket(Fn&& F);
template<typename Fn>
void DispatchUpstreamSocket(Fn&& F);
template<typename SocketT>
void DoReadFromClient(SocketT& ClientSocket);
template<typename SocketT>
void DoReadFromUpstream(SocketT& UpstreamSocket);
template<typename SocketT>
void DoForwardToUpstream(SocketT& UpstreamSocket, size_t BytesToWrite, uint64_t NewRequests);
template<typename SocketT>
void DoForwardToClient(SocketT& ClientSocket, size_t BytesToWrite);
template<typename SocketT>
void DoShutdownSocket(SocketT& Socket);
void Shutdown();
asio::ip::tcp::socket m_ClientTcpSocket;
asio::ip::tcp::socket m_UpstreamTcpSocket;
#if defined(ASIO_HAS_LOCAL_SOCKETS)
asio::local::stream_protocol::socket m_ClientUnixSocket;
asio::local::stream_protocol::socket m_UpstreamUnixSocket;
bool m_IsUnixClient = false;
bool m_IsUnixTarget = false;
#endif
std::string m_TargetHost;
uint16_t m_TargetPort;
std::string m_TargetUnixSocket;
TcpProxyService& m_Owner;
static constexpr size_t kBufferSize = 16 * 1024;
std::array<char, kBufferSize> m_ClientBuffer;
std::array<char, kBufferSize> m_UpstreamBuffer;
std::atomic<bool> m_ShutdownCalled{false};
std::string m_ClientLabel;
std::chrono::steady_clock::time_point m_StartTime;
std::atomic<uint64_t> m_BytesFromClient{0};
std::atomic<uint64_t> m_BytesToClient{0};
std::optional<HttpTrafficInspector> m_RequestInspector;
std::optional<HttpTrafficInspector> m_ResponseInspector;
std::unique_ptr<HttpTrafficRecorder> m_Recorder;
};
class TcpProxyService
{
public:
TcpProxyService(asio::io_context& IoContext, const ProxyMapping& Mapping);
void Start();
void Stop();
const ProxyMapping& GetMapping() const { return m_Mapping; }
uint64_t GetTotalConnections() const { return m_TotalConnections.load(std::memory_order_relaxed); }
uint64_t GetActiveConnections() const { return m_ActiveConnections.load(std::memory_order_relaxed); }
uint64_t GetPeakActiveConnections() const { return m_PeakActiveConnections.load(std::memory_order_relaxed); }
uint64_t GetTotalBytesFromClient() const { return m_TotalBytesFromClient.load(std::memory_order_relaxed); }
uint64_t GetTotalBytesToClient() const { return m_TotalBytesToClient.load(std::memory_order_relaxed); }
metrics::Meter& GetRequestMeter() { return m_RequestMeter; }
metrics::Meter& GetBytesMeter() { return m_BytesMeter; }
// Returns a snapshot of active sessions under a shared lock.
std::vector<std::shared_ptr<TcpProxySession>> GetActiveSessions() const;
void SetRecording(bool Enabled, const std::string& Dir);
bool IsRecording() const { return m_RecordingEnabled.load(std::memory_order_relaxed); }
std::string GetRecordDir() const;
LoggerRef Log() { return m_Log; }
private:
friend class TcpProxySession;
void DoAccept();
#if defined(ASIO_HAS_LOCAL_SOCKETS)
void DoAcceptUnix();
#endif
void OnAcceptedSession(std::shared_ptr<TcpProxySession> Session);
LoggerRef m_Log;
ProxyMapping m_Mapping;
asio::io_context& m_IoContext;
asio::ip::tcp::acceptor m_TcpAcceptor;
asio::ip::tcp::endpoint m_ListenEndpoint;
#if defined(ASIO_HAS_LOCAL_SOCKETS)
asio::local::stream_protocol::acceptor m_UnixAcceptor;
#endif
bool m_Stopped = false;
void AddSession(std::shared_ptr<TcpProxySession> Session);
void RemoveSession(TcpProxySession* Session);
std::atomic<uint64_t> m_TotalConnections{0};
std::atomic<uint64_t> m_ActiveConnections{0};
std::atomic<uint64_t> m_PeakActiveConnections{0};
std::atomic<uint64_t> m_TotalBytesFromClient{0};
std::atomic<uint64_t> m_TotalBytesToClient{0};
metrics::Meter m_RequestMeter;
metrics::Meter m_BytesMeter;
mutable RwLock m_SessionsLock;
std::vector<std::shared_ptr<TcpProxySession>> m_Sessions;
std::atomic<bool> m_RecordingEnabled{false};
mutable RwLock m_RecordDirLock;
std::string m_RecordDir;
std::atomic<uint64_t> m_RecordSessionCounter{0};
};
} // namespace zen
|