89 lines
2.0 KiB
C++
89 lines
2.0 KiB
C++
#include "examples/socks4a/tunnel.h"
|
|
|
|
#include <malloc.h>
|
|
#include <stdio.h>
|
|
#include <sys/resource.h>
|
|
#include <unistd.h>
|
|
|
|
using namespace muduo;
|
|
using namespace muduo::net;
|
|
|
|
EventLoop* g_eventLoop;
|
|
InetAddress* g_serverAddr;
|
|
std::map<string, TunnelPtr> g_tunnels;
|
|
|
|
void onServerConnection(const TcpConnectionPtr& conn)
|
|
{
|
|
LOG_DEBUG << (conn->connected() ? "UP" : "DOWN");
|
|
if (conn->connected())
|
|
{
|
|
conn->setTcpNoDelay(true);
|
|
conn->stopRead();
|
|
TunnelPtr tunnel(new Tunnel(g_eventLoop, *g_serverAddr, conn));
|
|
tunnel->setup();
|
|
tunnel->connect();
|
|
g_tunnels[conn->name()] = tunnel;
|
|
}
|
|
else
|
|
{
|
|
assert(g_tunnels.find(conn->name()) != g_tunnels.end());
|
|
g_tunnels[conn->name()]->disconnect();
|
|
g_tunnels.erase(conn->name());
|
|
}
|
|
}
|
|
|
|
void onServerMessage(const TcpConnectionPtr& conn, Buffer* buf, Timestamp)
|
|
{
|
|
LOG_DEBUG << buf->readableBytes();
|
|
if (!conn->getContext().empty())
|
|
{
|
|
const TcpConnectionPtr& clientConn
|
|
= boost::any_cast<const TcpConnectionPtr&>(conn->getContext());
|
|
clientConn->send(buf);
|
|
}
|
|
}
|
|
|
|
void memstat()
|
|
{
|
|
malloc_stats();
|
|
}
|
|
|
|
int main(int argc, char* argv[])
|
|
{
|
|
if (argc < 4)
|
|
{
|
|
fprintf(stderr, "Usage: %s <host_ip> <port> <listen_port>\n", argv[0]);
|
|
}
|
|
else
|
|
{
|
|
LOG_INFO << "pid = " << getpid() << ", tid = " << CurrentThread::tid();
|
|
{
|
|
// set max virtual memory to 256MB.
|
|
size_t kOneMB = 1024*1024;
|
|
rlimit rl = { 256*kOneMB, 256*kOneMB };
|
|
setrlimit(RLIMIT_AS, &rl);
|
|
}
|
|
const char* ip = argv[1];
|
|
uint16_t port = static_cast<uint16_t>(atoi(argv[2]));
|
|
InetAddress serverAddr(ip, port);
|
|
g_serverAddr = &serverAddr;
|
|
|
|
uint16_t acceptPort = static_cast<uint16_t>(atoi(argv[3]));
|
|
InetAddress listenAddr(acceptPort);
|
|
|
|
EventLoop loop;
|
|
g_eventLoop = &loop;
|
|
loop.runEvery(3, memstat);
|
|
|
|
TcpServer server(&loop, listenAddr, "TcpRelay");
|
|
|
|
server.setConnectionCallback(onServerConnection);
|
|
server.setMessageCallback(onServerMessage);
|
|
|
|
server.start();
|
|
|
|
loop.loop();
|
|
}
|
|
}
|
|
|