使用libevent tcp 转发代理

共计 4780 个字符，预计需要花费 12 分钟才能阅读完成。

#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#ifdef WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <sys/socket.h>
#include <netinet/in.h>
#endif

#include <event2/bufferevent_ssl.h>
#include <event2/bufferevent.h>
#include <event2/buffer.h>
#include <event2/listener.h>
#include <event2/util.h>

static struct event_base *base;
static struct sockaddr_storage lsn_on_addr = {0};
static struct sockaddr_storage con_to_addr = {0};

#define MAX_OUTPUT (512 * 1024)
#define MAX_BUFFER_SIZE 1024 * 20
void __on_drained(struct bufferevent *bev, void *ctx);
void __on_recv(struct bufferevent *bev, void *ctx);
void __on_close(struct bufferevent *bev, void *ctx);
void __on_error(struct bufferevent *bev, short what, void *ctx);

void __on_accept(struct evconnlistener *listener, evutil_socket_t fd, struct sockaddr *a, int slen, void *p) {
    char ip[128] = {0};
    struct bufferevent *b_srv, *b_clt;

    /* Create two linked bufferevent objects: one to connect, one for the
     * new connection */
    b_clt = bufferevent_socket_new(base, fd, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_DEFER_CALLBACKS);

    b_srv = bufferevent_socket_new(base, -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_DEFER_CALLBACKS);

    assert(b_clt && b_srv);

    if (bufferevent_socket_connect(b_srv, (struct sockaddr *)&con_to_addr, sizeof(con_to_addr)) < 0) {
        perror("bufferevent_socket_connect");
        bufferevent_free(b_srv);
        bufferevent_free(b_clt);
        return;
    }

    printf("Accept from: %s, connect to: %s\n", evutil_inet_ntop(AF_INET, &((struct sockaddr_in *)a)->sin_addr, ip, sizeof(ip)),
           evutil_inet_ntop(AF_INET, &((struct sockaddr_in *)&con_to_addr)->sin_addr, ip, sizeof(ip)));

    bufferevent_setcb(b_clt, __on_recv, NULL, __on_error, b_srv);
    bufferevent_setcb(b_srv, __on_recv, NULL, __on_error, b_clt);

    bufferevent_enable(b_clt, EV_READ | EV_WRITE);
    bufferevent_enable(b_srv, EV_READ | EV_WRITE);
}

void __on_drained(struct bufferevent *bev, void *ctx) {

    struct bufferevent *partner = (struct bufferevent *)ctx;

    printf("%d no full\n", bufferevent_getfd(bev));
    /* We were choking the other side until we drained our outbuf a bit.
     * Now it seems drained. */
    bufferevent_setcb(bev, __on_recv, NULL, __on_error, partner);
    bufferevent_setwatermark(bev, EV_WRITE, 0, 0);
    if (partner) bufferevent_enable(partner, EV_READ);
}

void __on_recv(struct bufferevent *bev, void *ctx) {
    struct bufferevent *partner = (struct bufferevent *)ctx;
    char buf[MAX_BUFFER_SIZE] = {0};
    size_t ret = 0;
    while ((ret = bufferevent_read(bev, buf, sizeof(buf))) > 0) {
        bufferevent_write(partner, buf, ret);
        // printf(" >>> size : %d data: %s\n", ret, buf);
    }

    // struct bufferevent *partner = (struct bufferevent *)ctx;
    // struct evbuffer *src, *dst;
    // size_t len;
    // src = bufferevent_get_input(bev);
    // len = evbuffer_get_length(src);
    // if (!partner) {
    //     evbuffer_drain(src, len);
    //     return;
    // }
    // dst = bufferevent_get_output(partner);
    // evbuffer_add_buffer(dst, src);

    // if (evbuffer_get_length(dst) >= MAX_OUTPUT) {
    //     /* We're giving the other side data faster than it can
    //      * pass it on.  Stop reading here until we have drained the
    //      * other side to MAX_OUTPUT/2 bytes. */
    //     printf("%d is full\n", bufferevent_getfd(bev));
    //     bufferevent_setcb(partner, __on_recv, __on_drained, __on_error, bev);
    //     bufferevent_setwatermark(partner, EV_WRITE, MAX_OUTPUT / 2, MAX_OUTPUT);
    //     bufferevent_disable(bev, EV_READ);
    // }
}
void __on_close(struct bufferevent *bev, void *ctx) {
    struct evbuffer *b = bufferevent_get_output(bev);

    if (evbuffer_get_length(b) == 0) {
        printf("Close %2d done\n", bufferevent_getfd(bev));
        bufferevent_free(bev);
    }
}

void __on_error(struct bufferevent *bev, short what, void *ctx) {
    struct bufferevent *partner = (struct bufferevent *)ctx;

    if (what & (BEV_EVENT_EOF | BEV_EVENT_ERROR)) {
        if (what & BEV_EVENT_ERROR) {
            if (errno) perror("connection error");
        }

        if (partner) {
            /* Flush all pending data */
            __on_recv(bev, ctx);

            if (evbuffer_get_length(bufferevent_get_output(partner))) {
                /* We still have to flush data from the other
                 * side, but when that's done, close the other
                 * side. */
                bufferevent_setcb(partner, NULL, __on_close, __on_error, NULL);
                bufferevent_disable(partner, EV_READ);
            }
            else {
                /* We have nothing left to say to the other
                 * side; close it. */
                printf("Close %2d & %2d\n", bufferevent_getfd(partner), bufferevent_getfd(bev));
                bufferevent_free(partner);
            }
        }
        else {
            printf("Close %2d\n", bufferevent_getfd(bev));
        }

        bufferevent_free(bev);
    }
}
void syntax() {
    fputs("Syntax:\n", stderr);
    fputs("   ./tcp_proxy [-s] [-W] <listen-on-addr> <connect-to-addr>\n", stderr);
    fputs("Example:\n", stderr);
    fputs("   ./tcp-proxy 127.0.0.1:15001 127.0.0.1:5001\n", stderr);
    exit(1);
}

int main(int argc, char *argv[]) {
    int ret = -1;
    int con_len = sizeof(con_to_addr);
    int lsn_len = sizeof(lsn_on_addr);

    struct evconnlistener *listener = NULL;

    if (argc < 3) {
        syntax();
    }

    ret = evutil_parse_sockaddr_port(argv[2], (struct sockaddr *)&con_to_addr, &con_len);
    assert(0 == ret);

    ret = evutil_parse_sockaddr_port(argv[1], (struct sockaddr *)&lsn_on_addr, &lsn_len);
    assert(0 == ret);

    base = event_base_new();
    assert(base);

    listener = evconnlistener_new_bind(base, __on_accept, NULL, LEV_OPT_CLOSE_ON_FREE | LEV_OPT_CLOSE_ON_EXEC | LEV_OPT_REUSEABLE, -1,
                                       (struct sockaddr *)&lsn_on_addr, lsn_len);
    assert(listener);

    event_base_dispatch(base);

    evconnlistener_free(listener);
    event_base_free(base);
    exit(EXIT_SUCCESS);
}