pdns/sstuff.hh

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