forwarder.cc 7.4 KB

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