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