Merge common socket code

7 files changed, 255 insertions, 216 deletions

diff --git a/common/network/CMakeLists.txt b/common/network/CMakeLists.txt
index a63df96b..d00ca452 100644
--- a/common/network/CMakeLists.txt
+++ b/common/network/CMakeLists.txt
@@ -1,6 +1,7 @@
 include_directories(${CMAKE_SOURCE_DIR}/common)
 
 set(NETWORK_SOURCES
+  Socket.cxx
   TcpSocket.cxx)
 
 if(NOT WIN32)
diff --git a/common/network/Socket.cxx b/common/network/Socket.cxx
new file mode 100644
index 00000000..9dd8bfe1
--- /dev/null
+++ b/common/network/Socket.cxx
@@ -0,0 +1,183 @@
+/* Copyright (C) 2002-2005 RealVNC Ltd.  All Rights Reserved.
+ * 
+ * This is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation; either version 2 of the License, or
+ * (at your option) any later version.
+ * 
+ * This software is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ * 
+ * You should have received a copy of the GNU General Public License
+ * along with this software; if not, write to the Free Software
+ * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307,
+ * USA.
+ */
+
+#ifdef HAVE_CONFIG_H
+#include <config.h>
+#endif
+
+#ifdef WIN32
+//#include <io.h>
+#include <winsock2.h>
+#include <ws2tcpip.h>
+#define errorNumber WSAGetLastError()
+#else
+#define errorNumber errno
+#define closesocket close
+#include <sys/socket.h>
+#endif
+
+#include <unistd.h>
+#include <signal.h>
+#include <fcntl.h>
+#include <errno.h>
+
+#include <network/Socket.h>
+
+using namespace network;
+
+// -=- Socket initialisation
+static bool socketsInitialised = false;
+void network::initSockets() {
+  if (socketsInitialised)
+    return;
+#ifdef WIN32
+  WORD requiredVersion = MAKEWORD(2,0);
+  WSADATA initResult;
+  
+  if (WSAStartup(requiredVersion, &initResult) != 0)
+    throw SocketException("unable to initialise Winsock2", errorNumber);
+#else
+  signal(SIGPIPE, SIG_IGN);
+#endif
+  socketsInitialised = true;
+}
+
+bool network::isSocketListening(int sock)
+{
+  int listening = 0;
+  socklen_t listening_size = sizeof(listening);
+  if (getsockopt(sock, SOL_SOCKET, SO_ACCEPTCONN,
+                 (char *)&listening, &listening_size) < 0)
+    return false;
+  return listening != 0;
+}
+
+Socket::Socket(int fd)
+  : instream(0), outstream(0),
+    isShutdown_(false), queryConnection(false)
+{
+  initSockets();
+  setFd(fd);
+}
+
+Socket::Socket()
+  : instream(0), outstream(0),
+    isShutdown_(false), queryConnection(false)
+{
+  initSockets();
+}
+
+Socket::~Socket()
+{
+  if (instream && outstream)
+    closesocket(getFd());
+  delete instream;
+  delete outstream;
+}
+
+// if shutdown() is overridden then the override MUST call on to here
+void Socket::shutdown()
+{
+  isShutdown_ = true;
+  ::shutdown(getFd(), 2);
+}
+
+bool Socket::isShutdown() const
+{
+  return isShutdown_;
+}
+
+// Was there a "?" in the ConnectionFilter used to accept this Socket?
+void Socket::setRequiresQuery()
+{
+  queryConnection = true;
+}
+
+bool Socket::requiresQuery() const
+{
+  return queryConnection;
+}
+
+void Socket::setFd(int fd)
+{
+#ifndef WIN32
+  // - By default, close the socket on exec()
+  fcntl(fd, F_SETFD, FD_CLOEXEC);
+#endif
+
+  instream = new rdr::FdInStream(fd);
+  outstream = new rdr::FdOutStream(fd);
+  isShutdown_ = false;
+}
+
+SocketListener::SocketListener(int fd)
+  : fd(fd), filter(0)
+{
+  initSockets();
+}
+
+SocketListener::SocketListener()
+  : fd(-1), filter(0)
+{
+  initSockets();
+}
+
+SocketListener::~SocketListener()
+{
+  if (fd != -1)
+    closesocket(fd);
+}
+
+void SocketListener::shutdown()
+{
+#ifdef WIN32
+  closesocket(fd);
+  fd = -1;
+#else
+  ::shutdown(fd, 2);
+#endif
+}
+
+Socket* SocketListener::accept() {
+  int new_sock = -1;
+
+  // Accept an incoming connection
+  if ((new_sock = ::accept(fd, 0, 0)) < 0)
+    throw SocketException("unable to accept new connection", errorNumber);
+
+  // Create the socket object & check connection is allowed
+  Socket* s = createSocket(new_sock);
+  if (filter && !filter->verifyConnection(s)) {
+    delete s;
+    return NULL;
+  }
+
+  return s;
+}
+
+void SocketListener::listen(int sock)
+{
+  // - Set it to be a listening socket
+  if (::listen(sock, 5) < 0) {
+    int e = errorNumber;
+    closesocket(sock);
+    throw SocketException("unable to set socket to listening mode", e);
+  }
+
+  fd = sock;
+}
diff --git a/common/network/Socket.h b/common/network/Socket.h
index 39ac39b4..bfda8a57 100644
--- a/common/network/Socket.h
+++ b/common/network/Socket.h
@@ -30,24 +30,22 @@
 
 namespace network {
 
+  void initSockets();
+
+  bool isSocketListening(int sock);
+
   class Socket {
   public:
-    Socket(int fd)
-      : instream(new rdr::FdInStream(fd)),
-      outstream(new rdr::FdOutStream(fd)),
-      isShutdown_(false),
-      queryConnection(false) {}
-    virtual ~Socket() {
-      delete instream;
-      delete outstream;
-    }
+    Socket(int fd);
+    virtual ~Socket();
+
     rdr::FdInStream &inStream() {return *instream;}
     rdr::FdOutStream &outStream() {return *outstream;}
     int getFd() {return outstream->getFd();}
 
-    // if shutdown() is overridden then the override MUST call on to here
-    virtual void shutdown() {isShutdown_ = true;}
-    bool isShutdown() const {return isShutdown_;}
+    void shutdown();
+    bool isShutdown() const;
+
     virtual bool cork(bool enable) = 0;
 
     // information about the remote end of the socket
@@ -55,15 +53,15 @@ namespace network {
     virtual char* getPeerEndpoint() = 0; // <address>::<port>
 
     // Was there a "?" in the ConnectionFilter used to accept this Socket?
-    void setRequiresQuery() {queryConnection = true;}
-    bool requiresQuery() const {return queryConnection;}
+    void setRequiresQuery();
+    bool requiresQuery() const;
 
   protected:
-    Socket() : instream(0), outstream(0),
-      isShutdown_(false), queryConnection(false) {}
-    Socket(rdr::FdInStream* i, rdr::FdOutStream* o)
-      : instream(i), outstream(o),
-      isShutdown_(false), queryConnection(false) {}
+    Socket();
+
+    void setFd(int fd);
+
+  private:
     rdr::FdInStream* instream;
     rdr::FdOutStream* outstream;
     bool isShutdown_;
@@ -78,22 +76,32 @@ namespace network {
 
   class SocketListener {
   public:
-    SocketListener() : fd(0), filter(0) {}
-    virtual ~SocketListener() {}
+    SocketListener(int fd);
+    virtual ~SocketListener();
 
     // shutdown() stops the socket from accepting further connections
-    virtual void shutdown() = 0;
+    void shutdown();
 
     // accept() returns a new Socket object if there is a connection
     // attempt in progress AND if the connection passes the filter
     // if one is installed.  Otherwise, returns 0.
-    virtual Socket* accept() = 0;
+    Socket* accept();
 
     virtual int getMyPort() = 0;
 
     // setFilter() applies the specified filter to all new connections
     void setFilter(ConnectionFilter* f) {filter = f;}
     int getFd() {return fd;}
+
+  protected:
+    SocketListener();
+
+    void listen(int fd);
+
+    // createSocket() should create a new socket of the correct class
+    // for the given file descriptor
+    virtual Socket* createSocket(int fd) = 0;
+
   protected:
     int fd;
     ConnectionFilter* filter;
diff --git a/common/network/TcpSocket.cxx b/common/network/TcpSocket.cxx
index 555ce215..51d77c76 100644
--- a/common/network/TcpSocket.cxx
+++ b/common/network/TcpSocket.cxx
@@ -33,8 +33,6 @@
 #include <netinet/tcp.h>
 #include <netdb.h>
 #include <errno.h>
-#include <signal.h>
-#include <fcntl.h>
 #endif
 
 #include <stdlib.h>
@@ -97,30 +95,27 @@ int network::findFreeTcpPort (void)
   return ntohs(addr.sin_port);
 }
 
+int network::getSockPort(int sock)
+{
+  vnc_sockaddr_t sa;
+  socklen_t sa_size = sizeof(sa);
+  if (getsockname(sock, &sa.u.sa, &sa_size) < 0)
+    return 0;
 
-// -=- Socket initialisation
-static bool socketsInitialised = false;
-static void initSockets() {
-  if (socketsInitialised)
-    return;
-#ifdef WIN32
-  WORD requiredVersion = MAKEWORD(2,0);
-  WSADATA initResult;
-  
-  if (WSAStartup(requiredVersion, &initResult) != 0)
-    throw SocketException("unable to initialise Winsock2", errorNumber);
-#else
-  signal(SIGPIPE, SIG_IGN);
-#endif
-  socketsInitialised = true;
+  switch (sa.u.sa.sa_family) {
+  case AF_INET6:
+    return ntohs(sa.u.sin6.sin6_port);
+  default:
+    return ntohs(sa.u.sin.sin_port);
+  }
 }
 
-
 // -=- TcpSocket
 
-TcpSocket::TcpSocket(int sock)
-  : Socket(new FdInStream(sock), new FdOutStream(sock))
+TcpSocket::TcpSocket(int sock) : Socket(sock)
 {
+  // Disable Nagle's algorithm, to reduce latency
+  enableNagles(false);
 }
 
 TcpSocket::TcpSocket(const char *host, int port)
@@ -129,7 +124,6 @@ TcpSocket::TcpSocket(const char *host, int port)
   struct addrinfo *ai, *current, hints;
 
   // - Create a socket
-  initSockets();
 
   memset(&hints, 0, sizeof(struct addrinfo));
   hints.ai_family = AF_UNSPEC;
@@ -211,21 +205,11 @@ TcpSocket::TcpSocket(const char *host, int port)
       throw SocketException("unable connect to socket", err);
   }
 
-#ifndef WIN32
-  // - By default, close the socket on exec()
-  fcntl(sock, F_SETFD, FD_CLOEXEC);
-#endif
+  // Take proper ownership of the socket
+  setFd(sock);
 
   // Disable Nagle's algorithm, to reduce latency
-  enableNagles(sock, false);
-
-  // Create the input and output streams
-  instream = new FdInStream(sock);
-  outstream = new FdOutStream(sock);
-}
-
-TcpSocket::~TcpSocket() {
-  closesocket(getFd());
+  enableNagles(false);
 }
 
 char* TcpSocket::getPeerAddress() {
@@ -293,15 +277,9 @@ char* TcpSocket::getPeerEndpoint() {
   return buffer;
 }
 
-void TcpSocket::shutdown()
-{
-  Socket::shutdown();
-  ::shutdown(getFd(), 2);
-}
-
-bool TcpSocket::enableNagles(int sock, bool enable) {
+bool TcpSocket::enableNagles(bool enable) {
   int one = enable ? 0 : 1;
-  if (setsockopt(sock, IPPROTO_TCP, TCP_NODELAY,
+  if (setsockopt(getFd(), IPPROTO_TCP, TCP_NODELAY,
                  (char *)&one, sizeof(one)) < 0) {
     int e = errorNumber;
     vlog.error("unable to setsockopt TCP_NODELAY: %d", e);
@@ -321,34 +299,8 @@ bool TcpSocket::cork(bool enable) {
 #endif
 }
 
-bool TcpSocket::isListening(int sock)
+TcpListener::TcpListener(int sock) : SocketListener(sock)
 {
-  int listening = 0;
-  socklen_t listening_size = sizeof(listening);
-  if (getsockopt(sock, SOL_SOCKET, SO_ACCEPTCONN,
-                 (char *)&listening, &listening_size) < 0)
-    return false;
-  return listening != 0;
-}
-
-int TcpSocket::getSockPort(int sock)
-{
-  vnc_sockaddr_t sa;
-  socklen_t sa_size = sizeof(sa);
-  if (getsockname(sock, &sa.u.sa, &sa_size) < 0)
-    return 0;
-
-  switch (sa.u.sa.sa_family) {
-  case AF_INET6:
-    return ntohs(sa.u.sin6.sin6_port);
-  default:
-    return ntohs(sa.u.sin.sin_port);
-  }
-}
-
-TcpListener::TcpListener(int sock)
-{
-  fd = sock;
 }
 
 TcpListener::TcpListener(const struct sockaddr *listenaddr,
@@ -358,8 +310,6 @@ TcpListener::TcpListener(const struct sockaddr *listenaddr,
   vnc_sockaddr_t sa;
   int sock;
 
-  initSockets();
-
   if ((sock = socket (listenaddr->sa_family, SOCK_STREAM, 0)) < 0)
     throw SocketException("unable to create listening socket", errorNumber);
 
@@ -397,53 +347,11 @@ TcpListener::TcpListener(const struct sockaddr *listenaddr,
     throw SocketException("failed to bind socket", e);
   }
 
-  // - Set it to be a listening socket
-  if (listen(sock, 5) < 0) {
-    int e = errorNumber;
-    closesocket(sock);
-    throw SocketException("unable to set socket to listening mode", e);
-  }
-
-  fd = sock;
-}
-
-TcpListener::~TcpListener() {
-  closesocket(fd);
+  listen(sock);
 }
 
-void TcpListener::shutdown()
-{
-#ifdef WIN32
-  closesocket(getFd());
-#else
-  ::shutdown(getFd(), 2);
-#endif
-}
-
-
-Socket*
-TcpListener::accept() {
-  int new_sock = -1;
-
-  // Accept an incoming connection
-  if ((new_sock = ::accept(fd, 0, 0)) < 0)
-    throw SocketException("unable to accept new connection", errorNumber);
-
-#ifndef WIN32
-  // - By default, close the socket on exec()
-  fcntl(new_sock, F_SETFD, FD_CLOEXEC);
-#endif
-
-  // Disable Nagle's algorithm, to reduce latency
-  TcpSocket::enableNagles(new_sock, false);
-
-  // Create the socket object & check connection is allowed
-  TcpSocket* s = new TcpSocket(new_sock);
-  if (filter && !filter->verifyConnection(s)) {
-    delete s;
-    return 0;
-  }
-  return s;
+Socket* TcpListener::createSocket(int fd) {
+  return new TcpSocket(fd);
 }
 
 void TcpListener::getMyAddresses(std::list<char*>* result) {
@@ -489,7 +397,7 @@ void TcpListener::getMyAddresses(std::list<char*>* result) {
 }
 
 int TcpListener::getMyPort() {
-  return TcpSocket::getSockPort(getFd());
+  return getSockPort(getFd());
 }
 
 
diff --git a/common/network/TcpSocket.h b/common/network/TcpSocket.h
index 87d8e695..eb6c0958 100644
--- a/common/network/TcpSocket.h
+++ b/common/network/TcpSocket.h
@@ -48,35 +48,33 @@ namespace network {
   /* Tunnelling support. */
   int findFreeTcpPort (void);
 
+  int getSockPort(int sock);
+
   class TcpSocket : public Socket {
   public:
     TcpSocket(int sock);
     TcpSocket(const char *name, int port);
-    virtual ~TcpSocket();
 
     virtual char* getPeerAddress();
     virtual char* getPeerEndpoint();
 
-    virtual void shutdown();
     virtual bool cork(bool enable);
 
-    static bool enableNagles(int sock, bool enable);
-    static bool isListening(int sock);
-    static int getSockPort(int sock);
+  protected:
+    bool enableNagles(bool enable);
   };
 
   class TcpListener : public SocketListener {
   public:
     TcpListener(const struct sockaddr *listenaddr, socklen_t listenaddrlen);
     TcpListener(int sock);
-    virtual ~TcpListener();
-
-    virtual void shutdown();
-    virtual Socket* accept();
 
     virtual int getMyPort();
 
     static void getMyAddresses(std::list<char*>* result);
+
+  protected:
+    virtual Socket* createSocket(int fd);
   };
 
   void createLocalTcpListeners(std::list<SocketListener*> *listeners,
diff --git a/common/network/UnixSocket.cxx b/common/network/UnixSocket.cxx
index 8a223e40..bfabc141 100644
--- a/common/network/UnixSocket.cxx
+++ b/common/network/UnixSocket.cxx
@@ -26,8 +26,6 @@
 #include <sys/un.h>
 #include <unistd.h>
 #include <errno.h>
-#include <signal.h>
-#include <fcntl.h>
 #include <stdlib.h>
 #include <stddef.h>
 
@@ -39,20 +37,9 @@ using namespace rdr;
 
 static rfb::LogWriter vlog("UnixSocket");
 
-// -=- Socket initialisation
-static bool socketsInitialised = false;
-static void initSockets() {
-  if (socketsInitialised)
-    return;
-  signal(SIGPIPE, SIG_IGN);
-  socketsInitialised = true;
-}
-
-
 // -=- UnixSocket
 
-UnixSocket::UnixSocket(int sock)
-  : Socket(new FdInStream(sock), new FdOutStream(sock))
+UnixSocket::UnixSocket(int sock) : Socket(sock)
 {
 }
 
@@ -66,7 +53,6 @@ UnixSocket::UnixSocket(const char *path)
     throw SocketException("socket path is too long", ENAMETOOLONG);
 
   // - Create a socket
-  initSockets();
   sock = socket(AF_UNIX, SOCK_STREAM, 0);
   if (sock == -1)
     throw SocketException("unable to create socket", errno);
@@ -85,16 +71,7 @@ UnixSocket::UnixSocket(const char *path)
   if (result == -1)
     throw SocketException("unable connect to socket", err);
 
-  // - By default, close the socket on exec()
-  fcntl(sock, F_SETFD, FD_CLOEXEC);
-
-  // Create the input and output streams
-  instream = new FdInStream(sock);
-  outstream = new FdOutStream(sock);
-}
-
-UnixSocket::~UnixSocket() {
-  close(getFd());
+  setFd(sock);
 }
 
 char* UnixSocket::getPeerAddress() {
@@ -132,12 +109,6 @@ char* UnixSocket::getPeerEndpoint() {
   return getPeerAddress();
 }
 
-void UnixSocket::shutdown()
-{
-  Socket::shutdown();
-  ::shutdown(getFd(), 2);
-}
-
 bool UnixSocket::cork(bool enable)
 {
   return true;
@@ -153,13 +124,9 @@ UnixListener::UnixListener(const char *path, int mode)
     throw SocketException("socket path is too long", ENAMETOOLONG);
 
   // - Create a socket
-  initSockets();
   if ((fd = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
     throw SocketException("unable to create listening socket", errno);
 
-  // - By default, close the socket on exec()
-  fcntl(fd, F_SETFD, FD_CLOEXEC);
-
   // - Delete existing socket (ignore result)
   unlink(path);
 
@@ -183,12 +150,7 @@ UnixListener::UnixListener(const char *path, int mode)
     throw SocketException("unable to set socket mode", err);
   }
 
-  // - Set it to be a listening socket
-  if (listen(fd, 5) < 0) {
-    err = errno;
-    close(fd);
-    throw SocketException("unable to set socket to listening mode", err);
-  }
+  listen(fd);
 }
 
 UnixListener::~UnixListener()
@@ -196,31 +158,12 @@ UnixListener::~UnixListener()
   struct sockaddr_un addr;
   socklen_t salen = sizeof(addr);
 
-  close(fd);
-
   if (getsockname(getFd(), (struct sockaddr *)&addr, &salen) == 0)
     unlink(addr.sun_path);
 }
 
-void UnixListener::shutdown()
-{
-  ::shutdown(getFd(), 2);
-}
-
-
-Socket*
-UnixListener::accept() {
-  int new_sock = -1;
-
-  // Accept an incoming connection
-  if ((new_sock = ::accept(fd, 0, 0)) < 0)
-    throw SocketException("unable to accept new connection", errno);
-
-  // - By default, close the socket on exec()
-  fcntl(new_sock, F_SETFD, FD_CLOEXEC);
-
-  // - Create the socket object
-  return new UnixSocket(new_sock);
+Socket* UnixListener::createSocket(int fd) {
+  return new UnixSocket(fd);
 }
 
 int UnixListener::getMyPort() {
diff --git a/common/network/UnixSocket.h b/common/network/UnixSocket.h
index 48c5b63f..1ffca456 100644
--- a/common/network/UnixSocket.h
+++ b/common/network/UnixSocket.h
@@ -37,12 +37,10 @@ namespace network {
   public:
     UnixSocket(int sock);
     UnixSocket(const char *name);
-    virtual ~UnixSocket();
 
     virtual char* getPeerAddress();
     virtual char* getPeerEndpoint();
 
-    virtual void shutdown();
     virtual bool cork(bool enable);
   };
 
@@ -51,10 +49,10 @@ namespace network {
     UnixListener(const char *listenaddr, int mode);
     virtual ~UnixListener();
 
-    virtual void shutdown();
-    virtual Socket* accept();
-
     int getMyPort();
+
+  protected:
+    virtual Socket* createSocket(int fd);
   };
 
 }