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