pdns/sstuff.hh

   1 /*
   2  * This file is part of PowerDNS or dnsdist.
   3  * Copyright -- PowerDNS.COM B.V. and its contributors
   4  *
   5  * This program is free software; you can redistribute it and/or modify
   6  * it under the terms of version 2 of the GNU General Public License as
   7  * published by the Free Software Foundation.
   8  *
   9  * In addition, for the avoidance of any doubt, permission is granted to
  10  * link this program with OpenSSL and to (re)distribute the binaries
  11  * produced as the result of such linking.
  12  *
  13  * This program is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16  * GNU General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU General Public License
  19  * along with this program; if not, write to the Free Software
  20  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  21  */
  22 #ifndef SSTUFF_HH
  23 #define SSTUFF_HH
  24
  25 #include <string>
  26 #include <sstream>
  27 #include <iostream>
  28 #include "iputils.hh"
  29 #include <errno.h>
  30 #include <sys/types.h>
  31 #include <unistd.h>
  32 #include <sys/socket.h>
  33 #include <netinet/in.h>
  34 #include <arpa/inet.h>
  35 #include <sys/select.h>
  36 #include <fcntl.h>
  37 #include <stdexcept>
  38
  39 #include <boost/utility.hpp>
  40 #include <csignal>
  41 #include "namespaces.hh"
  42 #include "namespaces.hh"
  43
  44
  45 typedef int ProtocolType; //!< Supported protocol types
  46
  47 //! Representation of a Socket and many of the Berkeley functions available
  48 class Socket : public boost::noncopyable
  49 {
  50   Socket(int fd): d_socket(fd)
  51   {
  52   }
  53
  54 public:
  55   //! Construct a socket of specified address family and socket type.
  56   Socket(int af, int st, ProtocolType pt=0)
  57   {
  58     if((d_socket=socket(af, st, pt))<0)
  59       throw NetworkError(strerror(errno));
  60     setCloseOnExec(d_socket);
  61   }
  62
  63   Socket(Socket&& rhs): d_buffer(std::move(rhs.d_buffer)), d_socket(rhs.d_socket)
  64   {
  65     rhs.d_socket = -1;
  66   }
  67
  68   ~Socket()
  69   {
  70     try {
  71       if (d_socket != -1) {
  72         closesocket(d_socket);
  73       }
  74     }
  75     catch(const PDNSException& e) {
  76     }
  77   }
  78
  79   //! If the socket is capable of doing so, this function will wait for a connection
  80   std::unique_ptr<Socket> accept()
  81   {
  82     struct sockaddr_in remote;
  83     socklen_t remlen=sizeof(remote);
  84     memset(&remote, 0, sizeof(remote));
  85     int s=::accept(d_socket, reinterpret_cast<sockaddr *>(&remote), &remlen);
  86     if(s<0) {
  87       if(errno==EAGAIN)
  88         return nullptr;
  89
  90       throw NetworkError("Accepting a connection: "+string(strerror(errno)));
  91     }
  92
  93     return std::unique_ptr<Socket>(new Socket(s));
  94   }
  95
  96   //! Get remote address
  97   bool getRemote(ComboAddress &remote) {
  98     socklen_t remotelen=sizeof(remote);
  99     return (getpeername(d_socket, reinterpret_cast<struct sockaddr *>(&remote), &remotelen) >= 0);
 100   }
 101
 102   //! Check remote address against netmaskgroup ng
 103   bool acl(const NetmaskGroup &ng)
 104   {
 105     ComboAddress remote;
 106     if (getRemote(remote))
 107       return ng.match(remote);
 108
 109     return false;
 110   }
 111
 112   //! Set the socket to non-blocking
 113   void setNonBlocking()
 114   {
 115     ::setNonBlocking(d_socket);
 116   }
 117
 118   //! Set the socket to blocking
 119   void setBlocking()
 120   {
 121     ::setBlocking(d_socket);
 122   }
 123
 124   void setReuseAddr()
 125   {
 126     try {
 127       ::setReuseAddr(d_socket);
 128     } catch (const PDNSException &e) {
 129       throw NetworkError(e.reason);
 130     }
 131   }
 132
 133   //! Bind the socket to a specified endpoint
 134   void bind(const ComboAddress &local, bool reuseaddr=true)
 135   {
 136     int tmp=1;
 137     if(reuseaddr && setsockopt(d_socket, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<char*>(&tmp), sizeof tmp)<0)
 138       throw NetworkError(string("Setsockopt failed: ")+strerror(errno));
 139
 140     if(::bind(d_socket, reinterpret_cast<const struct sockaddr *>(&local), local.getSocklen())<0)
 141       throw NetworkError("While binding: "+string(strerror(errno)));
 142   }
 143
 144   //! Connect the socket to a specified endpoint
 145   void connect(const ComboAddress &ep, int timeout=0)
 146   {
 147     SConnectWithTimeout(d_socket, ep, timeout);
 148   }
 149
 150
 151   //! For datagram sockets, receive a datagram and learn where it came from
 152   /** For datagram sockets, receive a datagram and learn where it came from
 153       \param dgram Will be filled with the datagram
 154       \param ep Will be filled with the origin of the datagram */
 155   void recvFrom(string &dgram, ComboAddress &ep)
 156   {
 157     socklen_t remlen = sizeof(ep);
 158     ssize_t bytes;
 159     d_buffer.resize(s_buflen);
 160     if((bytes=recvfrom(d_socket, &d_buffer[0], s_buflen, 0, reinterpret_cast<sockaddr *>(&ep) , &remlen)) <0)
 161       throw NetworkError("After recvfrom: "+string(strerror(errno)));
 162
 163     dgram.assign(d_buffer, 0, static_cast<size_t>(bytes));
 164   }
 165
 166   bool recvFromAsync(string &dgram, ComboAddress &ep)
 167   {
 168     struct sockaddr_in remote;
 169     socklen_t remlen = sizeof(remote);
 170     ssize_t bytes;
 171     d_buffer.resize(s_buflen);
 172     if((bytes=recvfrom(d_socket, &d_buffer[0], s_buflen, 0, reinterpret_cast<sockaddr *>(&remote), &remlen))<0) {
 173       if(errno!=EAGAIN) {
 174         throw NetworkError("After async recvfrom: "+string(strerror(errno)));
 175       }
 176       else {
 177         return false;
 178       }
 179     }
 180     dgram.assign(d_buffer, 0, static_cast<size_t>(bytes));
 181     return true;
 182   }
 183
 184
 185   //! For datagram sockets, send a datagram to a destination
 186   void sendTo(const char* msg, size_t len, const ComboAddress &ep)
 187   {
 188     if(sendto(d_socket, msg, len, 0, reinterpret_cast<const sockaddr *>(&ep), ep.getSocklen())<0)
 189       throw NetworkError("After sendto: "+string(strerror(errno)));
 190   }
 191
 192   //! For connected datagram sockets, send a datagram
 193   void send(const std::string& msg)
 194   {
 195     if(::send(d_socket, msg.c_str(), msg.size(), 0)<0)
 196       throw NetworkError("After send: "+string(strerror(errno)));
 197   }
 198
 199
 200   /** For datagram sockets, send a datagram to a destination
 201       \param dgram The datagram
 202       \param ep The intended destination of the datagram */
 203   void sendTo(const string &dgram, const ComboAddress &ep)
 204   {
 205     sendTo(dgram.c_str(), dgram.length(), ep);
 206   }
 207
 208
 209   //! Write this data to the socket, taking care that all bytes are written out
 210   void writen(const string &data)
 211   {
 212     if(data.empty())
 213       return;
 214
 215     size_t toWrite=data.length();
 216     ssize_t res;
 217     const char *ptr=data.c_str();
 218
 219     do {
 220       res=::send(d_socket, ptr, toWrite, 0);
 221       if(res<0)
 222         throw NetworkError("Writing to a socket: "+string(strerror(errno)));
 223       if(!res)
 224         throw NetworkError("EOF on socket");
 225       toWrite -= static_cast<size_t>(res);
 226       ptr += static_cast<size_t>(res);
 227     } while(toWrite);
 228
 229   }
 230
 231   //! tries to write toWrite bytes from ptr to the socket
 232   /** tries to write toWrite bytes from ptr to the socket, but does not make sure they al get written out
 233       \param ptr Location to write from
 234       \param toWrite number of bytes to try
 235   */
 236   size_t tryWrite(const char *ptr, size_t toWrite)
 237   {
 238     ssize_t res;
 239     res=::send(d_socket,ptr,toWrite,0);
 240     if(res==0)
 241       throw NetworkError("EOF on writing to a socket");
 242
 243     if(res>0)
 244       return res;
 245
 246     if(errno==EAGAIN)
 247       return 0;
 248
 249     throw NetworkError("Writing to a socket: "+string(strerror(errno)));
 250   }
 251
 252   //! Writes toWrite bytes from ptr to the socket
 253   /** Writes toWrite bytes from ptr to the socket. Returns how many bytes were written */
 254   size_t write(const char *ptr, size_t toWrite)
 255   {
 256     ssize_t res;
 257     res=::send(d_socket,ptr,toWrite,0);
 258     if(res<0) {
 259       throw NetworkError("Writing to a socket: "+string(strerror(errno)));
 260     }
 261     return res;
 262   }
 263
 264   void writenWithTimeout(const void *buffer, size_t n, int timeout)
 265   {
 266     size_t bytes=n;
 267     const char *ptr = reinterpret_cast<const char*>(buffer);
 268     ssize_t ret;
 269     while(bytes) {
 270       ret=::write(d_socket, ptr, bytes);
 271       if(ret < 0) {
 272         if(errno==EAGAIN) {
 273           ret=waitForRWData(d_socket, false, timeout, 0);
 274           if(ret < 0)
 275             throw NetworkError("Waiting for data write");
 276           if(!ret)
 277             throw NetworkError("Timeout writing data");
 278           continue;
 279         }
 280         else
 281           throw NetworkError("Writing data: "+stringerror());
 282       }
 283       if(!ret) {
 284         throw NetworkError("Did not fulfill TCP write due to EOF");
 285       }
 286
 287       ptr += static_cast<size_t>(ret);
 288       bytes -= static_cast<size_t>(ret);
 289     }
 290   }
 291
 292   //! reads one character from the socket
 293   int getChar()
 294   {
 295     char c;
 296
 297     ssize_t res=::recv(d_socket,&c,1,0);
 298     if(res)
 299       return c;
 300     return -1;
 301   }
 302
 303   void getline(string &data)
 304   {
 305     data="";
 306     int c;
 307     while((c=getChar())!=-1) {
 308       data+=(char)c;
 309       if(c=='\n')
 310         break;
 311     }
 312   }
 313
 314   //! Reads a block of data from the socket to a string
 315   void read(string &data)
 316   {
 317     d_buffer.resize(s_buflen);
 318     ssize_t res=::recv(d_socket, &d_buffer[0], s_buflen, 0);
 319     if(res<0)
 320       throw NetworkError("Reading from a socket: "+string(strerror(errno)));
 321     data.assign(d_buffer, 0, static_cast<size_t>(res));
 322   }
 323
 324   //! Reads a block of data from the socket to a block of memory
 325   size_t read(char *buffer, size_t bytes)
 326   {
 327     ssize_t res=::recv(d_socket, buffer, bytes, 0);
 328     if(res<0)
 329       throw NetworkError("Reading from a socket: "+string(strerror(errno)));
 330     return static_cast<size_t>(res);
 331   }
 332
 333   ssize_t readWithTimeout(char* buffer, size_t n, int timeout)
 334   {
 335     int err = waitForRWData(d_socket, true, timeout, 0);
 336
 337     if(err == 0)
 338       throw NetworkError("timeout reading");
 339     if(err < 0)
 340       throw NetworkError("nonblocking read failed: "+string(strerror(errno)));
 341
 342     return read(buffer, n);
 343   }
 344
 345   //! Sets the socket to listen with a default listen backlog of 10 pending connections
 346   void listen(unsigned int length=10)
 347   {
 348     if(::listen(d_socket,length)<0)
 349       throw NetworkError("Setting socket to listen: "+string(strerror(errno)));
 350   }
 351
 352   //! Returns the internal file descriptor of the socket
 353   int getHandle() const
 354   {
 355     return d_socket;
 356   }
 357
 358 private:
 359   static const size_t s_buflen{4096};
 360   std::string d_buffer;
 361   int d_socket;
 362 };
 363
 364
 365 #endif