PageRenderTime 52ms CodeModel.GetById 23ms RepoModel.GetById 0ms app.codeStats 0ms

/apps/rtmp/src/rtmp_socket.erl

https://github.com/MicronXD/erlyvideo
Erlang | 702 lines | 387 code | 152 blank | 163 comment | 0 complexity | 0d566d4cdb8f0c1abd68f2cc6167ff7f MD5 | raw file
Possible License(s): GPL-3.0, BSD-3-Clause, MIT, MPL-2.0-no-copyleft-exception
  1. %%% @author Max Lapshin <max@maxidoors.ru> [http://erlyvideo.org]
  2. %%% @copyright 2010 Max Lapshin
  3. %%% @doc RTMP socket module.
  4. %%% Designed to look like rtmp mode of usual TCP socket. If you have used {packet, http}, you will
  5. %%% find many common behaviours.
  6. %%%
  7. %%% When working on server side, you should accept Socket::port() with:
  8. %%% <pre><code>
  9. %%% {ok, RTMP} = rtmp_socket:accept(Socket).
  10. %%% receive
  11. %%% {rtmp, RTMP, connected} ->
  12. %%% rtmp_socket:setopts(RTMP, [{active, once}]),
  13. %%% loop(RTMP)
  14. %%% end
  15. %%% loop(RTMP) ->
  16. %%% receive
  17. %%% {rtmp, RTMP, disconnect, Statistics} ->
  18. %%% ok;
  19. %%% {rtmp, RTMP, #rtmp_message{} = Message} ->
  20. %%% io:format("Message: ~p~n", [Message]),
  21. %%% loop(RTMP)
  22. %%% end.
  23. %%% </code></pre>
  24. %%%
  25. %%% You are strongly advised to use {active, once} mode, because in any other case it is very easy
  26. %%% to crash whole your server with OOM killer.
  27. %%% @reference See <a href="http://erlyvideo.org/rtmp" target="_top">http://erlyvideo.org/rtmp</a> for more information
  28. %%% @end
  29. %%%
  30. %%% This file is part of erlang-rtmp.
  31. %%%
  32. %%% erlang-rtmp is free software: you can redistribute it and/or modify
  33. %%% it under the terms of the GNU General Public License as published by
  34. %%% the Free Software Foundation, either version 3 of the License, or
  35. %%% (at your option) any later version.
  36. %%%
  37. %%% erlang-rtmp is distributed in the hope that it will be useful,
  38. %%% but WITHOUT ANY WARRANTY; without even the implied warranty of
  39. %%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  40. %%% GNU General Public License for more details.
  41. %%%
  42. %%% You should have received a copy of the GNU General Public License
  43. %%% along with erlang-rtmp. If not, see <http://www.gnu.org/licenses/>.
  44. %%%
  45. %%%---------------------------------------------------------------------------------------
  46. -module(rtmp_socket).
  47. -author('Max Lapshin <max@maxidoors.ru>').
  48. -include("../include/rtmp.hrl").
  49. -include("rtmp_private.hrl").
  50. -version(1.1).
  51. -export([accept/1, connect/1, start_link/1, getopts/2, setopts/2, getstat/2, getstat/1, send/2, get_socket/1]).
  52. -export([status/3, status/4, prepare_status/2, prepare_status/3, invoke/2, invoke/4, prepare_invoke/3, notify/4, prepare_notify/3]).
  53. -export([start_socket/2, start_server/3, start_server/4, set_socket/2]).
  54. %% gen_fsm callbacks
  55. -export([init/1, handle_event/3, handle_sync_event/4, handle_info/3, terminate/3, code_change/4]).
  56. -export([wait_for_socket_on_server/2, wait_for_socket_on_client/2, handshake_c1/2, handshake_c3/2, handshake_s1/2, loop/2, loop/3]).
  57. %% @spec (Port::integer(), Name::atom(), Callback::atom()) -> {ok, Pid::pid()}
  58. %% @doc Starts RTMP listener on port Port, registered under name Name with callback module Callback.
  59. %% Callback must export one function: create_client/1
  60. %% create_client(RTMPSocket::pid()) -> {ok, Pid::pid()}
  61. %%
  62. %% This function receives RTMPSocket, that is ready to send messages and after this callback function returns, this socket
  63. %% will send rtmp_message as it is defined in overview.
  64. %% @end
  65. -spec(start_server(Port::integer(), Name::atom(), Callback::atom()) -> {ok, Pid::pid()}).
  66. start_server(Port, Name, Callback) ->
  67. rtmp_sup:start_rtmp_listener(Port, Name, Callback, []).
  68. start_server(Port, Name, Callback, Args) ->
  69. rtmp_sup:start_rtmp_listener(Port, Name, Callback, Args).
  70. %% @spec (Socket::port()) -> {ok, RTMP::pid()}
  71. %% @doc Accepts client connection on socket Socket, starts RTMP decoder, passes socket to it
  72. %% and returns pid of newly created RTMP socket.
  73. %% @end
  74. -spec(accept(Socket::port()) -> {ok, RTMPSocket::pid()}).
  75. accept(Socket) ->
  76. {ok, Pid} = start_socket(accept, Socket),
  77. setopts(Pid, [{consumer,self()}]),
  78. {ok,Pid}.
  79. %% @spec (Socket::port()) -> {ok, RTMP::pid()}
  80. %% @doc Accepts client connection on socket Socket, starts RTMP decoder, passes socket to it
  81. %% and returns pid of newly created RTMP socket.
  82. %% @end
  83. -spec(connect(Socket::port()|string()) -> {ok, RTMPSocket::pid()}).
  84. connect(ServerSpec) when is_list(ServerSpec) ->
  85. {_Proto,_Auth,Host,Port,_Path,_Query} = http_uri2:parse(ServerSpec),
  86. {ok, Socket} = gen_tcp:connect(Host, Port, [binary, {active, false}, {packet, raw}]),
  87. connect(Socket);
  88. connect(Socket) when is_port(Socket) ->
  89. {ok, Pid} = start_socket(connect, Socket),
  90. setopts(Pid, [{consumer,self()}]),
  91. {ok,Pid}.
  92. %% @spec (Type::accept|connect, Socket::port) -> {ok, RTMP::pid()}
  93. %% @doc Starts RTMP socket with provided consumer and inititiate server or client connection
  94. %% @end
  95. -spec(start_socket(Type::connect|accept, Socket::port()) -> {ok, RTMP::pid()}).
  96. start_socket(Type, Socket) ->
  97. {ok, RTMP} = rtmp_sup:start_rtmp_socket(Type),
  98. case Socket of
  99. _ when is_port(Socket) -> gen_tcp:controlling_process(Socket, RTMP);
  100. _ -> ok
  101. end,
  102. set_socket(RTMP, Socket),
  103. {ok, RTMP}.
  104. set_socket(RTMP, Socket) ->
  105. gen_fsm:send_event(RTMP, {socket, Socket}).
  106. get_socket(RTMP) ->
  107. gen_fsm:sync_send_all_state_event(RTMP, get_socket, ?RTMP_TIMEOUT).
  108. %% @private
  109. start_link(Type) ->
  110. gen_fsm:start_link(?MODULE, [Type], []).
  111. %% @spec (RTMP::pid(), Options::[{Key, Value}]|{Key, Value}) -> ok
  112. %% @doc Just the same as {@link inet:setopts/2. inet:setopts/2} this function changes state of
  113. %% rtmp socket.<br/>
  114. %% Available options:
  115. %% <ul><li><code>chunk_size</code> - change outgoing chunk size</li>
  116. %% <li><code>window_size</code> - ask remote client to send read acknowlegement after WindowSize bytes</li>
  117. %% <li><code>amf_version</code> - change AMF0 to AMF3, but only before first call</li>
  118. %% <li><code>consumer</code> - change messages consumer</li>
  119. %% <li><code>debug = true|false</code> - dump all important packets or no</li>
  120. %% </ul>
  121. %% @end
  122. -spec(setopts(RTMP::rtmp_socket_pid(), Options::[{Key::atom(), Value::any()}]) -> ok).
  123. setopts(RTMP, Options) ->
  124. gen_fsm:sync_send_all_state_event(RTMP, {setopts, Options}).
  125. %% @spec (RTMP::pid(), Options::[{Key, Value}]|{Key, Value}) -> ok
  126. %% @doc Just the same as {@link inet:getopts/2. inet:getopts/2} this function gets state of
  127. %% rtmp socket.<br/>
  128. %% Available options:
  129. %% <ul>
  130. %% <li><code>chunk_size</code> - get outgoing chunk size</li>
  131. %% <li><code>window_size</code> - get remote client read acknowlegement window size bytes</li>
  132. %% <li><code>amf_version</code> - get AMF0 or AMF3</li>
  133. %% <li><code>consumer</code> - get messages consumer</li>
  134. %% <li><code>address</code> - get remote client IP and port</li>
  135. %% </ul>
  136. %% @end
  137. -spec(getopts(RTMP::rtmp_socket_pid(), Options::[Key::atom()]) -> ok).
  138. getopts(RTMP, Options) ->
  139. gen_fsm:sync_send_all_state_event(RTMP, {getopts, Options}, ?RTMP_TIMEOUT).
  140. %% @spec (RTMP::pid(), Stats::[Key]) -> Values::[{Key,Value}]
  141. %% @doc Just the same as {@link inet:getstats/2. inet:getstats/2} this function gets statistics of
  142. %% rtmp socket.<br/>
  143. %% Available options:
  144. %% <ul>
  145. %% <li><code>recv_oct</code> - number of bytes received to the socket.</li>
  146. %% <li><code>send_oct</code> - number of bytes sent from the socket</li>
  147. %% </ul>
  148. %% @end
  149. -spec(getstat(RTMP::rtmp_socket_pid(), Options::[Key::atom()]) -> ok).
  150. getstat(RTMP, Options) ->
  151. gen_fsm:sync_send_all_state_event(RTMP, {getstat, Options}, ?RTMP_TIMEOUT).
  152. %% @spec (RTMP::pid()) -> Values::[{Key,Value}]
  153. %% @doc Just the same as {@link inet:getstats/1. inet:getstats/1} this function gets statistics of
  154. %% rtmp socket.<br/>
  155. -spec(getstat(RTMP::rtmp_socket_pid()) -> ok).
  156. getstat(RTMP) ->
  157. gen_fsm:sync_send_all_state_event(RTMP, getstat, ?RTMP_TIMEOUT).
  158. %% @spec (RTMP::pid(), Message::rtmp_message()) -> ok
  159. %% @doc Sends message to client.
  160. %% @end
  161. -spec(send(RTMP::rtmp_socket_pid(), Message::rtmp_message()) -> ok).
  162. send(RTMP, Message) ->
  163. % io:format("Message ~p ~p ~p~n", [Message#rtmp_message.type, Message#rtmp_message.timestamp, Message#rtmp_message.stream_id]),
  164. % case process_info(RTMP, message_queue_len) of
  165. % {message_queue_len, Length} when Length < 20 -> RTMP ! Message;
  166. % {message_queue_len, Length} -> gen_fsm:sync_send_event(RTMP, Message, ?RTMP_TIMEOUT);
  167. % _ -> ok
  168. % end,
  169. RTMP ! Message,
  170. ok.
  171. notify(RTMP, StreamId, Name, Args) ->
  172. send(RTMP, prepare_notify(StreamId, Name, Args)).
  173. prepare_notify(StreamId, Name, Args) ->
  174. Arg = {object, lists:ukeymerge(1, [{level, <<"status">>}], lists:keysort(1, Args))},
  175. #rtmp_message{type = metadata, channel_id = rtmp_lib:channel_id(audio, StreamId), stream_id = StreamId, body = [Name, Arg], timestamp = same}.
  176. prepare_status(StreamId, Code) when is_list(Code) ->
  177. prepare_status(StreamId, list_to_binary(Code), <<"-">>);
  178. prepare_status(StreamId, Code) when is_binary(Code) ->
  179. prepare_status(StreamId, Code, <<"-">>).
  180. -spec(prepare_status(StreamId::non_neg_integer(), Code::any_string(), Description::any_string()) -> rtmp_message()).
  181. prepare_status(StreamId, Code, Description) ->
  182. Arg = {object, [
  183. {code, Code},
  184. {level, <<"status">>},
  185. {description, Description}
  186. ]},
  187. prepare_invoke(StreamId, onStatus, [Arg]).
  188. status(RTMP, StreamId, Code, Description) ->
  189. send(RTMP, prepare_status(StreamId, Code, Description)).
  190. -spec(status(RTMP::rtmp_socket_pid(), StreamId::integer(), Code::any_string()) -> ok).
  191. status(RTMP, StreamId, Code) when is_list(Code) ->
  192. status(RTMP, StreamId, list_to_binary(Code), <<"-">>);
  193. status(RTMP, StreamId, Code) when is_binary(Code) ->
  194. status(RTMP, StreamId, Code, <<"-">>).
  195. prepare_invoke(StreamId, Command, Args) when is_integer(StreamId) ->
  196. AMF = #rtmp_funcall{
  197. command = Command,
  198. type = invoke,
  199. id = 0,
  200. stream_id = StreamId,
  201. args = [null | Args ]},
  202. #rtmp_message{stream_id = StreamId, type = invoke, body = AMF}.
  203. invoke(RTMP, StreamId, Command, Args) ->
  204. send(RTMP, prepare_invoke(StreamId, Command, Args)).
  205. invoke(RTMP, #rtmp_funcall{stream_id = StreamId} = AMF) ->
  206. send(RTMP, #rtmp_message{stream_id = StreamId, type = invoke, body = AMF}).
  207. %% @private
  208. init([accept]) ->
  209. {ok, wait_for_socket_on_server, #rtmp_socket{channels = {}, out_channels = {}, active = false}, ?RTMP_TIMEOUT};
  210. init([connect]) ->
  211. {ok, wait_for_socket_on_client, #rtmp_socket{channels = {}, out_channels = {}, active = false}, ?RTMP_TIMEOUT}.
  212. %% @private
  213. wait_for_socket_on_server(timeout, State) ->
  214. {stop, normal, State};
  215. wait_for_socket_on_server({socket, Socket}, #rtmp_socket{} = State) when is_port(Socket) ->
  216. case inet:setopts(Socket, [{active, once}, {packet, raw}, binary]) of
  217. ok ->
  218. {ok, {IP, Port}} = inet:peername(Socket),
  219. {next_state, handshake_c1, State#rtmp_socket{socket = Socket, address = IP, port = Port}, ?RTMP_TIMEOUT};
  220. {error, _Error} ->
  221. {stop, normal, State}
  222. end;
  223. wait_for_socket_on_server({socket, Socket}, #rtmp_socket{} = State) when is_pid(Socket) ->
  224. erlang:monitor(process, Socket),
  225. {next_state, handshake_c1, State#rtmp_socket{socket = Socket}, ?RTMP_TIMEOUT}.
  226. %% @private
  227. -spec(wait_for_socket_on_client(Message::any(), Socket::rtmp_socket()) -> no_return()).
  228. wait_for_socket_on_client(timeout, State) ->
  229. {stop, normal, State};
  230. wait_for_socket_on_client({socket, Socket}, #rtmp_socket{} = State) ->
  231. inet:setopts(Socket, [{active, once}, {packet, raw}, binary]),
  232. {ok, {IP, Port}} = inet:peername(Socket),
  233. State1 = State#rtmp_socket{socket = Socket, address = IP, port = Port},
  234. send_data(State1, rtmp_handshake:c1()),
  235. {next_state, handshake_s1, State1, ?RTMP_TIMEOUT}.
  236. %% @private
  237. handshake_c1(timeout, State) ->
  238. {stop, normal, State}.
  239. %% @private
  240. handshake_c3(timeout, State) ->
  241. {stop, normal, State}.
  242. %% @private
  243. handshake_s1(timeout, State) ->
  244. {stop, normal, State}.
  245. %% @private
  246. loop(timeout, #rtmp_socket{pinged = false} = State) ->
  247. State1 = send_data(State, #rtmp_message{type = ping}),
  248. {next_state, loop, State1#rtmp_socket{pinged = true}, ?RTMP_TIMEOUT};
  249. loop(timeout, #rtmp_socket{consumer = Consumer} = State) ->
  250. Consumer ! {rtmp, self(), timeout},
  251. {stop, normal, State}.
  252. %% @private
  253. loop(#rtmp_message{} = Message, _From, State) ->
  254. State1 = send_data(State, Message),
  255. {reply, ok, loop, State1, ?RTMP_TIMEOUT}.
  256. -type(rtmp_option() ::active|amf_version|chunk_size|window_size|client_buffer|address).
  257. % -type(rtmp_option_value() ::{rtmp_option(), any()}).
  258. -spec(get_options(State::rtmp_socket(), Key::rtmp_option()|[rtmp_option()]) -> term()).
  259. get_options(State, active) ->
  260. {active, State#rtmp_socket.active};
  261. get_options(State, amf_version) ->
  262. {amf_version, State#rtmp_socket.amf_version};
  263. get_options(State, chunk_size) ->
  264. {chunk_size, State#rtmp_socket.server_chunk_size};
  265. get_options(State, window_size) ->
  266. {window_size, State#rtmp_socket.window_size};
  267. get_options(State, client_buffer) ->
  268. {client_buffer, State#rtmp_socket.client_buffer};
  269. get_options(State, debug) ->
  270. {debug, State#rtmp_socket.debug};
  271. get_options(State, address) ->
  272. {address, {State#rtmp_socket.address, State#rtmp_socket.port}};
  273. get_options(_State, []) ->
  274. [];
  275. get_options(State, [Key | Options]) ->
  276. [get_options(State, Key) | get_options(State, Options)].
  277. get_stat(State) ->
  278. get_stat(State, [recv_oct, send_oct]).
  279. get_stat(State, recv_oct) ->
  280. {recv_oct, State#rtmp_socket.bytes_read};
  281. get_stat(State, send_oct) ->
  282. {send_oct, State#rtmp_socket.bytes_sent};
  283. get_stat(_State, []) ->
  284. [];
  285. get_stat(State, [Key | Options]) ->
  286. [get_stat(State, Key) | get_stat(State, Options)].
  287. set_options(State, [{amf_version, Version} | Options]) ->
  288. set_options(State#rtmp_socket{amf_version = Version}, Options);
  289. set_options(#rtmp_socket{socket = Socket, buffer = Data} = State, [{active, Active} | Options]) ->
  290. State1 = flush_send(State#rtmp_socket{active = Active}),
  291. State2 = case Active of
  292. false ->
  293. inet:setopts(Socket, [{active, false}]),
  294. State1;
  295. true ->
  296. inet:setopts(Socket, [{active, true}]),
  297. State1;
  298. once when size(Data) > 0 ->
  299. handle_rtmp_data(State1);
  300. once ->
  301. activate_socket(Socket),
  302. State1
  303. end,
  304. set_options(State2, Options);
  305. set_options(#rtmp_socket{} = State, [{debug, Debug} | Options]) ->
  306. io:format("Set debug to ~p~n", [Debug]),
  307. set_options(State#rtmp_socket{debug = Debug}, Options);
  308. set_options(#rtmp_socket{consumer = undefined} = State, [{consumer, Consumer} | Options]) ->
  309. erlang:monitor(process, Consumer),
  310. set_options(State#rtmp_socket{consumer = Consumer}, Options);
  311. set_options(State, [{chunk_size, ChunkSize} | Options]) ->
  312. State1 = send_data(State, #rtmp_message{type = chunk_size, body = ChunkSize}),
  313. set_options(State1#rtmp_socket{server_chunk_size = ChunkSize}, Options);
  314. set_options(State, []) -> State.
  315. %%-------------------------------------------------------------------------
  316. %% Func: handle_event/3
  317. %% Returns: {next_state, NextStateName, NextStateData} |
  318. %% {next_state, NextStateName, NextStateData, Timeout} |
  319. %% {stop, Reason, NewStateData}
  320. %% @private
  321. %%-------------------------------------------------------------------------
  322. handle_event({setopts, Options}, StateName, State) ->
  323. NewState = set_options(State, Options),
  324. {next_state, StateName, NewState, ?RTMP_TIMEOUT};
  325. handle_event(Event, StateName, StateData) ->
  326. {stop, {StateName, undefined_event, Event}, StateData}.
  327. %%-------------------------------------------------------------------------
  328. %% Func: handle_sync_event/4
  329. %% Returns: {next_state, NextStateName, NextStateData} |
  330. %% {next_state, NextStateName, NextStateData, Timeout} |
  331. %% {reply, Reply, NextStateName, NextStateData} |
  332. %% {reply, Reply, NextStateName, NextStateData, Timeout} |
  333. %% {stop, Reason, NewStateData} |
  334. %% {stop, Reason, Reply, NewStateData}
  335. %% @private
  336. %%-------------------------------------------------------------------------
  337. handle_sync_event({getopts, Options}, _From, StateName, State) ->
  338. {reply, get_options(State, Options), StateName, State, ?RTMP_TIMEOUT};
  339. handle_sync_event({getstat, Options}, _From, StateName, State) ->
  340. {reply, get_stat(State, Options), StateName, State, ?RTMP_TIMEOUT};
  341. handle_sync_event(getstat, _From, StateName, State) ->
  342. {reply, get_stat(State), StateName, State, ?RTMP_TIMEOUT};
  343. handle_sync_event(get_socket, _From, StateName, #rtmp_socket{socket = Socket} = State) ->
  344. {reply, Socket, StateName, State, ?RTMP_TIMEOUT};
  345. handle_sync_event({setopts, Options}, _From, StateName, State) ->
  346. NewState = set_options(State, Options),
  347. {reply, ok, StateName, NewState, ?RTMP_TIMEOUT};
  348. handle_sync_event(Event, _From, StateName, StateData) ->
  349. {stop, {StateName, undefined_event, Event}, StateData}.
  350. %%-------------------------------------------------------------------------
  351. %% Func: handle_info/3
  352. %% Returns: {next_state, NextStateName, NextStateData} |
  353. %% {next_state, NextStateName, NextStateData, Timeout} |
  354. %% {stop, Reason, NewStateData}
  355. %% @private
  356. %%-------------------------------------------------------------------------
  357. handle_info({tcp, Socket, Data}, handshake_c1, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead} = State) when size(Buffer) + size(Data) < ?HS_BODY_LEN + 1 ->
  358. activate_socket(Socket),
  359. {next_state, handshake_c1, State#rtmp_socket{buffer = <<Buffer/binary, Data/binary>>, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
  360. handle_info({tcp, Socket, Data}, handshake_c1, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead} = State) ->
  361. activate_socket(Socket),
  362. <<Handshake:(?HS_BODY_LEN+1)/binary, Rest/binary>> = <<Buffer/binary, Data/binary>>,
  363. State1 = case rtmp_handshake:server(Handshake) of
  364. {uncrypted, Reply} ->
  365. send_data(State, Reply);
  366. {crypted, Reply, KeyIn, KeyOut} ->
  367. NewState = send_data(State, Reply),
  368. ?D({"Established RTMPE connection"}),
  369. NewState#rtmp_socket{key_in = KeyIn, key_out = KeyOut}
  370. end,
  371. {next_state, 'handshake_c3', State1#rtmp_socket{buffer = Rest, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
  372. handle_info({tcp, Socket, Data}, handshake_c3, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead} = State) when size(Buffer) + size(Data) < ?HS_BODY_LEN ->
  373. activate_socket(Socket),
  374. {next_state, handshake_c3, State#rtmp_socket{buffer = <<Buffer/binary, Data/binary>>, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
  375. handle_info({tcp, Socket, Data}, handshake_c3, #rtmp_socket{socket=Socket, consumer = Consumer, buffer = Buffer, bytes_read = BytesRead, active = Active, key_in = KeyIn} = State) ->
  376. <<_HandShakeC3:?HS_BODY_LEN/binary, CryptedData/binary>> = <<Buffer/binary, Data/binary>>,
  377. Consumer ! {rtmp, self(), connected},
  378. case Active of
  379. true -> inet:setopts(Socket, [{active, true}]);
  380. _ -> ok
  381. end,
  382. {NewKeyIn, Rest} = case KeyIn of
  383. undefined -> {undefined, CryptedData};
  384. _ -> rtmpe:crypt(KeyIn, CryptedData)
  385. end,
  386. % ?D({decrypt, Rest, CryptedData == element(2, rtmpe:crypt(NewKeyIn, Rest))}),
  387. State1 = State#rtmp_socket{bytes_read = BytesRead + size(Data), key_in = NewKeyIn, buffer = Rest},
  388. case Rest of
  389. <<>> -> {next_state, loop, State1, ?RTMP_TIMEOUT};
  390. _ -> {next_state, loop, handle_rtmp_data(State1), ?RTMP_TIMEOUT}
  391. end;
  392. handle_info({tcp, Socket, Data}, handshake_s1, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead} = State) when size(Buffer) + size(Data) < ?HS_BODY_LEN*2 + 1 ->
  393. activate_socket(Socket),
  394. {next_state, handshake_s1, State#rtmp_socket{buffer = <<Buffer/binary, Data/binary>>, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
  395. handle_info({tcp, Socket, Data}, handshake_s1, #rtmp_socket{socket=Socket, consumer = Consumer, buffer = Buffer, bytes_read = BytesRead, active = Active} = State) ->
  396. <<?HS_UNCRYPTED, S1:?HS_BODY_LEN/binary, _S2:?HS_BODY_LEN/binary, Rest/binary>> = <<Buffer/binary, Data/binary>>,
  397. send_data(State, rtmp_handshake:c2(S1)),
  398. Consumer ! {rtmp, self(), connected},
  399. case Active of
  400. true -> inet:setopts(Socket, [{active, true}]);
  401. _ -> ok
  402. end,
  403. {next_state, loop, handle_rtmp_data(State#rtmp_socket{bytes_read = BytesRead + size(Data), buffer = Rest}), ?RTMP_TIMEOUT};
  404. handle_info({tcp, Socket, CryptedData}, loop, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead, bytes_unack = BytesUnack, key_in = KeyIn} = State) ->
  405. State1 = flush_send(State),
  406. {NewKeyIn, Data} = case KeyIn of
  407. undefined -> {undefined, CryptedData};
  408. _ -> rtmpe:crypt(KeyIn, CryptedData)
  409. end,
  410. {next_state, loop, handle_rtmp_data(State1#rtmp_socket{bytes_read = BytesRead + size(Data), bytes_unack = BytesUnack + size(Data), key_in = NewKeyIn, buffer = <<Buffer/binary, Data/binary>>}), ?RTMP_TIMEOUT};
  411. handle_info({tcp_closed, Socket}, _StateName, #rtmp_socket{socket = Socket, consumer = Consumer} = StateData) ->
  412. Consumer ! {rtmp, self(), disconnect, get_stat(StateData)},
  413. {stop, normal, StateData};
  414. handle_info(#rtmp_message{} = Message, loop, State) ->
  415. State1 = send_data(State, Message),
  416. State2 = flush_send(State1),
  417. {next_state, loop, State2, ?RTMP_TIMEOUT};
  418. handle_info({rtmpt, RTMPT, alive}, StateName, #rtmp_socket{socket = RTMPT} = State) ->
  419. {next_state, StateName, State#rtmp_socket{pinged = false}, ?RTMP_TIMEOUT};
  420. handle_info({rtmpt, RTMPT, Data}, StateName, State) ->
  421. handle_info({tcp, RTMPT, Data}, StateName, State);
  422. handle_info({'DOWN', _, process, _Client, _Reason}, _StateName, State) ->
  423. {stop, normal, State};
  424. handle_info(_Info, StateName, StateData) ->
  425. error_logger:error_msg("Unknown message to rtmp socket: ~p ~p ~p~n", [_Info, StateName, StateData]),
  426. {next_state, StateName, StateData, ?RTMP_TIMEOUT}.
  427. % flush_send(State) -> flush_send([], State).
  428. flush_send(State) ->
  429. receive
  430. #rtmp_message{} = Message ->
  431. NewState = send_data(State, Message),
  432. % ?D({out, Message}),
  433. % {NewState, Data} = rtmp:encode(State, Message),
  434. % flush_send([Data | Packet], NewState)
  435. flush_send(NewState)
  436. after
  437. 0 -> State
  438. end.
  439. activate_socket(Socket) when is_port(Socket) ->
  440. inet:setopts(Socket, [{active, once}]);
  441. activate_socket(Socket) when is_pid(Socket) ->
  442. ok.
  443. send_data(#rtmp_socket{sent_audio_notify = false} = Socket,
  444. #rtmp_message{type = audio, timestamp = DTS, stream_id = StreamId, body = Body} = Message) when size(Body) > 0 ->
  445. State1 = send_data(Socket#rtmp_socket{sent_audio_notify = true}, rtmp_lib:empty_audio(StreamId, DTS)),
  446. send_data(State1, Message#rtmp_message{ts_type = new});
  447. send_data(#rtmp_socket{sent_video_notify = false} = Socket, #rtmp_message{type = video, timestamp = DTS, stream_id = StreamId} = Message) ->
  448. Msg = [
  449. #rtmp_message{type = video, channel_id = rtmp_lib:channel_id(video, StreamId), timestamp = DTS, stream_id = StreamId, body = <<87,0>>, ts_type = new},
  450. #rtmp_message{type = video, channel_id = rtmp_lib:channel_id(video, StreamId), timestamp = DTS, stream_id = StreamId, body = <<23,2,0,0,0>>, ts_type = new},
  451. #rtmp_message{type = video, channel_id = rtmp_lib:channel_id(video, StreamId), timestamp = DTS, stream_id = StreamId, body = <<87,1>>, ts_type = delta}
  452. ],
  453. State2 = lists:foldl(fun(M, State1) ->
  454. send_data(State1, M)
  455. end, Socket#rtmp_socket{sent_video_notify = true}, Msg),
  456. send_data(State2, Message);
  457. send_data(#rtmp_socket{sent_audio_notify = true} = Socket, #rtmp_message{type = stream_end} = Message) ->
  458. send_data(Socket#rtmp_socket{sent_audio_notify = false}, Message);
  459. send_data(#rtmp_socket{sent_video_notify = true} = Socket, #rtmp_message{type = stream_end} = Message) ->
  460. send_data(Socket#rtmp_socket{sent_video_notify = false}, Message);
  461. send_data(#rtmp_socket{socket = Socket, key_out = KeyOut} = State, Message) ->
  462. case State#rtmp_socket.debug of
  463. true -> print_rtmp_message(out, Message);
  464. _ -> ok
  465. end,
  466. {NewState, Data} = case Message of
  467. #rtmp_message{} ->
  468. % ?D({Socket,Message#rtmp_message.type, Message#rtmp_message.timestamp, Message#rtmp_message.ts_type}),
  469. rtmp:encode(State, Message);
  470. _ ->
  471. {State, Message}
  472. end,
  473. {NewKeyOut, Crypt} = case KeyOut of
  474. undefined -> {undefined, Data};
  475. _ -> rtmpe:crypt(KeyOut, Data)
  476. end,
  477. % (catch rtmp_stat_collector:out_bytes(self(), iolist_size(Crypt))),
  478. if
  479. is_port(Socket) ->
  480. gen_tcp:send(Socket, Crypt);
  481. is_pid(Socket) ->
  482. rtmpt:write(Socket, Crypt)
  483. end,
  484. NewState#rtmp_socket{key_out = NewKeyOut}.
  485. handle_rtmp_data(#rtmp_socket{buffer = <<>>} = State) ->
  486. % ?D({empty_input, erlang:get_stacktrace()}),
  487. State;
  488. handle_rtmp_data(#rtmp_socket{bytes_unack = Bytes, window_size = Window} = State) when Bytes >= Window ->
  489. State1 = send_data(State, #rtmp_message{type = ack_read}),
  490. handle_rtmp_data(State1);
  491. handle_rtmp_data(#rtmp_socket{buffer = Data} = State) ->
  492. got_rtmp_message(rtmp:decode(State, Data)).
  493. print_rtmp_message(InOut, #rtmp_message{channel_id = Channel, ts_type = TSType, timestamp = TS, type = Type, stream_id = StreamId, body = Body}) ->
  494. DecodedBody = case Type of
  495. video when size(Body) > 10 -> erlang:setelement(5, flv:decode_video_tag(Body), size(Body));
  496. audio when size(Body) > 0 -> erlang:setelement(7, flv:decode_audio_tag(Body), size(Body));
  497. _ -> Body
  498. end,
  499. io:format("~p ~p ~p ~p ~p ~p ~p~n", [InOut, Channel, TSType, TS, Type, StreamId, DecodedBody]),
  500. ok;
  501. print_rtmp_message(InOut, Msg) ->
  502. io:format("~p ~s~n", [InOut, io_lib_pretty_limited:print(Msg, 20)]),
  503. ok.
  504. got_rtmp_message({#rtmp_socket{debug = true}, Msg, _} = Message) ->
  505. print_rtmp_message(in, Msg),
  506. handle_rtmp_message(Message);
  507. got_rtmp_message(Decoded) ->
  508. handle_rtmp_message(Decoded).
  509. handle_rtmp_message({#rtmp_socket{consumer = Consumer} = State, #rtmp_message{type = window_size, body = Size} = Message, Rest}) ->
  510. Consumer ! {rtmp, self(), Message},
  511. rtmp_message_sent(State#rtmp_socket{window_size = Size, buffer = Rest});
  512. handle_rtmp_message({#rtmp_socket{consumer = Consumer, pinged = true} = State, #rtmp_message{type = pong} = Message, Rest}) ->
  513. Consumer ! {rtmp, self(), Message},
  514. rtmp_message_sent(State#rtmp_socket{pinged = false, buffer = Rest});
  515. handle_rtmp_message({#rtmp_socket{consumer = Consumer} = State, #rtmp_message{type = ping, body = Timestamp} = Message, Rest}) ->
  516. Consumer ! {rtmp, self(), Message},
  517. send_data(State, #rtmp_message{type = pong, body = Timestamp}),
  518. rtmp_message_sent(State#rtmp_socket{buffer = Rest});
  519. handle_rtmp_message({#rtmp_socket{consumer = Consumer} = State, Message, Rest}) ->
  520. Consumer ! {rtmp, self(), Message},
  521. rtmp_message_sent(State#rtmp_socket{buffer = Rest});
  522. handle_rtmp_message({#rtmp_socket{socket=Socket} = State, Rest}) ->
  523. activate_socket(Socket),
  524. State#rtmp_socket{buffer = Rest}.
  525. rtmp_message_sent(#rtmp_socket{active = true,socket = Socket} = State) ->
  526. activate_socket(Socket),
  527. handle_rtmp_data(State);
  528. rtmp_message_sent(#rtmp_socket{active = _} = State) ->
  529. handle_rtmp_data(State);
  530. rtmp_message_sent(State) ->
  531. State.
  532. %%-------------------------------------------------------------------------
  533. %% Func: terminate/3
  534. %% Purpose: Shutdown the fsm
  535. %% Returns: any
  536. %% @private
  537. %%-------------------------------------------------------------------------
  538. terminate(_Reason, _StateName, #rtmp_socket{socket=Socket}) when is_pid(Socket) ->
  539. gen_server:cast(Socket, close),
  540. ok;
  541. terminate(_Reason, _StateName, _State) ->
  542. ok.
  543. %%-------------------------------------------------------------------------
  544. %% Func: code_change/4
  545. %% Purpose: Convert process state when code is changed
  546. %% Returns: {ok, NewState, NewStateData}
  547. %% @private
  548. %%-------------------------------------------------------------------------
  549. code_change(_OldVersion, _StateName, _State, _Extra) ->
  550. ok.