Haskell Lol Server

gistfile1.hs

import System.IO 
import Network.Socket
import Network (PortID(..))
import System.Environment
import Control.Concurrent
import Control.Monad
import Control.Monad.Fix (fix)
import Data.Char
import Data.List.Split
import System.IO.Error
import qualified Data.Map as Map
import Data.List (elemIndex,intercalate)
import Control.Exception.Base (evaluate)


policyFileL = "<?xml version='1.0'?><cross-domain-policy><allow-access-from domain='*' to-ports='"
policyFileR = "' /></cross-domain-policy>\0"

{--
 - Represents an instruction request for relay to carry out
 -
 - Commands:
 - push 		-> pushes {contents} to the user with id {args}
 - connect 		-> maps a socket connection to a user with id {args}
 - listConnection	-> lists all users currently connected
 - writePolicyFile	-> writes a cross-domain policy file to allow the Flash socket to connect
 -}
data Instruction = Instruction { command :: String, args :: String, contents :: String } deriving (Show)


main = do
	args <- getArgs -- accepts the port # as its only argument
	runServer $ PortNumber $ fromIntegral (read $ head args :: Int)


runServer :: PortID -> IO ()
runServer (PortNumber port) = do

	sock <- socket AF_INET Stream defaultProtocol
	setSocketOption sock ReuseAddr 1
	putStrLn ("Binding on port: " ++ show port)
	bindSocket sock (SockAddrInet port iNADDR_ANY)
	listen sock 5 
	mutexMap <- newMVar Map.empty -- maps user ID's to their associatesd socket handles
	acceptConnections sock mutexMap (show port)




acceptConnections :: Socket -> MVar (Map.Map String Handle) -> String -> IO ()
acceptConnections sock mutexMap port = forever $ do

	(connection, address) <- accept sock
	putStrLn ("Client connected at: " ++ show address)

	-- manage each connection in a separate thread
	forkIO $ handleConnection connection mutexMap port




handleConnection :: Socket -> MVar (Map.Map String Handle) -> String -> IO ()
handleConnection connection mutexMap port = do

	connectionHandle <- socketToHandle connection ReadWriteMode
	hSetBuffering connectionHandle LineBuffering

	-- Flash's socket library passes a request for a crossdomain policy file
	-- terminatd with a null char; All the handle methods for extracting from 
	-- the buffer in whole are blocking on a newline, so we have to peek and 
	-- see if there's content along with the incoming connection
	firstChar <- try (hLookAhead connectionHandle)
	case firstChar of 
		Right '<' ->
			-- print out the policy file and exit
			hPutStrLn connectionHandle (policyFileL ++ port ++ policyFileR)
		
		otherwise -> do
			-- runs until user disconnects
			handleMessage connectionHandle mutexMap port
			
			-- if we're this far, we need to try to remove the connection
			-- from the map
			connMap <- takeMVar mutexMap
			--tester <- (evaluate (show connMap))
			--putStrLn tester
			putMVar mutexMap $! removeByValue connectionHandle connMap 

	putStrLn "Closing connection"
	hClose connectionHandle




handleMessage :: Handle -> MVar (Map.Map String Handle) -> String -> IO ()
handleMessage connectionHandle mutexMap port = do

	request <- try (hGetLine connectionHandle)
	case request of
		Left e ->
			if isEOFError e
				then return () -- this is how we disconnect the user
				else ioError e
		Right serialInstruction -> do
			let maybeInstruction = getInstruction serialInstruction port
			case maybeInstruction of
				Nothing ->
					return ()
				Just instruction -> do
					handleInstruction instruction connectionHandle mutexMap
			handleMessage connectionHandle mutexMap port




handleInstruction :: Instruction -> Handle -> MVar (Map.Map String Handle) -> IO ()
handleInstruction (Instruction command args contents) connectionHandle mutexMap = do

	case command of
		"connect" -> do
			putStrLn ("Connecting " ++ args)
			connMap <- takeMVar mutexMap
			putMVar mutexMap (Map.insert args connectionHandle connMap)
		"push" -> do
			connMap <- takeMVar mutexMap
			let target = Map.lookup args connMap
			case target of
				Nothing ->
					return()
				Just connection -> do
					putStrLn ("Pushing: " ++ contents)
					push <- try (hPutStrLn connection contents)
					case push of
						Left _ -> putStrLn "Error pushing"
						Right _ -> putStrLn "Successfully pushed"
			putMVar mutexMap connMap
		"listConnected" -> do
			connMap <- takeMVar mutexMap
			mapM_ (\(k,v) -> hPutStrLn connectionHandle k) (Map.toList connMap)
			putMVar mutexMap connMap
		"writePolicyFile" -> hPutStrLn connectionHandle args
		otherwise ->
			return ()





getInstruction :: String -> String -> Maybe Instruction
getInstruction "" _ = Nothing
getInstruction "<policy-file-request/>" port = Just (Instruction "writePolicyFile" (policyFileL ++ port ++ policyFileR) "")
getInstruction stream _
	| (length instr) > 2 = Just ((\(x:xs:xss) -> Instruction x xs (intercalate ":" xss)) instr)
	| otherwise = Nothing
	where
	instr = splitOn ":" stream


removeByValue :: (Ord a, Eq b) => b -> (Map.Map a b) -> (Map.Map a b)
removeByValue value kvpairs = do
	let ei = elemIndex value (Map.elems kvpairs)
	case ei of
		Nothing -> kvpairs
		Just keyIndex -> Map.delete ((Map.keys kvpairs) !! keyIndex) kvpairs