shaka-packager/tools/android/forwarder2/forwarder.cc

// Copyright (c) 2012 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 "tools/android/forwarder2/forwarder.h"

#include "base/basictypes.h"
#include "base/bind.h"
#include "base/logging.h"
#include "base/memory/ref_counted.h"
#include "base/posix/eintr_wrapper.h"
#include "base/single_thread_task_runner.h"
#include "tools/android/forwarder2/socket.h"

namespace forwarder2 {
namespace {

// Helper class to buffer reads and writes from one socket to another.
class BufferedCopier {
 public:
  // Does NOT own the pointers.
  BufferedCopier(Socket* socket_from,
                 Socket* socket_to)
      : socket_from_(socket_from),
        socket_to_(socket_to),
        bytes_read_(0),
        write_offset_(0) {
  }

  bool AddToReadSet(fd_set* read_fds) {
    if (bytes_read_ == 0)
      return socket_from_->AddFdToSet(read_fds);
    return false;
  }

  bool AddToWriteSet(fd_set* write_fds) {
    if (write_offset_ < bytes_read_)
      return socket_to_->AddFdToSet(write_fds);
    return false;
  }

  bool TryRead(const fd_set& read_fds) {
    if (!socket_from_->IsFdInSet(read_fds))
      return false;
    if (bytes_read_ != 0)  // Can't read.
      return false;
    int ret = socket_from_->Read(buffer_, kBufferSize);
    if (ret > 0) {
      bytes_read_ = ret;
      return true;
    }
    return false;
  }

  bool TryWrite(const fd_set& write_fds) {
    if (!socket_to_->IsFdInSet(write_fds))
      return false;
    if (write_offset_ >= bytes_read_)  // Nothing to write.
      return false;
    int ret = socket_to_->Write(buffer_ + write_offset_,
                                bytes_read_ - write_offset_);
    if (ret > 0) {
      write_offset_ += ret;
      if (write_offset_ == bytes_read_) {
        write_offset_ = 0;
        bytes_read_ = 0;
      }
      return true;
    }
    return false;
  }

 private:
  // Not owned.
  Socket* socket_from_;
  Socket* socket_to_;

  // A big buffer to let our file-over-http bridge work more like real file.
  static const int kBufferSize = 1024 * 128;
  int bytes_read_;
  int write_offset_;
  char buffer_[kBufferSize];

  DISALLOW_COPY_AND_ASSIGN(BufferedCopier);
};

// Internal class that wraps a helper thread to forward traffic between
// |socket1| and |socket2|. After creating a new instance, call its Start()
// method to launch operations. Thread stops automatically if one of the socket
// disconnects, but ensures that all buffered writes to the other, still alive,
// socket, are written first. When this happens, the instance will delete itself
// automatically.
// Note that the instance will always be destroyed on the same thread that
// created it.
class Forwarder {
 public:
  Forwarder(scoped_ptr<Socket> socket1, scoped_ptr<Socket> socket2)
      : socket1_(socket1.Pass()),
        socket2_(socket2.Pass()),
        destructor_runner_(base::MessageLoopProxy::current()),
        thread_("ForwarderThread") {
  }

  void Start() {
    thread_.Start();
    thread_.message_loop_proxy()->PostTask(
        FROM_HERE,
        base::Bind(&Forwarder::ThreadHandler, base::Unretained(this)));
  }

 private:
  void ThreadHandler() {
    const int nfds = Socket::GetHighestFileDescriptor(*socket1_, *socket2_) + 1;
    fd_set read_fds;
    fd_set write_fds;

    // Copy from socket1 to socket2
    BufferedCopier buffer1(socket1_.get(), socket2_.get());
    // Copy from socket2 to socket1
    BufferedCopier buffer2(socket2_.get(), socket1_.get());

    bool run = true;
    while (run) {
      FD_ZERO(&read_fds);
      FD_ZERO(&write_fds);

      buffer1.AddToReadSet(&read_fds);
      buffer2.AddToReadSet(&read_fds);
      buffer1.AddToWriteSet(&write_fds);
      buffer2.AddToWriteSet(&write_fds);

      if (HANDLE_EINTR(select(nfds, &read_fds, &write_fds, NULL, NULL)) <= 0) {
        PLOG(ERROR) << "select";
        break;
      }
      // When a socket in the read set closes the connection, select() returns
      // with that socket descriptor set as "ready to read".  When we call
      // TryRead() below, it will return false, but the while loop will continue
      // to run until all the write operations are finished, to make sure the
      // buffers are completely flushed out.

      // Keep running while we have some operation to do.
      run = buffer1.TryRead(read_fds);
      run = run || buffer2.TryRead(read_fds);
      run = run || buffer1.TryWrite(write_fds);
      run = run || buffer2.TryWrite(write_fds);
    }

    // Note that the thread that |destruction_runner_| runs tasks on could be
    // temporarily blocked on I/O (e.g. select()) therefore it is safer to close
    // the sockets now rather than relying on the destructor.
    socket1_.reset();
    socket2_.reset();

    // Note that base::Thread must be destroyed on the thread it was created on.
    destructor_runner_->DeleteSoon(FROM_HERE, this);
  }

  scoped_ptr<Socket> socket1_;
  scoped_ptr<Socket> socket2_;
  scoped_refptr<base::SingleThreadTaskRunner> destructor_runner_;
  base::Thread thread_;
};

}  // namespace

void StartForwarder(scoped_ptr<Socket> socket1, scoped_ptr<Socket> socket2) {
  (new Forwarder(socket1.Pass(), socket2.Pass()))->Start();
}

}  // namespace forwarder2