1 files changed, 57 insertions, 10 deletions

diff --git a/src/SocketCore.h b/src/SocketCore.h
index 60f33279..b2dda052 100644
--- a/src/SocketCore.h
+++ b/src/SocketCore.h
@@ -50,6 +50,39 @@
 
 namespace aria2 {
 
+#ifndef __MINGW32__
+#  define SOCKET_ERRNO (errno)
+#else
+#  define SOCKET_ERRNO (WSAGetLastError())
+#endif // __MINGW32__
+
+#ifdef __MINGW32__
+#  define A2_EINPROGRESS WSAEWOULDBLOCK
+#  define A2_EWOULDBLOCK WSAEWOULDBLOCK
+#  define A2_EINTR WSAEINTR
+#  define A2_WOULDBLOCK(e) (e == WSAEWOULDBLOCK)
+#else // !__MINGW32__
+#  define A2_EINPROGRESS EINPROGRESS
+#  ifndef EWOULDBLOCK
+#    define EWOULDBLOCK EAGAIN
+#  endif // EWOULDBLOCK
+#  define A2_EWOULDBLOCK EWOULDBLOCK
+#  define A2_EINTR EINTR
+#  if EWOULDBLOCK == EAGAIN
+#    define A2_WOULDBLOCK(e) (e == EWOULDBLOCK)
+#  else // EWOULDBLOCK != EAGAIN
+#    define A2_WOULDBLOCK(e) (e == EWOULDBLOCK || e == EAGAIN)
+#  endif // EWOULDBLOCK != EAGAIN
+#endif   // !__MINGW32__
+
+#ifdef __MINGW32__
+#  define CLOSE(X) ::closesocket(X)
+#else
+#  define CLOSE(X) close(X)
+#endif // __MINGW32__
+
+std::string errorMsg(int errNum);
+
 #ifdef ENABLE_SSL
 class TLSContext;
 class TLSSession;
@@ -64,7 +97,7 @@ class SocketCore {
   friend bool operator!=(const SocketCore& s1, const SocketCore& s2);
   friend bool operator<(const SocketCore& s1, const SocketCore& s2);
 
-private:
+protected:
   // socket type defined in <sys/socket.h>
   int sockType_;
   // socket endpoint descriptor
@@ -121,7 +154,7 @@ public:
   // std::make_shared.
   SocketCore(sock_t sockfd, int sockType);
 
-  ~SocketCore();
+  virtual ~SocketCore();
 
   sock_t getSockfd() const { return sockfd_; }
 
@@ -193,6 +226,14 @@ public:
   Endpoint getPeerInfo() const;
 
   /**
+   * Stores peer's address to sockaddr.  len must be
+   * initialized to the size of sockaddr.  On success, address data is
+   * stored in sockaddr and actual size of address structure is stored
+   * in len.
+   */
+  void getPeerInfo(sockaddr_union& sockaddr, socklen_t& len) const;
+
+  /**
    * Accepts incoming connection on this socket.
    * You must call beginListen() before calling this method.
    * @return accepted socket.
@@ -208,8 +249,8 @@ public:
    * @param port service port number to connect to
    * @param tcpNodelay true to disable Nagle algorithm
    */
-  void establishConnection(const std::string& host, uint16_t port,
-                           bool tcpNodelay = true);
+  virtual void establishConnection(const std::string& host, uint16_t port,
+                                   bool tcpNodelay = true);
 
   void setNonBlockingMode();
 
@@ -221,7 +262,7 @@ public:
   /**
    * Closes the connection of this socket.
    */
-  void closeConnection();
+  virtual void closeConnection();
 
   /**
    * Checks whether this socket is available for writing.
@@ -230,7 +271,7 @@ public:
    * @return true if the socket is available for writing,
    * otherwise returns false.
    */
-  bool isWritable(time_t timeout);
+  virtual bool isWritable(time_t timeout);
 
   /**
    * Checks whether this socket is available for reading.
@@ -239,7 +280,7 @@ public:
    * @return true if the socket is available for reading,
    * otherwise returns false.
    */
-  bool isReadable(time_t timeout);
+  virtual bool isReadable(time_t timeout);
 
   /**
    * Writes data into this socket. data is a pointer pointing the first
@@ -253,7 +294,7 @@ public:
    * @param data data to write
    * @param len length of data
    */
-  ssize_t writeData(const void* data, size_t len);
+  virtual ssize_t writeData(const void* data, size_t len);
   ssize_t writeData(const std::string& msg)
   {
     return writeData(msg.c_str(), msg.size());
@@ -262,7 +303,10 @@ public:
   ssize_t writeData(const void* data, size_t len, const std::string& host,
                     uint16_t port);
 
-  ssize_t writeVector(a2iovec* iov, size_t iovcnt);
+  ssize_t writeData(const void* data, size_t len, const sockaddr* addr,
+                    socklen_t socklen, int flags);
+
+  virtual ssize_t writeVector(a2iovec* iov, size_t iovcnt);
 
   /**
    * Reads up to len bytes from this socket.
@@ -280,11 +324,14 @@ public:
    * @param len the maximum size data can store. This method assigns
    * the number of bytes read to len.
    */
-  void readData(void* data, size_t& len);
+  virtual void readData(void* data, size_t& len);
 
   // sender.addr will be numerihost assigned.
   ssize_t readDataFrom(void* data, size_t len, Endpoint& sender);
 
+  ssize_t readDataFrom(void* data, size_t len, sockaddr_union& sockaddr,
+                       socklen_t& sockaddrlen);
+
 #ifdef ENABLE_SSL
   // Performs TLS server side handshake. If handshake is completed,
   // returns true. If handshake has not been done yet, returns false.