-- $Id: echo.ghc,v 1.2 2001/05/01 20:19:52 doug Exp $
-- http://www.bagley.org/~doug/shootout/
-- Haskell echo/client server
-- written by Brian Gregor
-- compile with:
--      ghc -O -o echo -package net -package concurrent -package lang echo.hs

module Main where

import SocketPrim
import Concurrent
import System (getArgs,exitFailure)
import Exception(finally)
import MVar

server_sock = do
    s <- socket AF_INET Stream 6
    setSocketOption s ReuseAddr 1
    bindSocket s (SockAddrInet (mkPortNumber portnum) iNADDR_ANY)
    listen s 2
    return s

echo_server s = do
    (s', clientAddr) <- accept s
    proc <- read_data s' 0
    putStrLn ("server processed "++(show proc)++" bytes")
    sClose s'
    where
       read_data sock totalbytes = do
             (str,i) <- readSocket sock 19
         if (i >= 19) 
            then (do
              writ <- writeSocket sock str
              read_data sock $! (totalbytes+(length $! str)))
            else (return totalbytes)

local        = "127.0.0.1"        
message        = "Hello there sailor\n"
portnum     = 7001

client_sock = do
    s <- socket AF_INET Stream 6
    ia <- inet_addr local
    connect s (SockAddrInet (mkPortNumber portnum) ia)
    return s

echo_client n = do
    s <- client_sock
    drop <- server_echo s n
    sClose s
    where
      server_echo sock n 
          | n > 0 =(do 
                       writeSocket sock message
                   (str,i) <- readSocket sock 19
                       if (str /= message) then (do exitFailure)
                       else server_echo sock (n-1))
          | otherwise = (return [])

main = do 
     ~[n] <- getArgs
     -- server & client semaphores
         --get the server socket
         ssock <- server_sock 
     -- fork off the server
     s <- myForkIO (echo_server ssock)
     -- fork off the client
     c <- myForkIO (echo_client (read n::Int))
     -- let 'em run until they've signaled they're done
         join s ; join c

-- these are used to make the main thread wait until
-- the child threads have exited
myForkIO :: IO () -> IO (MVar ())
myForkIO io = do
  mvar <- newEmptyMVar
  forkIO (io `finally` putMVar mvar ())
  return mvar

join :: MVar () -> IO ()
join mvar = readMVar mvar