silence dnsdist debugging output, deal with partial reads/writes

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 1dc72625841eea4eb09500886f4d5a0a658a749a..4347e0fa515286c8776aa5d50343070272bf1b4a 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -204,7 +204,8 @@ void* statThread(void*)
     uint64_t numQueries=0;
     for(unsigned int n=0; n < g_numdownstreams; ++n) {
       DownstreamState& dss = g_dstates[n];
-      cout<<'\t'<<dss.remote.toStringWithPort()<<": "<<dss.outstanding<<" outstanding, "<<(dss.queries - prev[n].queries)/interval <<" qps"<<endl;
+      if(g_verbose)
+       cout<<'\t'<<dss.remote.toStringWithPort()<<": "<<dss.outstanding<<" outstanding, "<<(dss.queries - prev[n].queries)/interval <<" qps"<<endl;
       outstanding += dss.outstanding;
       prev[n].queries = dss.queries;
       numQueries += dss.queries;
@@ -218,7 +219,8 @@ void* statThread(void*)
        }         
       }
     }
-    cout<<outstanding<<" outstanding queries, " <<(numQueries - lastQueries)/interval <<" qps"<<endl;
+    if(g_verbose)
+      cout<<outstanding<<" outstanding queries, " <<(numQueries - lastQueries)/interval <<" qps"<<endl;
     lastQueries=numQueries;
   }
   return 0;
@@ -249,31 +251,34 @@ int getTCPDownstream(DownstreamState** ds)
 }
 
 bool getMsgLen(int fd, uint16_t* len)
+try
 {
   uint16_t raw;
-  int ret = read(fd, &raw, 2);
+  int ret = readn2(fd, &raw, 2);
   if(ret != 2)
     return false;
   *len = ntohs(raw);
   return true;
 }
+catch(...) {
+   return false;
+}
 
 bool putMsgLen(int fd, uint16_t len)
+try
 {
   uint16_t raw = htons(len);
-  int ret = write(fd, &raw, 2);
+  int ret = writen2(fd, &raw, 2);
   return ret==2;
 }
+catch(...) {
+  return false;
+}
 
 
 /* spawn as many of these as required, they call Accept on a socket on which they will accept queries, and 
    they will initiate downstream connections to service them. An attempt is made to reuse existing
    connections.
-
-   As it stands, this code is both dangerous (doesn't check for partial reads/writes) and inefficient  
-   (does split writes). 
-
-   In addition, we leak filedescriptors. 
 */
 void* tcpClientThread(void* p)
 {
@@ -293,8 +298,10 @@ void* tcpClientThread(void* p)
 
       if(dsock == -1)
        dsock = getTCPDownstream(&ds);
-      else if(g_verbose)
-       cout <<"Reusing existing TCP connection"<<endl;
+      else {
+       if(g_verbose)
+         cout <<"Reusing existing TCP connection"<<endl;
+      }
 
       uint16_t qlen, rlen;
       for(;;) {
@@ -303,7 +310,7 @@ void* tcpClientThread(void* p)
 
        ds->queries++;
        char query[qlen];
-       int ret = read(client, query, qlen);
+       readn2(client, query, qlen);
       
       retry:; 
        if(!putMsgLen(dsock, qlen)) {
@@ -314,7 +321,7 @@ void* tcpClientThread(void* p)
          goto retry;
        }
       
-       ret = write(dsock, query, qlen);
+       writen2(dsock, query, qlen);
       
        if(!getMsgLen(dsock, &rlen)) {
          if(g_verbose)
@@ -325,10 +332,10 @@ void* tcpClientThread(void* p)
        }
 
        char answerbuffer[rlen];
-       ret = read(dsock, answerbuffer, rlen);
+       readn2(dsock, answerbuffer, rlen);
       
        putMsgLen(client, rlen);
-       ret = write(client, answerbuffer, rlen);
+       writen2(client, answerbuffer, rlen);
       }
       if(g_verbose)
        cout<<"Closing client connection"<<endl;

diff --git a/pdns/misc.cc b/pdns/misc.cc
index 74db1b846508c2ff87b1774c72e422768393c5f2..08375617a2bb427d133f06ba9af97ef68e977f58 100644 (file)
--- a/pdns/misc.cc
+++ b/pdns/misc.cc
@@ -70,6 +70,28 @@ int writen2(int fd, const void *buf, size_t count)
   return count;
 }
 
+int readn2(int fd, void* buffer, unsigned int len)
+{
+  unsigned int pos=0;
+  int res;
+  for(;;) {
+    res = read(fd, (char*)buffer + pos, len - pos);
+    if(res == 0) 
+      throw runtime_error("EOF while writing message");
+    if(res < 0) {
+      if (errno == EAGAIN)
+        throw std::runtime_error("used writen2 on non-blocking socket, got EAGAIN");
+      else
+        unixDie("failed in writen2");
+    } 
+    
+    pos+=res;
+    if(pos == len)
+      break;
+  }
+  return len;
+}
+
 
 string nowTime()
 {

diff --git a/pdns/misc.hh b/pdns/misc.hh
index 34cc02f94fba24e73e837b672519f156a8522a8f..bc764bfb55f30c501afdd78783555326ea7adaf0 100644 (file)
--- a/pdns/misc.hh
+++ b/pdns/misc.hh
@@ -158,7 +158,7 @@ vstringtok (Container &container, string const &in,
 
 int writen2(int fd, const void *buf, size_t count);
 inline int writen2(int fd, const std::string &s) { return writen2(fd, s.data(), s.size()); }
-
+int readn2(int fd, void* buffer, unsigned int len);
 
 const string toLower(const string &upper);
 const string toLowerCanonic(const string &upper);