tools/android/forwarder2/forwarder.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "tools/android/forwarder2/forwarder.h"
   6
   7 #include "base/basictypes.h"
   8 #include "base/logging.h"
   9 #include "base/posix/eintr_wrapper.h"
  10 #include "tools/android/forwarder2/socket.h"
  11
  12 namespace forwarder2 {
  13 namespace {
  14
  15 const int kBufferSize = 32 * 1024;
  16
  17 }  // namespace
  18
  19
  20 // Helper class to buffer reads and writes from one socket to another.
  21 // Each implements a small buffer connected two one input socket, and
  22 // one output socket.
  23 //
  24 //   socket_from_ ---> [BufferedCopier] ---> socket_to_
  25 //
  26 // These objects are used in a pair to handle duplex traffic, as in:
  27 //
  28 //                   -------> [BufferedCopier_1] --->
  29 //                  |                                |
  30 //      socket_1   *                                  * socket_2
  31 //                  |                                |
  32 //                   <------ [BufferedCopier_2] <----
  33 //
  34 // When a BufferedCopier is in the READING state (see below), it only listens
  35 // to events on its input socket, and won't detect when its output socket
  36 // disconnects. To work around this, its peer will call its Close() method
  37 // when that happens.
  38
  39 class Forwarder::BufferedCopier {
  40  public:
  41   // Possible states:
  42   //    READING - Empty buffer and Waiting for input.
  43   //    WRITING - Data in buffer, and waiting for output.
  44   //    CLOSING - Like WRITING, but do not try to read after that.
  45   //    CLOSED  - Completely closed.
  46   //
  47   // State transitions are:
  48   //
  49   //   T01:  READING ---[receive data]---> WRITING
  50   //   T02:  READING ---[error on input socket]---> CLOSED
  51   //   T03:  READING ---[Close() call]---> CLOSED
  52   //
  53   //   T04:  WRITING ---[write partial data]---> WRITING
  54   //   T05:  WRITING ---[write all data]----> READING
  55   //   T06:  WRITING ---[error on output socket]----> CLOSED
  56   //   T07:  WRITING ---[Close() call]---> CLOSING
  57   //
  58   //   T08:  CLOSING ---[write partial data]---> CLOSING
  59   //   T09:  CLOSING ---[write all data]----> CLOSED
  60   //   T10:  CLOSING ---[Close() call]---> CLOSING
  61   //   T11:  CLOSING ---[error on output socket] ---> CLOSED
  62   //
  63   enum State {
  64     STATE_READING = 0,
  65     STATE_WRITING = 1,
  66     STATE_CLOSING = 2,
  67     STATE_CLOSED = 3,
  68   };
  69
  70   // Does NOT own the pointers.
  71   BufferedCopier(Socket* socket_from, Socket* socket_to)
  72       : socket_from_(socket_from),
  73         socket_to_(socket_to),
  74         bytes_read_(0),
  75         write_offset_(0),
  76         peer_(NULL),
  77         state_(STATE_READING) {}
  78
  79   // Sets the 'peer_' field pointing to the other BufferedCopier in a pair.
  80   void SetPeer(BufferedCopier* peer) {
  81     DCHECK(!peer_);
  82     peer_ = peer;
  83   }
  84
  85   bool is_closed() const { return state_ == STATE_CLOSED; }
  86
  87   // Gently asks to close a buffer. Called either by the peer or the forwarder.
  88   void Close() {
  89     switch (state_) {
  90       case STATE_READING:
  91         state_ = STATE_CLOSED;  // T03
  92         break;
  93       case STATE_WRITING:
  94         state_ = STATE_CLOSING;  // T07
  95         break;
  96       case STATE_CLOSING:
  97         break;  // T10
  98       case STATE_CLOSED:
  99         ;
 100     }
 101   }
 102
 103   // Call this before select(). This updates |read_fds|,
 104   // |write_fds| and |max_fd| appropriately *if* the buffer isn't closed.
 105   void PrepareSelect(fd_set* read_fds, fd_set* write_fds, int* max_fd) {
 106     int fd;
 107     switch (state_) {
 108       case STATE_READING:
 109         DCHECK(bytes_read_ == 0);
 110         DCHECK(write_offset_ == 0);
 111         fd = socket_from_->fd();
 112         if (fd < 0) {
 113           ForceClose();  // T02
 114           return;
 115         }
 116         FD_SET(fd, read_fds);
 117         break;
 118
 119       case STATE_WRITING:
 120       case STATE_CLOSING:
 121         DCHECK(bytes_read_ > 0);
 122         DCHECK(write_offset_ < bytes_read_);
 123         fd = socket_to_->fd();
 124         if (fd < 0) {
 125           ForceClose();  // T06
 126           return;
 127         }
 128         FD_SET(fd, write_fds);
 129         break;
 130
 131       case STATE_CLOSED:
 132         return;
 133     }
 134     *max_fd = std::max(*max_fd, fd);
 135   }
 136
 137   // Call this after a select() call to operate over the buffer.
 138   void ProcessSelect(const fd_set& read_fds, const fd_set& write_fds) {
 139     int fd;
 140     int ret;
 141     // With FORTIFY_SOURCE, FD_ISSET is implemented as a function that takes a
 142     // non-const fd_set*. Make a copy of the passed arguments so we can safely
 143     // take a reference.
 144     fd_set read_fds_copy = read_fds;
 145     fd_set write_fds_copy = write_fds;
 146     switch (state_) {
 147       case STATE_READING:
 148         fd = socket_from_->fd();
 149         if (fd < 0) {
 150           state_ = STATE_CLOSED;  // T02
 151           return;
 152         }
 153         if (!FD_ISSET(fd, &read_fds_copy))
 154           return;
 155
 156         ret = socket_from_->NonBlockingRead(buffer_, kBufferSize);
 157         if (ret <= 0) {
 158           ForceClose();  // T02
 159           return;
 160         }
 161         bytes_read_ = ret;
 162         write_offset_ = 0;
 163         state_ = STATE_WRITING;  // T01
 164         break;
 165
 166       case STATE_WRITING:
 167       case STATE_CLOSING:
 168         fd = socket_to_->fd();
 169         if (fd < 0) {
 170           ForceClose();  // T06 + T11
 171           return;
 172         }
 173         if (!FD_ISSET(fd, &write_fds_copy))
 174           return;
 175
 176         ret = socket_to_->NonBlockingWrite(buffer_ + write_offset_,
 177                                            bytes_read_ - write_offset_);
 178         if (ret <= 0) {
 179           ForceClose();  // T06 + T11
 180           return;
 181         }
 182
 183         write_offset_ += ret;
 184         if (write_offset_ < bytes_read_)
 185           return;  // T08 + T04
 186
 187         write_offset_ = 0;
 188         bytes_read_ = 0;
 189         if (state_ == STATE_CLOSING) {
 190           ForceClose();  // T09
 191           return;
 192         }
 193         state_ = STATE_READING;  // T05
 194         break;
 195
 196       case STATE_CLOSED:
 197         ;
 198     }
 199   }
 200
 201  private:
 202   // Internal method used to close the buffer and notify the peer, if any.
 203   void ForceClose() {
 204     if (peer_) {
 205       peer_->Close();
 206       peer_ = NULL;
 207     }
 208     state_ = STATE_CLOSED;
 209   }
 210
 211   // Not owned.
 212   Socket* socket_from_;
 213   Socket* socket_to_;
 214
 215   int bytes_read_;
 216   int write_offset_;
 217   BufferedCopier* peer_;
 218   State state_;
 219   char buffer_[kBufferSize];
 220
 221   DISALLOW_COPY_AND_ASSIGN(BufferedCopier);
 222 };
 223
 224 Forwarder::Forwarder(scoped_ptr<Socket> socket1,
 225                      scoped_ptr<Socket> socket2)
 226     : socket1_(socket1.Pass()),
 227       socket2_(socket2.Pass()),
 228       buffer1_(new BufferedCopier(socket1_.get(), socket2_.get())),
 229       buffer2_(new BufferedCopier(socket2_.get(), socket1_.get())) {
 230   buffer1_->SetPeer(buffer2_.get());
 231   buffer2_->SetPeer(buffer1_.get());
 232 }
 233
 234 Forwarder::~Forwarder() {
 235   DCHECK(thread_checker_.CalledOnValidThread());
 236 }
 237
 238 void Forwarder::RegisterFDs(fd_set* read_fds, fd_set* write_fds, int* max_fd) {
 239   DCHECK(thread_checker_.CalledOnValidThread());
 240   buffer1_->PrepareSelect(read_fds, write_fds, max_fd);
 241   buffer2_->PrepareSelect(read_fds, write_fds, max_fd);
 242 }
 243
 244 void Forwarder::ProcessEvents(const fd_set& read_fds, const fd_set& write_fds) {
 245   DCHECK(thread_checker_.CalledOnValidThread());
 246   buffer1_->ProcessSelect(read_fds, write_fds);
 247   buffer2_->ProcessSelect(read_fds, write_fds);
 248 }
 249
 250 bool Forwarder::IsClosed() const {
 251   DCHECK(thread_checker_.CalledOnValidThread());
 252   return buffer1_->is_closed() && buffer2_->is_closed();
 253 }
 254
 255 void Forwarder::Shutdown() {
 256   DCHECK(thread_checker_.CalledOnValidThread());
 257   buffer1_->Close();
 258   buffer2_->Close();
 259 }
 260
 261 }  // namespace forwarder2