Erang 聊天服务器（binary）

最近在学Erlang,写了个聊天服务器，后来又把消息改成用二进制传输，记录下源码

server.erl

%% Author: xieyijia
%% Created: 2011-11-7
%% Description: TODO: Add description to bin_chat_server
-module(bin_chat_server).
-import(lists,[foreach/2]).
-include_lib("stdlib/include/qlc.hrl").


%%
%% Exported Functions
%%
-export([start/0]).
-record(user, {
        id       %% 用户ID
        ,name   %% 用户名称
        ,passwd %% 用户登录密码
        ,login_times %% 登录次数
        ,chat_times   %% 聊天次数
        ,last_login   %% 最后一次登录时间
    }
).
-record(online_user, {id, name}). %%在线用户


-define(HEAD_SIZE, 4).
-define(UINT, 32/unsigned-little-integer).
-define(INT, 32/signed-little-integer).
-define(USHORT, 16/unsigned-little-integer).
-define(SHORT, 16/signed-little-integer).
-define(UTYTE, 8/unsigned-little-integer).
-define(BYTE, 8/signed-little-integer).


%% 协议消息头
-define(ADD_SOCKET, 100).
-define(LOGIN, 101).
-define(CHAT, 102).
-define(LOGOUT, 103).
-define(GET_NUM, 104).


%%
%% TODO: Add description of server_start/function_arity
%%　启动服务器
start() -> 
create_db_one_time(),
start_db(),
  init_tables(),
register(socket_store, spawn(fun() -> socket_store([]) end)),
{ok, Listen} = gen_tcp:listen(1234, [binary, {packet, 4}, {reuseaddr, true}, {active, true}]),
spawn(fun() -> par_connect(Listen) end),
io:format(" Sever start complete...~n").




%%
%% tcp server
%%
par_connect(Listen) ->
{ok, Socket} = gen_tcp:accept(Listen),
io:format(" Server accept Socket:~p ~n", [Socket]),
%% gen_tcp:close(Socket),
spawn(fun() -> par_connect(Listen) end),
loop_head(Socket).

loop_head(Socket) ->
receive
{tcp, Socket, Bin} ->
%% io:format("recv head binary = ~p~n", [Bin]),
<<Len:16,Other/binary>> = Bin,
io:format("Len = ~p, other = ~p~n", [Len, Other]),
case Len of 
?ADD_SOCKET ->
add_socket(Socket),
{B, L} = string_to_binary("ok"),
%% io:format("bin = ~p, len = ~p~n",[B, L]),
Reply = <<?ADD_SOCKET:16, L:16, B/binary>>,
io:format("Reply = ~p~n",[Reply]);
?LOGIN ->
<<ID:16, Length:16, Nick:Length/binary, P_len:16, Passwd:P_len/binary>> = Other,
%% io:format("id = ~p, nick = ~p, passwd = ~p ~n", [ID, Nick, Passwd]),
_Nick = list_to_atom(binary_to_list(Nick)),
_Passwd = list_to_integer(binary_to_list(Passwd)),
%% io:format("id = ~p, nick = ~p, passwd = ~p ~n", [ID, _Nick, _Passwd]),
Check = check_user_info(ID, _Passwd, _Nick),
io:format("check_user_info:~p~n", [Check]),
if
Check =:=[] -> 
{R, RL} = string_to_binary("failed"),
Reply = <<?LOGIN:16, RL:16, R:RL/binary>>;
true -> 
update_user_info(login, ID),
add_to_online_user(ID, _Nick),
{L, LL} = string_to_binary("ok"),
Reply = <<?LOGIN:16, LL:16, L:LL/binary>>,
io:format("user ~p login OK~n",[_Nick])
end;
?CHAT -> 
<<ID:16, Length:16, Nick:Length/binary, M_len:16, Msg:M_len/binary>> = Other,
update_user_info(chat, ID),
Pid = self(),
Slist = get_online_socket(Pid),
%% Nick = list_to_atom(binary_to_list(_Nick)),
send_msg_to_online_user(Slist,Nick,Msg),
{C, CL} = string_to_binary("ok"), 
Reply = <<?CHAT:16, CL:16, C:CL/binary>>,
io:format("~p say ~p OK~n", [Nick, Msg]);
?LOGOUT ->
<<ID:16>> = Other,
delete_user(ID),
{Logout, L} = string_to_binary("ok"),
Reply = <<?LOGOUT:16, L:16, Logout:L/binary>>,
io:format("id ~p logout OK~n", [ID]);
?GET_NUM ->
Olist = select_all(online_user),
Num = erlang:length(Olist),
io:format("Olist = ~p~n length = ~p~n", [Olist, Num]),
Reply = <<?GET_NUM:16, Num:16>>,
%% Reply = {get_online_num, Num},
io:format(" number is ~p OK~n", [Num])
end,
gen_tcp:send(Socket, [Reply]),
loop_head(Socket);
{error, closed} ->
io:format("recv error...~n")
end.


string_to_binary(Str) ->
%% io:format("come here...~p~n",[Str]),
Bin = list_to_binary(Str),
Len = byte_size(Bin),
%% io:format("bin = ~p, len = ~p~n", [Bin, Len]),
{Bin, Len}.


%% 
%% 创建数据医库
%%
create_db_one_time() ->
mnesia:create_schema([node()]),
mnesia:start(),
mnesia:create_table(user, [{attributes, record_info(fields, user)}]),
mnesia:create_table(online_user, [{attributes, record_info(fields, online_user)}]),
mnesia:stop().


start_db() ->
mnesia:start(),
mnesia:wait_for_tables([user, online_user], 2000).


example_users() ->
[{user, 1, aaa, 111111, 0, 0, {0,0,0,0,0,0}},
     {user, 2, bbb, 222222, 0, 0, {0,0,0,0,0,0}},
     {user, 3, ccc, 333333, 0, 0, {0,0,0,0,0,0}},
     {user, 4, ddd, 444444, 0, 0, {0,0,0,0,0,0}},
     {user, 5, eee, 555555, 0, 0, {0,0,0,0,0,0}}
].


init_tables() ->
mnesia:clear_table(user),
F = fun() ->
foreach(fun mnesia:write/1, example_users())
end,
mnesia:transaction(F).

do(Q) ->
F = fun() -> qlc:e(Q) end,
{atomic, Val} = mnesia:transaction(F),
Val.




%% 发送保存socket
add_socket(Socket) ->
socket_store ! {addsocket, Socket}.


%% 得到所有登录的socket
get_online_socket(Pid) ->
socket_store ! {get_socket, Pid},
receive
{socket_list, Socket_list} ->
Socket_list;
_ ->
io:format("get_online_socket error....~n")
end.


%% 保存所有登录的socket
socket_store(Socket_list) ->
receive
{addsocket, Socket} ->
New_socket_list = [{socket, Socket} | Socket_list],
socket_store(New_socket_list);
{get_socket, From} ->
From ! {socket_list, Socket_list},
socket_store(Socket_list)
end.


%% 
%% 逻辑处理
%%
check_user_info(ID, Passwd, Nick) ->
do(qlc:q([X || X <- mnesia:table(user),
 X#user.id =:= ID,
 X#user.passwd =:= Passwd,
 X#user.name =:= Nick])).
 


%%更新登录数，聊天次数，最后登录时间
update_user_info(login, ID) -> 
Old_user = lists:last(do(qlc:q([X || X <- mnesia:table(user), X#user.id =:= ID]))),
io:format("Old_user is: ~p ~n", [Old_user]),
#user{login_times=Ltimes} = Old_user,
io:format("login_times = ~p~n", [Ltimes]),
New_user = Old_user#user{login_times = Ltimes + 1, 
last_login = list_to_tuple(tuple_to_list(date())++tuple_to_list(time()))},
io:format("New_user is: ~p ~n", [New_user]),
F = fun() ->
mnesia:write(New_user)
end,
mnesia:transaction(F);
update_user_info(chat, ID) ->
Old_chat = lists:last(do(qlc:q([X || X <- mnesia:table(user), X#user.id =:= ID]))),
%% io:format("Old_chat is: ~p ~n", [Old_chat]),
#user{chat_times = Ctimes} = Old_chat,
%% io:format("chat_times = ~p~n", [Ctimes]),
New_chat = Old_chat#user{chat_times = Ctimes + 1},
F = fun() ->
mnesia:write(New_chat)
end,
mnesia:transaction(F).


%% 删除退出的用户
delete_user(ID) ->
Oid = {online_user, ID},
F = fun() ->
mnesia:delete(Oid)
end,
mnesia:transaction(F).


%%添加到online_user
add_to_online_user(ID, Nick) ->
Row = #online_user{id = ID, name = Nick},
io:format(" add to online_user:~p~n", [Row]),
F = fun() ->
mnesia:write(Row)
end,
mnesia:transaction(F).


%% 发送消息到其它登录的用户上
send_msg_to_online_user([], _, _) -> ok;
send_msg_to_online_user([H|T], Nick, Msg) ->
case H of
{socket, Socket} ->
Nick_len = byte_size(Nick),
Msg_len = byte_size(Msg),
Data = <<?CHAT:16, Nick_len:16, Nick:Nick_len/binary, Msg_len:16, Msg:Msg_len/binary>>,
gen_tcp:send(Socket, [Data]),
io:format("Server send chat data ~p~n", [Data]),
send_msg_to_online_user(T, Nick, Msg);
_ ->
io:format("send_msg_to_online_user erroe H = ~p~n", [H])
end.


%% 得到所有在线列表
select_all(Table) ->
do(qlc:q([X || X <- mnesia:table(Table)])).



client.erl

%% Author: xieyijia
%% Created: 2011-11-7
%% Description: TODO: 聊天客户端
-module(bin_chat_client).
-export([start/0]).
-compile(export_all).


%% 协议消息头
-define(ADD_SOCKET, 100).
-define(LOGIN, 101).
-define(CHAT, 102).
-define(LOGOUT, 103).
-define(GET_NUM, 104).


%%
%% TODO: Add description of start/function_arity
%%
start() -> 
start_send(),
spawn(fun() -> start_receive() end).


%%
%% Local Functions
%%
start_send() ->
{ok, Socket} = gen_tcp:connect("localhost", 1234, [binary, {packet, 4}]),
io:format("client connect Socket:~p...~n", [Socket]),
register(save_socket, spawn(fun() -> save_socket(Socket) end)),
io:format("client start send complete...~n").

start_receive() ->
{ok, Socket} = gen_tcp:connect("localhost", 1234, [binary, {packet, 4}]),
Add_socket = "add_socket",
{Bin, Len} = string_to_binary(Add_socket),
%% io:format("bin = ~p, len = ~p~n",[Bin, Len]),
Msg = <<?ADD_SOCKET:16, Len:16, Bin/binary>>,
gen_tcp:send(Socket, [Msg]),
receive_chat_msg(Socket),
io:format("client start receive complete...~n").


string_to_binary(Str) ->
Bin = list_to_binary(Str),
Len = byte_size(Bin),
{Bin, Len}.


save_socket(Socket) ->
receive
  {get_socket, Pid} ->
Pid ! {socket, Socket},
%% io:format("save_socket send socket:~p ok ~n", [Socket]),
save_socket(Socket);
_ ->
io:format("save_socket receive error msg ~n"),
save_socket(Socket)
end.


receive_msg(Socket) ->
receive
{tcp, Socket, Bin} ->
%% io:format("recv head binary = ~p~n", [Bin]),
<<TypeLen:16,Other/binary>> = Bin,
io:format("TypeLen = ~p, other = ~p~n", [TypeLen, Other]),
case TypeLen of
?LOGIN ->
<<Login_len:16, Login_result:Login_len/binary>> = Other,
io:format("client login, ~p~n", [Login_result]);
?LOGOUT ->
<<Logout_len:16, Logout_result:Logout_len/binary>> = Other,
io:format("client logout ~p~n", [Logout_result]);
?CHAT ->
<<Chat_len:16, Chat_result:Chat_len/binary>> = Other,
io:format("client chat ~p~n", [Chat_result]);
?GET_NUM ->
<<Num_len:16>> = Other,
io:format("client get online number is:~p ~n", [Num_len])
end;
_ ->
io:format("client receive error~n")
end.

receive_chat_msg(Socket) ->
receive
{tcp, Socket, Bin} ->
%% io:format("receive binary: ~p~n", [Bin]),
<<Length:16,Other/binary>> = Bin,
io:format("TypeLen = ~p, other = ~p~n", [Length, Other]),
case Length of
?CHAT ->
<<Name_len:16, _Name:Name_len/binary, Msg_len:16, _Msg:Msg_len/binary>> = Other,
io:format("receive ~p: say ~p~n", [_Name, _Msg]),
receive_chat_msg(Socket);
?ADD_SOCKET ->
<<Len:16, Result:Len/binary>> = Other,
io:format("server save the socket ~p~n", [Result]),
receive_chat_msg(Socket)
end
end.


login(ID, Passwd, Nick) ->
{_Nick, _NickLen} = string_to_binary(Nick),
{_Passwd, _PasswdLen} = string_to_binary(Passwd),
Msg = <<?LOGIN:16,ID:16, _NickLen:16, _Nick:_NickLen/binary, _PasswdLen:16, _Passwd:_PasswdLen/binary>>,
io:format("login send ~p~n", [Msg]),
go(Msg).


chat(ID, Nick, M) ->
{N, NL} = string_to_binary(Nick),
{_M, ML} = string_to_binary(M),
Msg = <<?CHAT:16, ID:16, NL:16, N:NL/binary, ML:16, _M:ML/binary>>,
go(Msg).


logout(ID) ->
Msg = <<?LOGOUT:16, ID:16>>,
go(Msg).


get_online_num() ->
Msg = <<?GET_NUM:16>>,
go(Msg).


go(Msg) ->
Pid = self(),
Socket = get_socket(Pid),
io:format("[go] get socket = :~p~n", [Socket]),
gen_tcp:send(Socket, [Msg]),
receive_msg(Socket).


get_socket(Pid) ->
save_socket ! {get_socket, Pid},
receive
{socket, Socket} ->
Socket;
_ ->
io:format("get socket error~n")
end.