updated

1 files changed, 99 insertions, 40 deletions

diff --git a/server.erl b/server.erl
index dbf4182..7fcd6ab 100644
--- a/server.erl
+++ b/server.erl
@@ -1,82 +1,141 @@
 -module(server).
--export([server/2, manager/1, broadcast/2, broadcast/3, worker/2, create_workers/2]).
+-export([server/0, server/2, worker/2, create_workers/2, acceptor/2, listener/1, manager/1, broadcast/2, broadcast/3, handle_data/3]).
 
+server() ->
+    server(1235, 10).
 server(Port, NumWorkers) ->
-    process_flag(trap_exit, true),
+    ets:new(client_table, [named_table, public, {read_concurrency, true}]),
+    ets:insert(client_table, {counter, 0}),
+    ets:insert(client_table, {server, "Server"}),
+
     {ok, ServerSock} = gen_tcp:listen(Port, [list]),
     spawn(?MODULE, create_workers, [NumWorkers, self()]),
     supervisor(ServerSock, []).
 
 create_workers(0, _) -> ok;
 create_workers(Num, Supervisor) ->
-    Supervisor ! create,
+    Supervisor ! worker_new,
     create_workers(Num - 1, Supervisor).
 
 supervisor(ServerSock, Pids) ->
     receive
-        create ->
-            io:format("Creating new worker\n"),
-            NewPid = spawn_link(?MODULE, worker, [ServerSock, Pids]),
-            ping_all(Pids, NewPid),
+        worker_new ->
+            {NewPid, _} = spawn_monitor(?MODULE, worker, [ServerSock, Pids]),
+            io:format("Creating new worker ~p\n", [NewPid]),
+
+            message_all(Pids, {ping, NewPid}),
             supervisor(ServerSock, [NewPid | Pids]);
-        {'EXIT', From, Reason} ->
-            io:format("~p exited due to ~p\n", [From, Reason]),
+
+        {'DOWN', _, process, From, Reason} ->
+            io:format("Worker ~p exited due to ~p\n", [From, Reason]),
             self() ! create,
-            supervisor(ServerSock, lists:delete(From, Pids))
+
+            message_all(lists:delete(From, Pids), {close, From}),
+            supervisor(ServerSock, lists:delete(From, Pids));
+        {announce, Data} ->
+            [Pid|_] = Pids,
+            Pid ! {broadcast, "Server", ok, Data},
+            supervisor(ServerSock, Pids)
     end.
 
-ping_all([Pid|Pids], NewPid) ->
-    Pid ! { ping, NewPid },
-    ping_all(Pids, NewPid);
-ping_all([], _) -> ok.
+message_all([Pid|Pids], Data) ->
+    Pid ! Data,
+    message_all(Pids, Data);
+message_all([], _) -> ok.
 
 worker(ServerSock, Pids) ->
     process_flag(trap_exit, false),
-    ManagerPid = self(),
-    spawn_link(
-        fun() ->
-           process_flag(trap_exit, false),
-           listener(ManagerPid, ServerSock)
-        end),
+    ListenerPid = spawn_link(?MODULE, listener, [self()]),
+    spawn_link(?MODULE, acceptor, [ListenerPid, ServerSock]),
     manager(Pids).
 
-listener(ManagerPid, ServerSock) ->
-    case gen_tcp:accept(ServerSock) of
-        {ok, Sock} ->
-            ok = gen_tcp:controlling_process(Sock, ManagerPid),
-            ManagerPid ! {ping, Sock},
-            listener(ManagerPid, ServerSock);
-        Other ->
-            exit(Other)
-    end.
-
 manager(Socks) ->
     receive
-        {tcp, Sock, Data} ->
-            broadcast(lists:delete(Sock, Socks), Data),
-            manager(Socks);
-        {tcp_closed, Sock} -> 
-            manager(lists:delete(Sock, Socks));
         {ping, Sock} ->
             manager([Sock | Socks]);
+        {close, Sock} -> 
+            manager(lists:delete(Sock, Socks));
 
         {repeat, Data} ->
             broadcast(Socks, Data, norepeat),
             manager(Socks);
-        Other ->
-            exit(Other)
+        {broadcast, FromIden, FromSock, Data} ->
+            NewData = FromIden ++ ": " ++ Data,
+            broadcast(lists:delete(FromSock, Socks), NewData),
+            manager(Socks);
+        {pm, FromIden, ToSock, Data} ->
+            NewData = "(pm) " ++ FromIden ++ ": " ++ Data,
+            broadcast([ToSock], NewData),
+            manager(Socks)
+    end.
+
+listener(ManagerPid) ->
+    receive
+        {tcp, Sock, Data} ->
+            spawn(?MODULE, handle_data, [ManagerPid, Sock, Data]),
+            listener(ManagerPid);
+        {tcp_opened, Sock} ->
+            ManagerPid ! {ping, Sock},
+            client_add(Sock),
+            listener(ManagerPid);
+        {tcp_closed, Sock} -> 
+            ManagerPid ! {close, Sock},
+            client_del(Sock),
+            listener(ManagerPid)
     end.
 
+acceptor(ListenerPid, ServerSock) ->
+    case gen_tcp:accept(ServerSock) of
+        {ok, Sock} ->
+            ok = gen_tcp:controlling_process(Sock, ListenerPid),
+            ListenerPid ! {tcp_opened, Sock},
+            acceptor(ListenerPid, ServerSock)
+    end.
+
+
+handle_data(ManagerPid, From, Data) ->
+    case catch parse_data(ManagerPid, From, Data) of
+        ok -> ok;
+        _ -> 
+            ManagerPid ! {pm, "Server", From, "ERR: Transaction Failed\n"}
+    end.
+
+parse_data(ManagerPid, From, Data) ->
+    case string:tokens(Data, " ") of
+        ["/pm" | [ToIden | _]] ->
+            NewData = string:slice(Data, string:length(ToIden) + 5),
+            ManagerPid ! {pm, client_iden(From), client_sock(ToIden), NewData};
+        ["\n"] -> ok;
+        _ -> 
+            ManagerPid ! {broadcast, client_iden(From), From, Data}
+    end, ok.
+    
+
+broadcast(Socks, Data) ->
+    broadcast(Socks, Data, ok).
 broadcast([Pid|Socks], Data, norepeat) when is_pid(Pid) ->
     broadcast(Socks, Data, norepeat);
 broadcast([Pid|Socks], Data, Opt) when is_pid(Pid) ->
     Pid ! {repeat, Data},
     broadcast(Socks, Data, Opt);
-
 broadcast([Sock|Socks], Data, Opt)  ->
     gen_tcp:send(Sock, Data),
     broadcast(Socks, Data, Opt);
 broadcast([], _, _) -> ok.
 
-broadcast(Socks, Data) ->
-    broadcast(Socks, Data, ok).
+client_add(Sock) ->
+    [{_, Count}|_] = ets:lookup(client_table, counter),
+    Iden = "Anon"++integer_to_list(Count),
+    ets:insert(client_table, {Iden, Sock}),
+    ets:insert(client_table, {Sock, Iden}),
+    ets:update_counter(client_table, counter, 1).
+client_del(Sock) ->
+    [{_, Iden}|_] = ets:lookup(client_table, Sock),
+    ets:delete(client_table, Sock),
+    ets:delete(client_table, Iden).
+client_iden(Sock) ->
+    [{_, Iden}|_] = ets:lookup(client_table, Sock),
+    Iden.
+client_sock(Iden) ->
+    [{_, Sock}|_] = ets:lookup(client_table, Iden),
+    Sock.