TcpStream.cc revision 24ca08a8
1efbfdb0cSShuo Chen#include "TcpStream.h"
2f41285cfSShuo Chen#include "InetAddress.h"
3efbfdb0cSShuo Chen
4355c2bf7SShuo Chen#include <errno.h>
5355c2bf7SShuo Chen#include <signal.h>
6355c2bf7SShuo Chen#include <unistd.h>
7355c2bf7SShuo Chen#include <sys/socket.h>
8355c2bf7SShuo Chen
9355c2bf7SShuo Chennamespace
10355c2bf7SShuo Chen{
11355c2bf7SShuo Chen
12355c2bf7SShuo Chenclass IgnoreSigPipe
13355c2bf7SShuo Chen{
14355c2bf7SShuo Chen public:
15355c2bf7SShuo Chen  IgnoreSigPipe()
16355c2bf7SShuo Chen  {
17355c2bf7SShuo Chen    ::signal(SIGPIPE, SIG_IGN);
18355c2bf7SShuo Chen  }
19355c2bf7SShuo Chen} initObj;
20355c2bf7SShuo Chen
21f41285cfSShuo Chenbool isSelfConnection(const Socket& sock)
22f41285cfSShuo Chen{
23f41285cfSShuo Chen  return sock.getLocalAddr() == sock.getPeerAddr();
24f41285cfSShuo Chen}
25f41285cfSShuo Chen
26355c2bf7SShuo Chen}
27355c2bf7SShuo Chen
28efbfdb0cSShuo ChenTcpStream::TcpStream(Socket&& sock)
29efbfdb0cSShuo Chen  : sock_(std::move(sock))
30efbfdb0cSShuo Chen{
31efbfdb0cSShuo Chen}
32355c2bf7SShuo Chen
33aafef3ccSShuo Chenint TcpStream::receiveAll(void* buf, int len)
34355c2bf7SShuo Chen{
35355c2bf7SShuo Chen  // FIXME: EINTR
36355c2bf7SShuo Chen  return ::recv(sock_.fd(), buf, len, MSG_WAITALL);
37355c2bf7SShuo Chen}
38355c2bf7SShuo Chen
39aafef3ccSShuo Chenint TcpStream::receiveSome(void* buf, int len)
40355c2bf7SShuo Chen{
41355c2bf7SShuo Chen  // FIXME: EINTR
42aafef3ccSShuo Chen  return sock_.read(buf, len);
43355c2bf7SShuo Chen}
44355c2bf7SShuo Chen
45355c2bf7SShuo Chenint TcpStream::sendAll(const void* buf, int len)
46355c2bf7SShuo Chen{
47355c2bf7SShuo Chen  int written = 0;
48355c2bf7SShuo Chen  while (written < len)
49355c2bf7SShuo Chen  {
507573bb11SShuo Chen    int nw = sock_.write(static_cast<const char*>(buf) + written, len - written);
517573bb11SShuo Chen    if (nw > 0)
52355c2bf7SShuo Chen    {
537573bb11SShuo Chen      written += nw;
54355c2bf7SShuo Chen    }
557573bb11SShuo Chen    else if (nw == 0)
56355c2bf7SShuo Chen    {
57355c2bf7SShuo Chen      break;
58355c2bf7SShuo Chen    }
597573bb11SShuo Chen    else if (errno != EINTR)  // nw < 0
60355c2bf7SShuo Chen    {
61355c2bf7SShuo Chen      break;
62355c2bf7SShuo Chen    }
63355c2bf7SShuo Chen  }
64355c2bf7SShuo Chen  return written;
65355c2bf7SShuo Chen}
66355c2bf7SShuo Chen
67355c2bf7SShuo Chenint TcpStream::sendSome(const void* buf, int len)
68355c2bf7SShuo Chen{
69355c2bf7SShuo Chen  // FIXME: EINTR
70aafef3ccSShuo Chen  return sock_.write(buf, len);
71355c2bf7SShuo Chen}
728f04b50cSShuo Chen
73569528b0SShuo Chenvoid TcpStream::setTcpNoDelay(bool on)
74a52ee0fdSShuo Chen{
75a52ee0fdSShuo Chen  sock_.setTcpNoDelay(on);
76a52ee0fdSShuo Chen}
77a52ee0fdSShuo Chen
78a52ee0fdSShuo Chenvoid TcpStream::shutdownWrite()
79a52ee0fdSShuo Chen{
80a52ee0fdSShuo Chen  sock_.shutdownWrite();
81a52ee0fdSShuo Chen}
82a52ee0fdSShuo Chen
838f04b50cSShuo ChenTcpStreamPtr TcpStream::connect(const InetAddress& serverAddr)
848f04b50cSShuo Chen{
85f41285cfSShuo Chen  return connectInternal(serverAddr, nullptr);
868f04b50cSShuo Chen}
878f04b50cSShuo Chen
88aafef3ccSShuo ChenTcpStreamPtr TcpStream::connect(const InetAddress& serverAddr, const InetAddress& localAddr)
898f04b50cSShuo Chen{
90f41285cfSShuo Chen  return connectInternal(serverAddr, &localAddr);
91f41285cfSShuo Chen}
92f41285cfSShuo Chen
93f41285cfSShuo ChenTcpStreamPtr TcpStream::connectInternal(const InetAddress& serverAddr, const InetAddress* localAddr)
94f41285cfSShuo Chen{
958f04b50cSShuo Chen  TcpStreamPtr stream;
9624ca08a8SShuo Chen  Socket sock(Socket::createTCP(serverAddr.family()));
97f41285cfSShuo Chen  if (localAddr)
98f41285cfSShuo Chen  {
99f41285cfSShuo Chen    sock.bindOrDie(*localAddr);
100f41285cfSShuo Chen  }
101f41285cfSShuo Chen  if (sock.connect(serverAddr) == 0 && !isSelfConnection(sock))
1028f04b50cSShuo Chen  {
1038f04b50cSShuo Chen    // FIXME: do poll(POLLOUT) to check errors
1048f04b50cSShuo Chen    stream.reset(new TcpStream(std::move(sock)));
1058f04b50cSShuo Chen  }
1068f04b50cSShuo Chen  return stream;
1078f04b50cSShuo Chen}
108