git.m6w6.name Git - awesomized/libmemcached/blob - libtest/client.cc

   1 /*  vim:expandtab:shiftwidth=2:tabstop=2:smarttab:
   2  *
   3  *  Data Differential YATL (i.e. libtest)  library
   4  *
   5  *  Copyright (C) 2012 Data Differential, http://datadifferential.com/
   6  *
   7  *  Redistribution and use in source and binary forms, with or without
   8  *  modification, are permitted provided that the following conditions are
   9  *  met:
  10  *
  11  *      * Redistributions of source code must retain the above copyright
  12  *  notice, this list of conditions and the following disclaimer.
  13  *
  14  *      * Redistributions in binary form must reproduce the above
  15  *  copyright notice, this list of conditions and the following disclaimer
  16  *  in the documentation and/or other materials provided with the
  17  *  distribution.
  18  *
  19  *      * The names of its contributors may not be used to endorse or
  20  *  promote products derived from this software without specific prior
  21  *  written permission.
  22  *
  23  *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  24  *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  25  *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  26  *  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  27  *  OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  28  *  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  29  *  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  30  *  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  31  *  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  32  *  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  33  *  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  34  *
  35  */
  36
  37 #include "libtest/yatlcon.h"
  38 #include <libtest/common.h>
  39
  40 #include <sys/types.h>
  41 #include <sys/socket.h>
  42 #include <unistd.h>
  43 #include <string>
  44
  45 #ifdef HAVE_POLL_H
  46 # include <poll.h>
  47 #endif
  48
  49 #ifndef HAVE_MSG_NOSIGNAL
  50 # define MSG_NOSIGNAL 0
  51 #endif
  52
  53 namespace libtest {
  54
  55 SimpleClient::SimpleClient(const std::string& hostname_, in_port_t port_) :
  56   _is_connected(false),
  57   _hostname(hostname_),
  58   _port(port_),
  59   sock_fd(INVALID_SOCKET),
  60   requested_message(1)
  61   {
  62   }
  63
  64 bool SimpleClient::ready(int event_)
  65 {
  66   struct pollfd fds[1];
  67   fds[0].fd= sock_fd;
  68   fds[0].events= event_;
  69   fds[0].revents= 0;
  70
  71   int timeout= 5000;
  72   if (_is_connected == false)
  73   {
  74     timeout= timeout * 30;
  75   }
  76
  77   int ready_fds= poll(fds, 1, timeout);
  78
  79   if (ready_fds == -1)
  80   {
  81     _error= strerror(errno);
  82     return false;
  83   }
  84   else if (ready_fds == 1)
  85   {
  86     if (fds[0].revents & (POLLERR | POLLHUP | POLLNVAL))
  87     {
  88       int err;
  89       socklen_t len= sizeof (err);
  90       // We replace errno with err if getsockopt() passes, but err has been
  91       // set.
  92       if (getsockopt(fds[0].fd, SOL_SOCKET, SO_ERROR, &err, &len) == 0)
  93       {
  94         // We check the value to see what happened wth the socket.
  95         if (err == 0)
  96         {
  97           _error= "getsockopt() returned no error but poll() indicated one existed";
  98           return false;
  99         }
 100         errno= err;
 101       }
 102       _error= strerror(errno);
 103
 104       return false;
 105     }
 106
 107     _is_connected= true;
 108     if (fds[0].revents & event_)
 109     {
 110       return true;
 111     }
 112   }
 113
 114   fatal_assert(ready_fds == 0);
 115   _error= "TIMEOUT";
 116
 117   return false;
 118 }
 119
 120 struct addrinfo* SimpleClient::lookup()
 121 {
 122   struct addrinfo *ai= NULL;
 123   struct addrinfo hints;
 124   memset(&hints, 0, sizeof(struct addrinfo));
 125   hints.ai_socktype= SOCK_STREAM;
 126   hints.ai_protocol= IPPROTO_TCP;
 127
 128   libtest::vchar_t service;
 129   service.resize(NI_MAXSERV);
 130   (void)snprintf(&service[0], service.size(), "%d", _port);
 131
 132   int getaddrinfo_error;
 133   if ((getaddrinfo_error= getaddrinfo(_hostname.c_str(), &service[0], &hints, &ai)) != 0)
 134   {
 135     if (getaddrinfo_error != EAI_SYSTEM)
 136     {
 137       _error= gai_strerror(getaddrinfo_error);
 138       return NULL;
 139     }
 140     else
 141     {
 142       _error= strerror(getaddrinfo_error);
 143       return NULL;
 144     }
 145   }
 146
 147   return ai;
 148 }
 149
 150 SimpleClient::~SimpleClient()
 151 {
 152   close_socket();
 153 }
 154
 155 void SimpleClient::close_socket()
 156 {
 157   if (sock_fd != INVALID_SOCKET)
 158   {
 159     close(sock_fd);
 160     sock_fd= INVALID_SOCKET;
 161   }
 162 }
 163
 164 bool SimpleClient::instance_connect()
 165 {
 166   _is_connected= false;
 167   struct addrinfo *ai;
 168   if ((ai= lookup()))
 169   {
 170     {
 171       struct addrinfo* address_info_next= ai;
 172
 173       while (address_info_next and sock_fd == INVALID_SOCKET)
 174       {
 175         if ((sock_fd= socket(address_info_next->ai_family, address_info_next->ai_socktype, address_info_next->ai_protocol)) != SOCKET_ERROR)
 176         {
 177           if (connect(sock_fd, address_info_next->ai_addr, address_info_next->ai_addrlen) == SOCKET_ERROR)
 178           {
 179             switch (errno)
 180             {
 181             case EINTR:
 182               close_socket();
 183               continue;
 184
 185             case EINPROGRESS: // nonblocking mode - first return
 186             case EALREADY: // nonblocking mode - subsequent returns
 187               continue; // Jump to while() and continue on
 188
 189
 190             case ECONNREFUSED:
 191             default:
 192               break;
 193             }
 194
 195             close_socket();
 196             _error= strerror(errno);
 197           }
 198         }
 199         else
 200         {
 201           FATAL(strerror(errno));
 202         }
 203         address_info_next= address_info_next->ai_next;
 204       }
 205
 206       freeaddrinfo(ai);
 207     }
 208
 209     if (sock_fd == INVALID_SOCKET)
 210     {
 211       fatal_assert(_error.size());
 212     }
 213
 214     return bool(sock_fd != INVALID_SOCKET);
 215   }
 216
 217   return false;
 218 }
 219
 220 bool SimpleClient::is_valid()
 221 {
 222   _error.clear();
 223   if (sock_fd == INVALID_SOCKET)
 224   {
 225     return instance_connect();
 226   }
 227
 228   return true;
 229 }
 230
 231 bool SimpleClient::message(const char* ptr, const size_t len)
 232 {
 233   if (is_valid())
 234   {
 235     if (ready(POLLOUT))
 236     {
 237       off_t offset= 0;
 238       do
 239       {
 240         ssize_t nw= send(sock_fd, ptr + offset, len - offset, MSG_NOSIGNAL);
 241         if (nw == -1)
 242         {
 243           if (errno != EINTR)
 244           {
 245             _error= strerror(errno);
 246             return false;
 247           }
 248         }
 249         else
 250         {
 251           offset += nw;
 252         }
 253       } while (offset < ssize_t(len));
 254
 255       return true;
 256     }
 257   }
 258
 259   fatal_assert(_error.size());
 260
 261   return false;
 262 }
 263
 264 bool SimpleClient::send_message(const std::string& arg)
 265 {
 266   if (message(arg.c_str(), arg.size()) == true)
 267   {
 268     return message("\r\n", 2);
 269   }
 270
 271   return false;
 272 }
 273
 274 bool SimpleClient::send_data(const libtest::vchar_t& message_, libtest::vchar_t& response_)
 275 {
 276   requested_message++;
 277   if (message(&message_[0], message_.size()))
 278   {
 279     return response(response_);
 280   }
 281
 282   return false;
 283 }
 284
 285 bool SimpleClient::send_message(const std::string& message_, std::string& response_)
 286 {
 287   requested_message++;
 288   if (send_message(message_))
 289   {
 290     return response(response_);
 291   }
 292
 293   return false;
 294 }
 295
 296 bool SimpleClient::response(libtest::vchar_t& response_)
 297 {
 298   response_.clear();
 299
 300   if (is_valid())
 301   {
 302     if (ready(POLLIN))
 303     {
 304       bool more= true;
 305       char buffer[2];
 306       buffer[1]= 0;
 307       do
 308       {
 309         ssize_t nr= recv(sock_fd, buffer, 1, MSG_NOSIGNAL);
 310         if (nr == -1)
 311         {
 312           if (errno != EINTR)
 313           {
 314             _error= strerror(errno);
 315             return false;
 316           }
 317         }
 318         else if (nr == 0)
 319         {
 320           close_socket();
 321           more= false;
 322         }
 323         else
 324         {
 325           response_.reserve(response_.size() + nr +1);
 326           fatal_assert(nr == 1);
 327           if (buffer[0] == '\n')
 328           {
 329             more= false;
 330           }
 331           response_.insert(response_.end(), buffer, buffer +nr);
 332         }
 333       } while (more);
 334
 335       return response_.size();
 336     }
 337   }
 338
 339   fatal_assert(_error.size());
 340   return false;
 341 }
 342
 343 bool SimpleClient::response(std::string& response_)
 344 {
 345   response_.clear();
 346
 347   if (is_valid())
 348   {
 349     if (ready(POLLIN))
 350     {
 351       bool more= true;
 352       char buffer[2];
 353       buffer[1]= 0;
 354       do
 355       {
 356         ssize_t nr= recv(sock_fd, buffer, 1, MSG_NOSIGNAL);
 357         if (nr == -1)
 358         {
 359           if (errno != EINTR)
 360           {
 361             _error= strerror(errno);
 362             return false;
 363           }
 364         }
 365         else if (nr == 0)
 366         {
 367           close_socket();
 368           more= false;
 369         }
 370         else
 371         {
 372           fatal_assert(nr == 1);
 373           if (buffer[0] == '\n')
 374           {
 375             more= false;
 376           }
 377           response_.append(buffer);
 378         }
 379       } while (more);
 380
 381       return response_.size();
 382     }
 383   }
 384
 385   fatal_assert(_error.size());
 386   return false;
 387 }
 388
 389 } // namespace libtest