Add recv and send, using WSARecv and WSASend

Author: joeyadams

Winsock.hsc

@@ -6,6 +6,11 @@ module Winsock (
     socket,
     connect,
     close,
+    recvBuf,
+    sendBuf,
+
+    recv,
+    send,
 ) where
 
 #include <windows.h>
@@ -26,6 +31,9 @@ import qualified IOCP.Manager as Manager
 
 import Control.Applicative      ((<$>))
 import Control.Monad            (void)
+import Data.ByteString          (ByteString)
+import Data.ByteString.Internal (createAndTrim)
+import Data.ByteString.Unsafe   (unsafeUseAsCStringLen)
 import Data.IORef
 import Data.Word
 import Foreign.C
@@ -40,6 +48,8 @@ import qualified System.Win32.Types as Win32
 newtype Socket = Socket IOCPHandle
     deriving Eq
 
+-- Note: Functions that take a 'Socket' expect WinSock to already be initialized.
+
 socket :: NS.Family -> NS.SocketType -> NS.ProtocolNumber -> IO Socket
 socket family stype protocol = do
     initWinsock
@@ -69,6 +79,40 @@ close (Socket ih) =
     Manager.closeWith ih $
     Win32.failIf_ (/= 0) "close" . c_closesocket . castHANDLEToSOCKET
 
+recvBuf :: Socket -> Ptr a -> Int -> IO Int
+recvBuf (Socket ih) buf len =
+    withIOCP ih 0 startCB completionCB
+  where
+    startCB h ol =
+        Win32.failIfFalse_ "recv" $
+        c_winsock_recv (castHANDLEToSOCKET h) (castPtr buf) (fromIntegral len) ol
+
+    completionCB err numBytes
+        | err == 0  = return (fromIntegral numBytes)
+        | otherwise = FFI.throwWinErr "recv" err
+
+sendBuf :: Socket -> Ptr a -> Int -> IO Int
+sendBuf (Socket ih) buf len =
+    withIOCP ih 0 startCB completionCB
+  where
+    startCB h ol =
+        Win32.failIfFalse_ "send" $
+        c_winsock_send (castHANDLEToSOCKET h) (castPtr buf) (fromIntegral len) ol
+
+    completionCB err numBytes
+        | err == 0  = return (fromIntegral numBytes)
+        | otherwise = FFI.throwWinErr "send" err
+
+recv :: Socket -> Int -> IO ByteString
+recv sock len =
+    createAndTrim len $ \buf ->
+    recvBuf sock buf len
+
+send :: Socket -> ByteString -> IO Int
+send sock bs =
+    unsafeUseAsCStringLen bs $ \(buf, len) ->
+    sendBuf sock buf len
+
 newtype Winsock = Winsock (Ptr ())
 
 getWinsock :: IO Winsock
@@ -97,3 +141,9 @@ foreign import ccall unsafe
 
 foreign import WINDOWS_CCONV safe "winsock2.h closesocket"
     c_closesocket :: SOCKET -> IO CInt
+
+foreign import ccall unsafe
+    c_winsock_recv :: SOCKET -> Ptr CChar -> #{type u_long} -> LPOVERLAPPED -> IO BOOL
+
+foreign import ccall unsafe
+    c_winsock_send :: SOCKET -> Ptr CChar -> #{type u_long} -> LPOVERLAPPED -> IO BOOL

cbits/Winsock.c

@@ -104,18 +104,6 @@ Winsock *c_winsock_init(void)
     return NULL;
 }
 
-static BOOL bind_inet(SOCKET sock, ULONG address, USHORT port)
-{
-    struct sockaddr_in addr;
-    int rc;
-    memset(&addr, 0, sizeof(addr));
-    addr.sin_family      = AF_INET;
-    addr.sin_addr.s_addr = address;
-    addr.sin_port        = port;
-    rc = bind(sock, (SOCKADDR *) &addr, sizeof(addr));
-    return (rc == 0);
-}
-
 BOOL c_winsock_connect(Winsock *winsock, SOCKET sock, SOCKADDR *addr, int addrLen, OVERLAPPED *ol)
 {
     // ConnectEx requires the socket to be initially bound.
@@ -134,3 +122,18 @@ BOOL c_winsock_connect(Winsock *winsock, SOCKET sock, SOCKADDR *addr, int addrLe
     BOOL ok = winsock->ConnectEx(sock, addr, addrLen, NULL, 0, NULL, ol);
     return (ok || WSAGetLastError() == ERROR_IO_PENDING);
 }
+
+BOOL c_winsock_recv(SOCKET sock, char *buf, u_long bufsize, OVERLAPPED *ol)
+{
+    WSABUF wsabuf = {.len = bufsize, .buf = buf};
+    DWORD flags = 0;
+    int rc = WSARecv(sock, &wsabuf, 1, NULL, &flags, ol, NULL);
+    return (rc == 0 || WSAGetLastError() == ERROR_IO_PENDING);
+}
+
+BOOL c_winsock_send(SOCKET sock, char *buf, u_long bufsize, OVERLAPPED *ol)
+{
+    WSABUF wsabuf = {.len = bufsize, .buf = buf};
+    int rc = WSASend(sock, &wsabuf, 1, NULL, 0, ol, NULL);
+    return (rc == 0 || WSAGetLastError() == ERROR_IO_PENDING);
+}

iocp.cabal

@@ -32,7 +32,9 @@ library
     build-tools: hsc2hs
 
     ghc-options: -Wall -fwarn-tabs -fno-warn-missing-signatures
+    cc-options: -Wall
 
     build-depends   : base == 4.*
+                    , bytestring
                     , Win32
                     , network