socks5.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>

#include <signal.h>
#include <sys/epoll.h>

#include <errno.h>

#undef EPOLLRDHUP

enum {
    DEBUG,
    INFORM,
    WARNING,
    ERROR,
    FATAL
};
int g_loglevel = INFORM;

#define PRINTF(level, format, ...) \
    do { \
        if (level < g_loglevel) break; \
        fprintf(stderr, "[%s:%d] " format, \
            __func__, __LINE__, ##__VA_ARGS__); \
    } while (0)

#define PDEBUG(format, ...) PRINTF(DEBUG, format, ##__VA_ARGS__)
#define PINFORM(format, ...) PRINTF(INFORM, format, ##__VA_ARGS__)
#define PWARNING(format, ...) PRINTF(WARNING, format, ##__VA_ARGS__)
#define PERROR(format, ...) PRINTF(ERROR, format, ##__VA_ARGS__)
#define PFATAL(format, ...) PRINTF(FATAL, format, ##__VA_ARGS__)

#define SAFE_CLOSE(fd) \
    do { \
        if (0 < fd) { \
            close(fd); \
        } \
    } while (0)

#define MAX_EVENTS 128
#define MAX_BUFFSIZE 2048

// socks5 server port
static int g_port = 9090;

static void help();
static int check_para(int argc, char **argv);
static int set_non_blocking(int fd);
static int socks5_sockopt(int fd);
//static int socks5_send(int fd, const void *buf, size_t len);
//static int socks5_recv(int fd, const void *buf, size_t len);
static int socks5_auth(int fd);

int main(int argc, char **argv)
{
    int efd = -1;
    struct epoll_event events[MAX_EVENTS];
    struct epoll_event event;
    int nfds;

    int sfd = -1;
    int cfd = -1;
    int rfd = -1;
    struct sockaddr_in addr;
    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);

    char buff[MAX_BUFFSIZE];
    int blen;
    int ret;
    int i;

    if (-1 == check_para(argc, argv)) {
        goto _err;
    }
    
    signal(SIGPIPE, SIG_IGN);
    
    efd = epoll_create(MAX_EVENTS);
    if (-1 == efd) {
        PERROR("epoll_create failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }
    memset(&events, 0, MAX_EVENTS * sizeof(struct epoll_event));

    sfd = socket(AF_INET, SOCK_STREAM, 0);
    if (-1 == sfd) {
        PERROR("socket failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }

    if (-1 == socks5_sockopt(sfd)) {
        PERROR("socks5_sockopt failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }
    if (-1 == set_non_blocking(sfd)) {
        PERROR("set_non_blocking failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }

    memset(&addr, 0, sizeof(struct sockaddr_in));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons(g_port);
    if (-1 == bind(sfd, (struct sockaddr *)&addr, sizeof(addr))) {
        PERROR("bind failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }

    if (-1 == listen(sfd, 5)) {
        PERROR("listen failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }

    event.data.fd = sfd;
    event.events = EPOLLIN;
    if (-1 == epoll_ctl(efd, EPOLL_CTL_ADD, sfd, &event)) {
        PERROR("epoll_ctl failed: [%d] %s\n", errno, strerror(errno));
        goto _err;
    }
    
    PINFORM("socks5 server start\n");
    PINFORM("log level [%d], server port [%d]\n", g_loglevel, g_port);
    
    for (;;) {
        nfds = epoll_wait(efd, events, MAX_EVENTS, -1);
        if (-1 == nfds) {
            PERROR("epoll_wait failed: [%d] %s\n", errno, strerror(errno));
            continue;
        }
        
        for (i = 0; i < nfds; i++) {
            //PDEBUG("Event [%x] on fd [%d]\n", events[i].events, events[i].data.fd);
            
            if (sfd == events[i].data.fd) {
                cfd = accept(sfd, (struct sockaddr *)&client_addr, &client_len);
                if (-1 == cfd) {
                    PERROR("accept failed: [%d] %s\n", errno, strerror(errno));
                    continue;
                }

                if ((rfd = socks5_auth(cfd)) <= 0) {
                    PDEBUG("socks5_auth failed!\n");
                    close(cfd);
                    continue;
                }

                socks5_sockopt(cfd);
                set_non_blocking(cfd);
                set_non_blocking(rfd);

                event.data.u64 = cfd | ((uint64_t)rfd << 32);
                event.events = EPOLLIN;
#ifdef EPOLLRDHUP
                event.events |= EPOLLRDHUP;
#endif                
                if (-1 == epoll_ctl(efd, EPOLL_CTL_ADD, cfd, &event)) {
                    PERROR("epoll_ctl failed: [%d] %s\n", errno, strerror(errno));
                } else {
                    event.data.u64 = rfd | ((uint64_t)cfd << 32);
                    if (-1 == epoll_ctl(efd, EPOLL_CTL_ADD, rfd, &event)) {
                        PERROR("epoll_ctl failed: [%d] %s\n", errno, strerror(errno));
                        epoll_ctl(efd, EPOLL_CTL_DEL, cfd, &event);
                        continue;
                    }
                    PDEBUG("epoll_ctl ADD fd [%d][%d]\n", cfd, rfd);
                }
            }
#ifdef EPOLLRDHUP
            else if ((EPOLLERR | EPOLLHUP | EPOLLRDHUP) & events[i].events) {
#else
            else if ((EPOLLERR | EPOLLHUP) & events[i].events) {
#endif                
                PDEBUG("fd [%d] closed\n", events[i].data.fd);
#if 1
                // close will autmatic do epoll_ctl EPOLL_CTL_DEL
                if (-1 == close(events[i].data.fd)) {
                    PERROR("close [%d] failed: [%d] %s\n", events[i].data.fd, errno, strerror(errno));
                }
                if (-1 == close(events[i].data.u64 >> 32)) {
                    PERROR("close [%d] failed: [%d] %s\n", events[i].data.u64 >> 32, errno, strerror(errno));
                }
#else
                // the argument &event is useless
                if (-1 == epoll_ctl(efd, EPOLL_CTL_DEL, events[i].data.fd, &event)) {
                    PERROR("epoll_ctl failed: [%d] %s\n", errno, strerror(errno));
                } else {
                    if (-1 == close(events[i].data.fd)) {
                        PERROR("close [%d] failed: [%d] %s\n", events[i].data.fd, errno, strerror(errno));
                    }
                }
#endif
            } else if (EPOLLIN & events[i].events) {
                ret = recv(events[i].data.fd, buff, MAX_BUFFSIZE, 0);
                // peer connection closed
                if (0 == ret) {
                    PDEBUG("fd [%d] closed\n", events[i].data.fd);
#if 1
                    // close will autmatic do epoll_ctl EPOLL_CTL_DEL
                    if (-1 == close(events[i].data.fd)) {
                        PERROR("close [%d] failed: [%d] %s\n", events[i].data.fd, errno, strerror(errno));
                    }
                    if (-1 == close(events[i].data.u64 >> 32)) {
                        PERROR("close [%d] failed: [%d] %s\n", events[i].data.u64 >> 32, errno, strerror(errno));
                    }
#else
                    // the argument &event is useless
                    if (-1 == epoll_ctl(efd, EPOLL_CTL_DEL, events[i].data.fd, &event)) {
                        PERROR("epoll_ctl failed: [%d] %s\n", errno, strerror(errno));
                    } else {
                        if (-1 == close(events[i].data.fd)) {
                            PERROR("close [%d] failed: [%d] %s\n", events[i].data.fd, errno, strerror(errno));
                        }
                    }
#endif
                } else if (-1 == ret) {
                    PERROR("recv failed: [%d] %s\n", errno, strerror(errno));
                } else {
                    if (-1 == send(events[i].data.u64 >> 32, buff, ret, 0)) {
                        PERROR("send failed: [%d] %s\n", errno, strerror(errno));
                    }
                }
            }
            else {
                PWARNING("unhandled event [%x] on fd [%d]\n", events[i].events, events[i].data.fd);
            }
        }
    }

    return 0;

_err:
    SAFE_CLOSE(efd);
    SAFE_CLOSE(sfd);

    return -1;
}

static void help()
{
    printf("Usage: socks5 [OPTION]\n\n");
    printf("Options:\n");
    printf(" -p <port>  socks5 service port, default 9090\n");
    printf(" -l <level> log level in [0-4], default 1\n");
    printf(" -d         run as a daemon\n");
    printf(" -h         print help information\n");
}

static int check_para(int argc, char **argv)
{
    int ch;
    int ret = 0;

    while ((ch = getopt(argc, argv, ":p:l:dh")) != -1) {
        switch (ch) {
            case 'd':
                daemon(1, 0);
                break;
            case 'p':
                g_port = (atoi(optarg) <= 0) ? g_port : atoi(optarg);
                break;
            case 'l':
                g_loglevel = (atoi(optarg) >= DEBUG&& atoi(optarg) <= FATAL)
                    ? atoi(optarg) : g_loglevel;
                break;
            case 'h':
            default:
                help();
                ret = -1;
                break;
        }
    }

    return ret;
}

static int set_non_blocking(int fd)
{
    int flags, s;

    flags = fcntl(fd, F_GETFL, 0);
    if (-1 == flags) {
        PERROR("fcntl failed: [%d] %s\n", errno, strerror(errno));
        return -1;
    }

    flags |= O_NONBLOCK;
    s = fcntl(fd, F_SETFL, flags);
    if (-1 == s) {
        PERROR("fcntl failed: [%d] %s\n", errno, strerror(errno));
        return -1;
    }

    return 0;
}

static int socks5_sockopt(int fd)
{
    struct timeval tmo = {0};
    int opt = 1;

    tmo.tv_sec = 2;
    tmo.tv_usec = 0;
    setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, (char *)&tmo, sizeof(tmo));
    setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, (char *)&tmo, sizeof(tmo));

#ifdef SO_NOSIGPIPE
    setsockopt(fd, SOL_SOCKET, SO_NOSIGPIPE, &opt, sizeof(opt));
#endif

    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    return 0;
}

static int socks5_auth(int fd)
{
    size_t len;
    char buff[MAX_BUFFSIZE];
    struct sockaddr_in addr;
    socklen_t addr_len = sizeof(addr);
    struct hostent *hptr;
    int rfd = -1;

    len = 2;
    if (len != recv(fd, buff, len, 0)) {
        goto _err1;
    }

    if (-1 == recv(fd, buff + 2, buff[1], 0)) {
        goto _err1;
    }

    // socks5 verison 0x05
    if (0x05 != buff[0] || 0x00 == buff[1]) {
        goto _err1;
    }

    // no auth
    send(fd, "\x05\x00", 2, 0);

    len = 4;
    if (len != recv(fd, buff, len, 0)) {
        goto _err2;
    }
    // must be connect
    if (0x05 != buff[0] || 0x01 != buff[1] || 0x00 != buff[2]) {
        goto _err2;
    }

    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;

    // address type: 0x01 => IPV4; 0x03 => DOMAIN
    if (0x01 == buff[3]) {
        len = 6;
        if (len != recv(fd, buff, len, 0)) {
            goto _err2;
        }

        memcpy(&(addr.sin_addr.s_addr), buff, 4);
        memcpy(&(addr.sin_port), buff + 4, 2);

        PDEBUG("connect: %s:%d\n",
            inet_ntoa(addr.sin_addr.s_addr), htons(addr.sin_port));
    } else if (0x03 == buff[3]) {
        len = 1;
        if (len != recv(fd, buff, len, 0)) {
            goto _err2;
        }
        len = buff[0];
        buff[len] = 0;
        if (len != recv(fd, buff, len, 0)) {
            goto _err2;
        }
        PDEBUG("connect: %s\n", buff);
        hptr = gethostbyname(buff);
        if (NULL == hptr 
            || AF_INET != hptr->h_addrtype
            || NULL == *(hptr->h_addr_list)) {
            PWARNING("gethostbyname(%s) error", buff);
            goto _err2;
        }
        memcpy(&(addr.sin_addr.s_addr), *(hptr->h_addr_list), 4);

        len = 2;
        if (len != recv(fd, buff, len, 0)) {
            goto _err2;
        }
        memcpy(&(addr.sin_port), buff, 2);
    }

    rfd = socket(AF_INET, SOCK_STREAM, 0);
    if (-1 == rfd) {
        PERROR("socket failed, [%d] %s\n", errno, strerror(errno));
        goto _err2;
    }

    socks5_sockopt(rfd);

    if (-1 == connect(rfd, (struct sockaddr *)&addr, sizeof(addr))) {
        PERROR("connect failed, [%d] %s\n", errno, strerror(errno));
        goto _err2;
    }

    addr_len = sizeof(addr);
    if (-1 == getpeername(rfd, (struct sockaddr *)&addr, (socklen_t *)&addr_len)) {
        PERROR("getpeername failed, [%d] %s\n", errno, strerror(errno));
        goto _err;
    }

    // reply remote address info
    memcpy(buff, "\x05\x00\x00\x01", 4);
    memcpy(buff + 4, &(addr.sin_addr.s_addr), 4);
    memcpy(buff + 8, &(addr.sin_port), 2);
    if (-1 == send(fd, buff, 10, 0)) {
        PERROR("send failed, [%d] %s\n", errno, strerror(errno));
        goto _err2;
    }

    return rfd;

_err1:
    send(fd, "\x05\xff", 2, 0);
    goto _err;
_err2:
    send(fd, "\x05\x01", 2, 0);
    goto _err;
_err:
    if (rfd > 0) close(rfd);

    return -1;
}