// Copyright 2012 The Chromium Authors
// 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 <utility>
#include "base/check.h"
#include "base/posix/eintr_wrapper.h"
#include "tools/android/forwarder2/socket.h"
namespace forwarder2 {
namespace {
const int kBufferSize = 32 * 1024;
} // namespace
// Helper class to buffer reads and writes from one socket to another.
// Each implements a small buffer connected two one input socket, and
// one output socket.
//
// socket_from_ ---> [BufferedCopier] ---> socket_to_
//
// These objects are used in a pair to handle duplex traffic, as in:
//
// -------> [BufferedCopier_1] --->
// | |
// socket_1 * * socket_2
// | |
// <------ [BufferedCopier_2] <----
//
// When a BufferedCopier is in the READING state (see below), it only listens
// to events on its input socket, and won't detect when its output socket
// disconnects. To work around this, its peer will call its Close() method
// when that happens.
class Forwarder::BufferedCopier {
public:
// Possible states:
// READING - Empty buffer and Waiting for input.
// WRITING - Data in buffer, and waiting for output.
// CLOSING - Like WRITING, but do not try to read after that.
// CLOSED - Completely closed.
//
// State transitions are:
//
// T01: READING ---[receive data]---> WRITING
// T02: READING ---[error on input socket]---> CLOSED
// T03: READING ---[Close() call]---> CLOSED
//
// T04: WRITING ---[write partial data]---> WRITING
// T05: WRITING ---[write all data]----> READING
// T06: WRITING ---[error on output socket]----> CLOSED
// T07: WRITING ---[Close() call]---> CLOSING
//
// T08: CLOSING ---[write partial data]---> CLOSING
// T09: CLOSING ---[write all data]----> CLOSED
// T10: CLOSING ---[Close() call]---> CLOSING
// T11: CLOSING ---[error on output socket] ---> CLOSED
//
enum State {
STATE_READING = 0,
STATE_WRITING = 1,
STATE_CLOSING = 2,
STATE_CLOSED = 3,
};
// 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),
peer_(NULL),
state_(STATE_READING) {}
BufferedCopier(const BufferedCopier&) = delete;
BufferedCopier& operator=(const BufferedCopier&) = delete;
// Sets the 'peer_' field pointing to the other BufferedCopier in a pair.
void SetPeer(BufferedCopier* peer) {
DCHECK(!peer_);
peer_ = peer;
}
bool is_closed() const { return state_ == STATE_CLOSED; }
// Gently asks to close a buffer. Called either by the peer or the forwarder.
void Close() {
switch (state_) {
case STATE_READING:
state_ = STATE_CLOSED; // T03
break;
case STATE_WRITING:
state_ = STATE_CLOSING; // T07
break;
case STATE_CLOSING:
break; // T10
case STATE_CLOSED:
break;
}
}
// Call this before select(). This updates |read_fds|,
// |write_fds| and |max_fd| appropriately *if* the buffer isn't closed.
void PrepareSelect(fd_set* read_fds, fd_set* write_fds, int* max_fd) {
int fd;
switch (state_) {
case STATE_READING:
DCHECK(bytes_read_ == 0);
DCHECK(write_offset_ == 0);
fd = socket_from_->fd();
if (fd < 0) {
ForceClose(); // T02
return;
}
FD_SET(fd, read_fds);
break;
case STATE_WRITING:
case STATE_CLOSING:
DCHECK(bytes_read_ > 0);
DCHECK(write_offset_ < bytes_read_);
fd = socket_to_->fd();
if (fd < 0) {
ForceClose(); // T06
return;
}
FD_SET(fd, write_fds);
break;
case STATE_CLOSED:
return;
}
*max_fd = std::max(*max_fd, fd);
}
// Call this after a select() call to operate over the buffer.
void ProcessSelect(const fd_set& read_fds, const fd_set& write_fds) {
int fd;
int ret;
// With FORTIFY_SOURCE, FD_ISSET is implemented as a function that takes a
// non-const fd_set*. Make a copy of the passed arguments so we can safely
// take a reference.
fd_set read_fds_copy = read_fds;
fd_set write_fds_copy = write_fds;
switch (state_) {
case STATE_READING:
fd = socket_from_->fd();
if (fd < 0) {
state_ = STATE_CLOSED; // T02
return;
}
if (!FD_ISSET(fd, &read_fds_copy))
return;
ret = socket_from_->NonBlockingRead(buffer_, kBufferSize);
if (ret <= 0) {
ForceClose(); // T02
return;
}
bytes_read_ = ret;
write_offset_ = 0;
state_ = STATE_WRITING; // T01
break;
case STATE_WRITING:
case STATE_CLOSING:
fd = socket_to_->fd();
if (fd < 0) {
ForceClose(); // T06 + T11
return;
}
if (!FD_ISSET(fd, &write_fds_copy))
return;
ret = socket_to_->NonBlockingWrite(buffer_ + write_offset_,
bytes_read_ - write_offset_);
if (ret <= 0) {
ForceClose(); // T06 + T11
return;
}
write_offset_ += ret;
if (write_offset_ < bytes_read_)
return; // T08 + T04
write_offset_ = 0;
bytes_read_ = 0;
if (state_ == STATE_CLOSING) {
ForceClose(); // T09
return;
}
state_ = STATE_READING; // T05
break;
case STATE_CLOSED:
break;
}
}
private:
// Internal method used to close the buffer and notify the peer, if any.
void ForceClose() {
if (peer_) {
peer_->Close();
peer_ = NULL;
}
state_ = STATE_CLOSED;
}
// Not owned.
Socket* socket_from_;
Socket* socket_to_;
int bytes_read_;
int write_offset_;
BufferedCopier* peer_;
State state_;
char buffer_[kBufferSize];
};
Forwarder::Forwarder(std::unique_ptr<Socket> socket1,
std::unique_ptr<Socket> socket2)
: socket1_(std::move(socket1)),
socket2_(std::move(socket2)),
buffer1_(new BufferedCopier(socket1_.get(), socket2_.get())),
buffer2_(new BufferedCopier(socket2_.get(), socket1_.get())) {
buffer1_->SetPeer(buffer2_.get());
buffer2_->SetPeer(buffer1_.get());
}
Forwarder::~Forwarder() {
DCHECK(thread_checker_.CalledOnValidThread());
}
void Forwarder::RegisterFDs(fd_set* read_fds, fd_set* write_fds, int* max_fd) {
DCHECK(thread_checker_.CalledOnValidThread());
buffer1_->PrepareSelect(read_fds, write_fds, max_fd);
buffer2_->PrepareSelect(read_fds, write_fds, max_fd);
}
void Forwarder::ProcessEvents(const fd_set& read_fds, const fd_set& write_fds) {
DCHECK(thread_checker_.CalledOnValidThread());
buffer1_->ProcessSelect(read_fds, write_fds);
buffer2_->ProcessSelect(read_fds, write_fds);
}
bool Forwarder::IsClosed() const {
DCHECK(thread_checker_.CalledOnValidThread());
return buffer1_->is_closed() && buffer2_->is_closed();
}
void Forwarder::Shutdown() {
DCHECK(thread_checker_.CalledOnValidThread());
buffer1_->Close();
buffer2_->Close();
}
} // namespace forwarder2