Web2socket proxy.

net/web2socket_proxy/web2socket_serv.cc

+// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "web2socket_serv.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include <vector>
+
+#include <arpa/inet.h>
+#include <fcntl.h>
+#include <netinet/in.h>
+#include <signal.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+
+#include "base/logging.h"
+#include "net/web2socket_proxy/web2socket_conn.h"
+#include "third_party/libevent/evdns.h"
+
+Web2SocketServ::Web2SocketServ(const std::string& origin,
+                               struct sockaddr* addr,
+                               int addr_len)
+    : origin_(origin),
+      addr_(addr),
+      addr_len_(addr_len),
+      evbase_(NULL),
+      listening_sock_(-1) {}
+
+Web2SocketServ::~Web2SocketServ() {
+  while (!conn_pool_.empty())
+    ZapConn(conn_pool_.back());
+  if (listening_sock_ >= 0) {
+    shutdown(listening_sock_, SHUT_RDWR);
+    close(listening_sock_);
+  }
+  if (evbase_)
+    event_base_free(evbase_);
+}
+
+void Web2SocketServ::Run() {
+  if (evbase_) {
+    // One run at a time is all we do.
+    return;
+  }
+
+  evbase_ = event_init();
+  if (!evbase_) {
+    perror("Couldn't create libevent base");
+    return;
+  }
+
+  listening_sock_ = socket(AF_INET, SOCK_STREAM, 0);
+  if (listening_sock_ < 0) {
+    perror("Failed to create socket");
+    return;
+  }
+  if (bind(listening_sock_, addr_, addr_len_)) {
+    perror("Failed to bind server socket");
+    return;
+  }
+  if (listen(listening_sock_, 12)) {
+    perror("Failed to listen server socket");
+    return;
+  }
+  {
+    int on = 1;
+    setsockopt(listening_sock_, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
+  }
+  if (!SetNonBlock(listening_sock_)) {
+    perror("Failed to go non block");
+    return;
+  }
+  {
+    struct event listen_event;
+    event_set(&listen_event, listening_sock_, EV_READ | EV_PERSIST,
+              &OnConnect, this);
+    event_base_set(evbase_, &listen_event);
+    if (event_add(&listen_event, NULL)) {
+      perror("Failed to add listening event");
+      return;
+    }
+  }
+  if (evdns_init()) {
+    perror("Failed to initialize evDNS");
+    return;
+  }
+  if (!IgnoreSigPipe())
+    return;
+  fprintf(stderr, "Starting event dispatch loop, accepting connections..\n");
+  event_base_dispatch(evbase_);
+  perror("Event dispatch loop terminated");
+}
+
+void Web2SocketServ::ZapConn(Conn* cs) {
+  RevMap::iterator rit = rev_map_.find(cs);
+  if (rit != rev_map_.end()) {
+    conn_pool_.erase(rit->second);
+    rev_map_.erase(rit);
+    delete cs;
+  }
+}
+
+void Web2SocketServ::MarkConnImportance(Conn* cs,
+                                        bool important) {
+  if (conn_pool_.size() < kConnPoolLimit / 4) {
+    // Fast common path.
+    return;
+  }
+  RevMap::iterator rit = rev_map_.find(cs);
+  if (rit != rev_map_.end()) {
+    ConnPool::iterator it = rit->second;
+    CHECK(*it == cs);
+    if (important && it == conn_pool_.begin()) {
+      // Already at the top.  Shortcut.
+      return;
+    }
+    conn_pool_.erase(it);
+  }
+  if (important) {
+    conn_pool_.push_front(cs);
+    rev_map_[cs] = conn_pool_.begin();
+  } else {
+    conn_pool_.push_back(cs);
+    rev_map_[cs] = conn_pool_.end();
+    --rev_map_[cs];
+  }
+}
+
+Conn* Web2SocketServ::GetFreshConn() {
+  if (conn_pool_.size() > kConnPoolLimit) {
+    // Connections overflow.  Shut those oldest not active.
+    ConnPool::iterator it = conn_pool_.end();
+    --it;
+    for (int i = conn_pool_.size() - kConnPoolLimit; i-- > 0;) {
+      // Shut may invalidate an iterator; hence postdecrement.
+      (*it--)->Shut();
+    }
+    if (conn_pool_.size() > kConnPoolLimit + 12) {
+      // Connections overflow.  Zap the oldest not active.
+      ZapConn(conn_pool_.back());
+    }
+  }
+  Conn* cs = new Conn(this);
+  conn_pool_.push_front(cs);
+  rev_map_[cs] = conn_pool_.begin();
+  return cs;
+}
+
+bool Web2SocketServ::IsConnSane(Conn* cs) {
+  return rev_map_.find(cs) != rev_map_.end();
+}
+
+void Web2SocketServ::OnConnect(int listening_sock,
+                               short event, void* ctx) {
+  Web2SocketServ* self = static_cast<Web2SocketServ*>(ctx);
+  Conn* cs = self->GetFreshConn();
+  cs->primchan().sock = accept(listening_sock, NULL, NULL);
+  if (cs->primchan().sock < 0
+      || !SetNonBlock(cs->primchan().sock)) {
+    self->ZapConn(cs);
+    // Read readiness was triggered on listening socket
+    // yet we failed to accept a connection; definitely weird.
+    sleep(1);
+    return;
+  }
+
+  cs->primchan().bev = bufferevent_new(cs->primchan().sock,
+                                       &Conn::OnPrimchanRead,
+                                       &Conn::OnPrimchanWrite,
+                                       &Conn::OnPrimchanError,
+                                       cs->token());
+  if (cs->primchan().bev == NULL) {
+    self->ZapConn(cs);
+    return;
+  }
+  bufferevent_base_set(self->evbase_, cs->primchan().bev);
+  bufferevent_setwatermark(cs->primchan().bev, EV_READ, 0, kReadBufferLimit);
+  if (bufferevent_enable(cs->primchan().bev, EV_READ | EV_WRITE)) {
+    self->ZapConn(cs);
+    return;
+  }
+}
+
+bool Web2SocketServ::SetNonBlock(int fd) {
+  int flags = fcntl(fd, F_GETFL, 0);
+  return flags >= 0 && fcntl(fd, F_SETFL, flags | O_NONBLOCK) == 0;
+}
+
+bool Web2SocketServ::IgnoreSigPipe() {
+  struct sigaction sa;
+  sa.sa_handler = SIG_IGN;
+  sa.sa_flags = 0;
+  if (sigemptyset(&sa.sa_mask) ||
+      sigaction(SIGPIPE, &sa, 0)) {
+    perror("Failed to disable sigpipe");
+    return false;
+  }
+  return true;
+}
+