mirror of https://github.com/aria2/aria2
524 lines
15 KiB
C++
524 lines
15 KiB
C++
/*
|
|
* Wslay - The WebSocket Library
|
|
*
|
|
* Copyright (c) 2011, 2012 Tatsuhiro Tsujikawa
|
|
*
|
|
* Permission is hereby granted, free of charge, to any person obtaining
|
|
* a copy of this software and associated documentation files (the
|
|
* "Software"), to deal in the Software without restriction, including
|
|
* without limitation the rights to use, copy, modify, merge, publish,
|
|
* distribute, sublicense, and/or sell copies of the Software, and to
|
|
* permit persons to whom the Software is furnished to do so, subject to
|
|
* the following conditions:
|
|
*
|
|
* The above copyright notice and this permission notice shall be
|
|
* included in all copies or substantial portions of the Software.
|
|
*
|
|
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
|
* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
|
|
* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
|
|
* NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
|
|
* LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
|
|
* OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
|
|
* WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
*/
|
|
// WebSocket Echo Server
|
|
// This is suitable for Autobahn server test.
|
|
// g++ -Wall -O2 -g -o echoserv echoserv.cc -L../lib/.libs -I../lib/includes -lwslay -lnettle
|
|
// $ export LD_LIBRARY_PATH=../lib/.libs
|
|
// $ ./a.out 9000
|
|
#include <sys/types.h>
|
|
#include <sys/socket.h>
|
|
#include <netdb.h>
|
|
#include <unistd.h>
|
|
#include <fcntl.h>
|
|
#include <sys/epoll.h>
|
|
#include <netinet/in.h>
|
|
#include <netinet/tcp.h>
|
|
#include <signal.h>
|
|
|
|
#include <cassert>
|
|
#include <cstdio>
|
|
#include <cerrno>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <string>
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <set>
|
|
#include <iomanip>
|
|
#include <fstream>
|
|
|
|
#include <nettle/base64.h>
|
|
#include <nettle/sha.h>
|
|
#include <wslay/wslay.h>
|
|
|
|
int create_listen_socket(const char *service) {
|
|
struct addrinfo hints;
|
|
int sfd = -1;
|
|
int r;
|
|
memset(&hints, 0, sizeof(struct addrinfo));
|
|
hints.ai_family = AF_UNSPEC;
|
|
hints.ai_socktype = SOCK_STREAM;
|
|
hints.ai_flags = AI_PASSIVE | AI_ADDRCONFIG;
|
|
struct addrinfo *res;
|
|
r = getaddrinfo(0, service, &hints, &res);
|
|
if (r != 0) {
|
|
std::cerr << "getaddrinfo: " << gai_strerror(r) << std::endl;
|
|
return -1;
|
|
}
|
|
for (struct addrinfo *rp = res; rp; rp = rp->ai_next) {
|
|
sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
|
|
if (sfd == -1) {
|
|
continue;
|
|
}
|
|
int val = 1;
|
|
if (setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, &val,
|
|
static_cast<socklen_t>(sizeof(val))) == -1) {
|
|
continue;
|
|
}
|
|
if (bind(sfd, rp->ai_addr, rp->ai_addrlen) == 0) {
|
|
break;
|
|
}
|
|
close(sfd);
|
|
}
|
|
freeaddrinfo(res);
|
|
if (listen(sfd, 16) == -1) {
|
|
perror("listen");
|
|
close(sfd);
|
|
return -1;
|
|
}
|
|
return sfd;
|
|
}
|
|
|
|
int make_non_block(int fd) {
|
|
int flags, r;
|
|
while ((flags = fcntl(fd, F_GETFL, 0)) == -1 && errno == EINTR)
|
|
;
|
|
if (flags == -1) {
|
|
return -1;
|
|
}
|
|
while ((r = fcntl(fd, F_SETFL, flags | O_NONBLOCK)) == -1 && errno == EINTR)
|
|
;
|
|
if (r == -1) {
|
|
return -1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
std::string sha1(const std::string &src) {
|
|
sha1_ctx ctx;
|
|
sha1_init(&ctx);
|
|
sha1_update(&ctx, src.size(), reinterpret_cast<const uint8_t *>(src.c_str()));
|
|
uint8_t temp[SHA1_DIGEST_SIZE];
|
|
sha1_digest(&ctx, SHA1_DIGEST_SIZE, temp);
|
|
std::string res(&temp[0], &temp[SHA1_DIGEST_SIZE]);
|
|
return res;
|
|
}
|
|
|
|
std::string base64(const std::string &src) {
|
|
base64_encode_ctx ctx;
|
|
base64_encode_init(&ctx);
|
|
int dstlen = BASE64_ENCODE_RAW_LENGTH(src.size());
|
|
char *dst = new char[dstlen];
|
|
base64_encode_raw(dst, src.size(),
|
|
reinterpret_cast<const uint8_t *>(src.c_str()));
|
|
std::string res(&dst[0], &dst[dstlen]);
|
|
delete[] dst;
|
|
return res;
|
|
}
|
|
|
|
std::string create_acceptkey(const std::string &clientkey) {
|
|
std::string s = clientkey + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
|
|
return base64(sha1(s));
|
|
}
|
|
|
|
class EventHandler {
|
|
public:
|
|
virtual ~EventHandler() {}
|
|
virtual int on_read_event() = 0;
|
|
virtual int on_write_event() = 0;
|
|
virtual bool want_read() = 0;
|
|
virtual bool want_write() = 0;
|
|
virtual int fd() const = 0;
|
|
virtual bool finish() = 0;
|
|
virtual EventHandler *next() = 0;
|
|
};
|
|
|
|
ssize_t send_callback(wslay_event_context_ptr ctx, const uint8_t *data,
|
|
size_t len, int flags, void *user_data);
|
|
ssize_t recv_callback(wslay_event_context_ptr ctx, uint8_t *data, size_t len,
|
|
int flags, void *user_data);
|
|
void on_msg_recv_callback(wslay_event_context_ptr ctx,
|
|
const struct wslay_event_on_msg_recv_arg *arg,
|
|
void *user_data);
|
|
|
|
class EchoWebSocketHandler : public EventHandler {
|
|
public:
|
|
EchoWebSocketHandler(int fd) : fd_(fd) {
|
|
struct wslay_event_callbacks callbacks = {
|
|
recv_callback,
|
|
send_callback,
|
|
NULL, /* genmask_callback */
|
|
NULL, /* on_frame_recv_start_callback */
|
|
NULL, /* on_frame_recv_callback */
|
|
NULL, /* on_frame_recv_end_callback */
|
|
on_msg_recv_callback};
|
|
wslay_event_context_server_init(&ctx_, &callbacks, this);
|
|
}
|
|
virtual ~EchoWebSocketHandler() {
|
|
wslay_event_context_free(ctx_);
|
|
shutdown(fd_, SHUT_WR);
|
|
close(fd_);
|
|
}
|
|
virtual int on_read_event() {
|
|
if (wslay_event_recv(ctx_) == 0) {
|
|
return 0;
|
|
} else {
|
|
return -1;
|
|
}
|
|
}
|
|
virtual int on_write_event() {
|
|
if (wslay_event_send(ctx_) == 0) {
|
|
return 0;
|
|
} else {
|
|
return -1;
|
|
}
|
|
}
|
|
ssize_t send_data(const uint8_t *data, size_t len, int flags) {
|
|
ssize_t r;
|
|
int sflags = 0;
|
|
#ifdef MSG_MORE
|
|
if (flags & WSLAY_MSG_MORE) {
|
|
sflags |= MSG_MORE;
|
|
}
|
|
#endif // MSG_MORE
|
|
while ((r = send(fd_, data, len, sflags)) == -1 && errno == EINTR)
|
|
;
|
|
return r;
|
|
}
|
|
ssize_t recv_data(uint8_t *data, size_t len, int flags) {
|
|
ssize_t r;
|
|
while ((r = recv(fd_, data, len, 0)) == -1 && errno == EINTR)
|
|
;
|
|
return r;
|
|
}
|
|
virtual bool want_read() { return wslay_event_want_read(ctx_); }
|
|
virtual bool want_write() { return wslay_event_want_write(ctx_); }
|
|
virtual int fd() const { return fd_; }
|
|
virtual bool finish() { return !want_read() && !want_write(); }
|
|
virtual EventHandler *next() { return 0; }
|
|
|
|
private:
|
|
int fd_;
|
|
wslay_event_context_ptr ctx_;
|
|
};
|
|
|
|
ssize_t send_callback(wslay_event_context_ptr ctx, const uint8_t *data,
|
|
size_t len, int flags, void *user_data) {
|
|
EchoWebSocketHandler *sv = (EchoWebSocketHandler *)user_data;
|
|
ssize_t r = sv->send_data(data, len, flags);
|
|
if (r == -1) {
|
|
if (errno == EAGAIN || errno == EWOULDBLOCK) {
|
|
wslay_event_set_error(ctx, WSLAY_ERR_WOULDBLOCK);
|
|
} else {
|
|
wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE);
|
|
}
|
|
}
|
|
return r;
|
|
}
|
|
|
|
ssize_t recv_callback(wslay_event_context_ptr ctx, uint8_t *data, size_t len,
|
|
int flags, void *user_data) {
|
|
EchoWebSocketHandler *sv = (EchoWebSocketHandler *)user_data;
|
|
ssize_t r = sv->recv_data(data, len, flags);
|
|
if (r == -1) {
|
|
if (errno == EAGAIN || errno == EWOULDBLOCK) {
|
|
wslay_event_set_error(ctx, WSLAY_ERR_WOULDBLOCK);
|
|
} else {
|
|
wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE);
|
|
}
|
|
} else if (r == 0) {
|
|
wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE);
|
|
r = -1;
|
|
}
|
|
return r;
|
|
}
|
|
|
|
void on_msg_recv_callback(wslay_event_context_ptr ctx,
|
|
const struct wslay_event_on_msg_recv_arg *arg,
|
|
void *user_data) {
|
|
if (!wslay_is_ctrl_frame(arg->opcode)) {
|
|
struct wslay_event_msg msgarg = {arg->opcode, arg->msg, arg->msg_length};
|
|
wslay_event_queue_msg(ctx, &msgarg);
|
|
}
|
|
}
|
|
|
|
class HttpHandshakeSendHandler : public EventHandler {
|
|
public:
|
|
HttpHandshakeSendHandler(int fd, const std::string &accept_key)
|
|
: fd_(fd),
|
|
resheaders_("HTTP/1.1 101 Switching Protocols\r\n"
|
|
"Upgrade: websocket\r\n"
|
|
"Connection: Upgrade\r\n"
|
|
"Sec-WebSocket-Accept: " +
|
|
accept_key +
|
|
"\r\n"
|
|
"\r\n"),
|
|
off_(0) {}
|
|
virtual ~HttpHandshakeSendHandler() {
|
|
if (fd_ != -1) {
|
|
shutdown(fd_, SHUT_WR);
|
|
close(fd_);
|
|
}
|
|
}
|
|
virtual int on_read_event() { return 0; }
|
|
virtual int on_write_event() {
|
|
while (1) {
|
|
size_t len = resheaders_.size() - off_;
|
|
if (len == 0) {
|
|
break;
|
|
}
|
|
ssize_t r;
|
|
while ((r = write(fd_, resheaders_.c_str() + off_, len)) == -1 &&
|
|
errno == EINTR)
|
|
;
|
|
if (r == -1) {
|
|
if (errno == EAGAIN || errno == EWOULDBLOCK) {
|
|
break;
|
|
} else {
|
|
perror("write");
|
|
return -1;
|
|
}
|
|
} else {
|
|
off_ += r;
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
virtual bool want_read() { return false; }
|
|
virtual bool want_write() { return true; }
|
|
virtual int fd() const { return fd_; }
|
|
virtual bool finish() { return off_ == resheaders_.size(); }
|
|
virtual EventHandler *next() {
|
|
if (finish()) {
|
|
int fd = fd_;
|
|
fd_ = -1;
|
|
return new EchoWebSocketHandler(fd);
|
|
} else {
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
private:
|
|
int fd_;
|
|
std::string headers_;
|
|
std::string resheaders_;
|
|
size_t off_;
|
|
};
|
|
|
|
class HttpHandshakeRecvHandler : public EventHandler {
|
|
public:
|
|
HttpHandshakeRecvHandler(int fd) : fd_(fd) {}
|
|
virtual ~HttpHandshakeRecvHandler() {
|
|
if (fd_ != -1) {
|
|
close(fd_);
|
|
}
|
|
}
|
|
virtual int on_read_event() {
|
|
char buf[4096];
|
|
ssize_t r;
|
|
std::string client_key;
|
|
while (1) {
|
|
while ((r = read(fd_, buf, sizeof(buf))) == -1 && errno == EINTR)
|
|
;
|
|
if (r == -1) {
|
|
if (errno == EWOULDBLOCK || errno == EAGAIN) {
|
|
break;
|
|
} else {
|
|
perror("read");
|
|
return -1;
|
|
}
|
|
} else if (r == 0) {
|
|
std::cerr << "http_upgrade: Got EOF" << std::endl;
|
|
return -1;
|
|
} else {
|
|
headers_.append(buf, buf + r);
|
|
if (headers_.size() > 8192) {
|
|
std::cerr << "Too large http header" << std::endl;
|
|
return -1;
|
|
}
|
|
}
|
|
}
|
|
if (headers_.find("\r\n\r\n") != std::string::npos) {
|
|
std::string::size_type keyhdstart;
|
|
if (headers_.find("Upgrade: websocket\r\n") == std::string::npos ||
|
|
headers_.find("Connection: Upgrade\r\n") == std::string::npos ||
|
|
(keyhdstart = headers_.find("Sec-WebSocket-Key: ")) ==
|
|
std::string::npos) {
|
|
std::cerr << "http_upgrade: missing required headers" << std::endl;
|
|
return -1;
|
|
}
|
|
keyhdstart += 19;
|
|
std::string::size_type keyhdend = headers_.find("\r\n", keyhdstart);
|
|
client_key = headers_.substr(keyhdstart, keyhdend - keyhdstart);
|
|
accept_key_ = create_acceptkey(client_key);
|
|
}
|
|
return 0;
|
|
}
|
|
virtual int on_write_event() { return 0; }
|
|
virtual bool want_read() { return true; }
|
|
virtual bool want_write() { return false; }
|
|
virtual int fd() const { return fd_; }
|
|
virtual bool finish() { return !accept_key_.empty(); }
|
|
virtual EventHandler *next() {
|
|
if (finish()) {
|
|
int fd = fd_;
|
|
fd_ = -1;
|
|
return new HttpHandshakeSendHandler(fd, accept_key_);
|
|
} else {
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
private:
|
|
int fd_;
|
|
std::string headers_;
|
|
std::string accept_key_;
|
|
};
|
|
|
|
class ListenEventHandler : public EventHandler {
|
|
public:
|
|
ListenEventHandler(int fd) : fd_(fd), cfd_(-1) {}
|
|
virtual ~ListenEventHandler() {
|
|
close(fd_);
|
|
close(cfd_);
|
|
}
|
|
virtual int on_read_event() {
|
|
if (cfd_ != -1) {
|
|
close(cfd_);
|
|
}
|
|
while ((cfd_ = accept(fd_, 0, 0)) == -1 && errno == EINTR)
|
|
;
|
|
if (cfd_ == -1) {
|
|
perror("accept");
|
|
}
|
|
return 0;
|
|
}
|
|
virtual int on_write_event() { return 0; }
|
|
virtual bool want_read() { return true; }
|
|
virtual bool want_write() { return false; }
|
|
virtual int fd() const { return fd_; }
|
|
virtual bool finish() { return false; }
|
|
virtual EventHandler *next() {
|
|
if (cfd_ != -1) {
|
|
int val = 1;
|
|
int fd = cfd_;
|
|
cfd_ = -1;
|
|
if (make_non_block(fd) == -1 ||
|
|
setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &val,
|
|
(socklen_t)sizeof(val)) == -1) {
|
|
close(fd);
|
|
return 0;
|
|
}
|
|
return new HttpHandshakeRecvHandler(fd);
|
|
} else {
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
private:
|
|
int fd_;
|
|
int cfd_;
|
|
};
|
|
|
|
int ctl_epollev(int epollfd, int op, EventHandler *handler) {
|
|
epoll_event ev;
|
|
memset(&ev, 0, sizeof(ev));
|
|
int events = 0;
|
|
if (handler->want_read()) {
|
|
events |= EPOLLIN;
|
|
}
|
|
if (handler->want_write()) {
|
|
events |= EPOLLOUT;
|
|
}
|
|
ev.events = events;
|
|
ev.data.ptr = handler;
|
|
return epoll_ctl(epollfd, op, handler->fd(), &ev);
|
|
}
|
|
|
|
void reactor(int sfd) {
|
|
std::set<EventHandler *> handlers;
|
|
ListenEventHandler *listen_handler = new ListenEventHandler(sfd);
|
|
handlers.insert(listen_handler);
|
|
int epollfd = epoll_create(16);
|
|
if (epollfd == -1) {
|
|
perror("epoll_create");
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
if (ctl_epollev(epollfd, EPOLL_CTL_ADD, listen_handler) == -1) {
|
|
perror("epoll_ctl");
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
static const size_t MAX_EVENTS = 64;
|
|
epoll_event events[MAX_EVENTS];
|
|
while (1) {
|
|
int nfds = epoll_wait(epollfd, events, MAX_EVENTS, -1);
|
|
if (nfds == -1) {
|
|
perror("epoll_wait");
|
|
return;
|
|
}
|
|
for (int n = 0; n < nfds; ++n) {
|
|
EventHandler *eh = (EventHandler *)events[n].data.ptr;
|
|
if (((events[n].events & EPOLLIN) && eh->on_read_event() == -1) ||
|
|
((events[n].events & EPOLLOUT) && eh->on_write_event() == -1) ||
|
|
(events[n].events & (EPOLLERR | EPOLLHUP))) {
|
|
handlers.erase(eh);
|
|
delete eh;
|
|
} else {
|
|
EventHandler *next = eh->next();
|
|
if (next) {
|
|
handlers.insert(next);
|
|
if (ctl_epollev(epollfd, EPOLL_CTL_ADD, next) == -1) {
|
|
if (errno == EEXIST) {
|
|
if (ctl_epollev(epollfd, EPOLL_CTL_MOD, next) == -1) {
|
|
perror("epoll_ctl");
|
|
delete next;
|
|
}
|
|
} else {
|
|
perror("epoll_ctl");
|
|
delete next;
|
|
}
|
|
}
|
|
}
|
|
if (eh->finish()) {
|
|
handlers.erase(eh);
|
|
delete eh;
|
|
} else {
|
|
if (ctl_epollev(epollfd, EPOLL_CTL_MOD, eh) == -1) {
|
|
perror("epoll_ctl");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
if (argc < 2) {
|
|
std::cerr << "Usage: " << argv[0] << " PORT" << std::endl;
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
struct sigaction act;
|
|
memset(&act, 0, sizeof(struct sigaction));
|
|
act.sa_handler = SIG_IGN;
|
|
sigaction(SIGPIPE, &act, 0);
|
|
int sfd = create_listen_socket(argv[1]);
|
|
if (sfd == -1) {
|
|
std::cerr << "Failed to create server socket" << std::endl;
|
|
exit(EXIT_FAILURE);
|
|
}
|
|
std::cout << "WebSocket echo server, listening on " << argv[1] << std::endl;
|
|
reactor(sfd);
|
|
}
|