/apps/rtmp/src/rtmp_socket.erl

  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
 52-export([accept/1, connect/1, start_link/1, getopts/2, setopts/2, getstat/2, getstat/1, send/2, get_socket/1]).
 53-export([status/3, status/4, prepare_status/2, prepare_status/3, invoke/2, invoke/4, prepare_invoke/3, notify/4, prepare_notify/3]).
 54
 55-export([start_socket/2, start_server/3, start_server/4, set_socket/2]).
 56  
 57
 58
 59%% gen_fsm callbacks
 60-export([init/1, handle_event/3, handle_sync_event/4, handle_info/3, terminate/3, code_change/4]).
 61
 62-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]).
 63
 64%% @spec (Port::integer(), Name::atom(), Callback::atom()) -> {ok, Pid::pid()}
 65%% @doc Starts RTMP listener on port Port, registered under name Name with callback module Callback.
 66%% Callback must export one function: create_client/1
 67%% create_client(RTMPSocket::pid()) -> {ok, Pid::pid()}
 68%%
 69%% This function receives RTMPSocket, that is ready to send messages and after this callback function returns, this socket
 70%% will send rtmp_message as it is defined in overview.
 71%% @end
 72-spec(start_server(Port::integer(), Name::atom(), Callback::atom()) -> {ok, Pid::pid()}).
 73start_server(Port, Name, Callback) ->
 74  rtmp_sup:start_rtmp_listener(Port, Name, Callback, []).
 75
 76start_server(Port, Name, Callback, Args) ->
 77  rtmp_sup:start_rtmp_listener(Port, Name, Callback, Args).
 78
 79
 80%% @spec (Socket::port()) -> {ok, RTMP::pid()}
 81%% @doc Accepts client connection on socket Socket, starts RTMP decoder, passes socket to it
 82%% and returns pid of newly created RTMP socket.
 83%% @end
 84-spec(accept(Socket::port()) -> {ok, RTMPSocket::pid()}).
 85accept(Socket) ->
 86  {ok, Pid} = start_socket(accept, Socket),
 87  setopts(Pid, [{consumer,self()}]),
 88  {ok,Pid}.
 89  
 90%% @spec (Socket::port()) -> {ok, RTMP::pid()}
 91%% @doc Accepts client connection on socket Socket, starts RTMP decoder, passes socket to it
 92%% and returns pid of newly created RTMP socket.
 93%% @end
 94-spec(connect(Socket::port()|string()) -> {ok, RTMPSocket::pid()}).
 95connect(ServerSpec) when is_list(ServerSpec) ->
 96  {_Proto,_Auth,Host,Port,_Path,_Query} = http_uri2:parse(ServerSpec),
 97  {ok, Socket} = gen_tcp:connect(Host, Port, [binary, {active, false}, {packet, raw}]),
 98  connect(Socket);
 99
100connect(Socket) when is_port(Socket) ->
101  {ok, Pid} = start_socket(connect, Socket),
102  setopts(Pid, [{consumer,self()}]),
103  {ok,Pid}.
104
105%% @spec (Type::accept|connect, Socket::port) -> {ok, RTMP::pid()}
106%% @doc Starts RTMP socket with provided consumer and inititiate server or client connection
107%% @end
108-spec(start_socket(Type::connect|accept, Socket::port()) -> {ok, RTMP::pid()}).
109start_socket(Type, Socket) ->
110  {ok, RTMP} = rtmp_sup:start_rtmp_socket(Type),
111  case Socket of
112    _ when is_port(Socket) -> gen_tcp:controlling_process(Socket, RTMP);
113    _ -> ok
114  end,
115  set_socket(RTMP, Socket),
116  {ok, RTMP}.
117
118
119set_socket(RTMP, Socket) ->
120  gen_fsm:send_event(RTMP, {socket, Socket}).
121
122get_socket(RTMP) ->
123  gen_fsm:sync_send_all_state_event(RTMP, get_socket, ?RTMP_TIMEOUT).
124  
125  
126%% @private
127start_link(Type) ->
128  gen_fsm:start_link(?MODULE, [Type], []).
129  
130  
131%% @spec (RTMP::pid(), Options::[{Key, Value}]|{Key, Value}) -> ok
132%% @doc Just the same as {@link inet:setopts/2. inet:setopts/2} this function changes state of 
133%% rtmp socket.<br/>
134%%  Available options:
135%%  <ul><li><code>chunk_size</code> - change outgoing chunk size</li>
136%%  <li><code>window_size</code> - ask remote client to send read acknowlegement after WindowSize bytes</li>
137%%  <li><code>amf_version</code> - change AMF0 to AMF3, but only before first call</li>
138%%  <li><code>consumer</code> - change messages consumer</li>
139%%  <li><code>debug = true|false</code> - dump all important packets or no</li>
140%% </ul>
141%% @end
142-spec(setopts(RTMP::rtmp_socket_pid(), Options::[{Key::atom(), Value::any()}]) -> ok).
143setopts(RTMP, Options) ->
144  gen_fsm:sync_send_all_state_event(RTMP, {setopts, Options}).
145
146%% @spec (RTMP::pid(), Options::[{Key, Value}]|{Key, Value}) -> ok
147%% @doc Just the same as {@link inet:getopts/2. inet:getopts/2} this function gets state of 
148%% rtmp socket.<br/>
149%%  Available options:
150%%  <ul>
151%%  <li><code>chunk_size</code> - get outgoing chunk size</li>
152%%  <li><code>window_size</code> - get remote client read acknowlegement window size bytes</li>
153%%  <li><code>amf_version</code> - get AMF0 or AMF3</li>
154%%  <li><code>consumer</code> - get messages consumer</li>
155%%  <li><code>address</code> - get remote client IP and port</li>
156%% </ul>
157%% @end
158-spec(getopts(RTMP::rtmp_socket_pid(), Options::[Key::atom()]) -> ok).
159getopts(RTMP, Options) ->
160  gen_fsm:sync_send_all_state_event(RTMP, {getopts, Options}, ?RTMP_TIMEOUT).
161
162%% @spec (RTMP::pid(), Stats::[Key]) -> Values::[{Key,Value}]
163%% @doc Just the same as {@link inet:getstats/2. inet:getstats/2} this function gets statistics of 
164%% rtmp socket.<br/>
165%%  Available options:
166%%  <ul>
167%%  <li><code>recv_oct</code> - number of bytes received to the socket.</li>
168%%  <li><code>send_oct</code> - number of bytes sent from the socket</li>
169%% </ul>
170%% @end
171-spec(getstat(RTMP::rtmp_socket_pid(), Options::[Key::atom()]) -> ok).
172getstat(RTMP, Options) ->
173  gen_fsm:sync_send_all_state_event(RTMP, {getstat, Options}, ?RTMP_TIMEOUT).
174
175%% @spec (RTMP::pid()) -> Values::[{Key,Value}]
176%% @doc Just the same as {@link inet:getstats/1. inet:getstats/1} this function gets statistics of 
177%% rtmp socket.<br/>
178-spec(getstat(RTMP::rtmp_socket_pid()) -> ok).
179getstat(RTMP) ->
180  gen_fsm:sync_send_all_state_event(RTMP, getstat, ?RTMP_TIMEOUT).
181
182  
183%% @spec (RTMP::pid(), Message::rtmp_message()) -> ok
184%% @doc Sends message to client.
185%% @end
186-spec(send(RTMP::rtmp_socket_pid(), Message::rtmp_message()) -> ok).
187send(RTMP, Message) ->
188  % io:format("Message ~p ~p ~p~n", [Message#rtmp_message.type, Message#rtmp_message.timestamp, Message#rtmp_message.stream_id]),
189  % case process_info(RTMP, message_queue_len) of
190  %   {message_queue_len, Length} when Length < 20 -> RTMP ! Message;
191  %   {message_queue_len, Length} -> gen_fsm:sync_send_event(RTMP, Message, ?RTMP_TIMEOUT);
192  %   _ -> ok
193  % end,
194  RTMP ! Message,
195  ok.
196  
197
198
199notify(RTMP, StreamId, Name, Args) ->
200  send(RTMP, prepare_notify(StreamId, Name, Args)).
201
202prepare_notify(StreamId, Name, Args) ->
203  Arg = {object, lists:ukeymerge(1, [{level, <<"status">>}], lists:keysort(1, Args))},
204  #rtmp_message{type = metadata, channel_id = rtmp_lib:channel_id(audio, StreamId), stream_id = StreamId, body = [Name, Arg], timestamp = same}.
205
206
207prepare_status(StreamId, Code) when is_list(Code) ->
208  prepare_status(StreamId, list_to_binary(Code), <<"-">>);
209
210prepare_status(StreamId, Code) when is_binary(Code) ->
211  prepare_status(StreamId, Code, <<"-">>).
212
213
214-spec(prepare_status(StreamId::non_neg_integer(), Code::any_string(), Description::any_string()) -> rtmp_message()).
215prepare_status(StreamId, Code, Description) ->
216  Arg = {object, [
217    {code, Code}, 
218    {level, <<"status">>}, 
219    {description, Description}
220  ]},
221  prepare_invoke(StreamId, onStatus, [Arg]).
222
223
224status(RTMP, StreamId, Code, Description) ->
225  send(RTMP, prepare_status(StreamId, Code, Description)).
226
227-spec(status(RTMP::rtmp_socket_pid(), StreamId::integer(), Code::any_string()) -> ok).
228status(RTMP, StreamId, Code) when is_list(Code) ->
229  status(RTMP, StreamId, list_to_binary(Code), <<"-">>);
230
231status(RTMP, StreamId, Code) when is_binary(Code) ->
232  status(RTMP, StreamId, Code, <<"-">>).
233
234  
235prepare_invoke(StreamId, Command, Args) when is_integer(StreamId) ->
236  AMF = #rtmp_funcall{
237      command = Command,
238        type = invoke,
239        id = 0,
240        stream_id = StreamId,
241        args = [null | Args ]},
242  #rtmp_message{stream_id = StreamId, type = invoke, body = AMF}.
243
244
245invoke(RTMP, StreamId, Command, Args) ->
246  send(RTMP, prepare_invoke(StreamId, Command, Args)).
247
248invoke(RTMP, #rtmp_funcall{stream_id = StreamId} = AMF) ->
249  send(RTMP, #rtmp_message{stream_id = StreamId, type = invoke, body = AMF}).
250
251
252  
253%% @private  
254init([accept]) ->
255  {ok, wait_for_socket_on_server, #rtmp_socket{channels = {}, out_channels = {}, active = false}, ?RTMP_TIMEOUT};
256
257init([connect]) ->
258  {ok, wait_for_socket_on_client, #rtmp_socket{channels = {}, out_channels = {}, active = false}, ?RTMP_TIMEOUT}.
259
260
261
262%% @private 
263
264wait_for_socket_on_server(timeout, State) ->
265  {stop, normal, State};
266
267wait_for_socket_on_server({socket, Socket}, #rtmp_socket{} = State) when is_port(Socket) ->
268  case inet:setopts(Socket, [{active, once}, {packet, raw}, binary]) of
269    ok ->
270      {ok, {IP, Port}} = inet:peername(Socket),
271      {next_state, handshake_c1, State#rtmp_socket{socket = Socket, address = IP, port = Port}, ?RTMP_TIMEOUT};
272    {error, _Error} ->
273      {stop, normal, State}
274  end;
275      
276wait_for_socket_on_server({socket, Socket}, #rtmp_socket{} = State) when is_pid(Socket) ->
277  erlang:monitor(process, Socket),
278  {next_state, handshake_c1, State#rtmp_socket{socket = Socket}, ?RTMP_TIMEOUT}.
279
280
281%% @private  
282-spec(wait_for_socket_on_client(Message::any(), Socket::rtmp_socket()) -> no_return()).
283
284wait_for_socket_on_client(timeout, State) ->
285  {stop, normal, State};
286
287wait_for_socket_on_client({socket, Socket}, #rtmp_socket{} = State) ->
288  inet:setopts(Socket, [{active, once}, {packet, raw}, binary]),
289  {ok, {IP, Port}} = inet:peername(Socket),
290  State1 = State#rtmp_socket{socket = Socket, address = IP, port = Port},
291  send_data(State1, rtmp_handshake:c1()),
292  {next_state, handshake_s1, State1, ?RTMP_TIMEOUT}.
293
294%% @private  
295handshake_c1(timeout, State) ->
296  {stop, normal, State}.
297
298%% @private  
299handshake_c3(timeout, State) ->
300  {stop, normal, State}.
301
302%% @private  
303handshake_s1(timeout, State) ->
304  {stop, normal, State}.
305  
306%% @private  
307loop(timeout, #rtmp_socket{pinged = false} = State) ->
308  State1 = send_data(State, #rtmp_message{type = ping}),
309  {next_state, loop, State1#rtmp_socket{pinged = true}, ?RTMP_TIMEOUT};
310  
311loop(timeout, #rtmp_socket{consumer = Consumer} = State) ->
312  Consumer ! {rtmp, self(), timeout},
313  {stop, normal, State}.
314
315
316%% @private
317loop(#rtmp_message{} = Message, _From, State) ->
318  State1 = send_data(State, Message),
319  {reply, ok, loop, State1, ?RTMP_TIMEOUT}.
320
321
322-type(rtmp_option() ::active|amf_version|chunk_size|window_size|client_buffer|address).
323% -type(rtmp_option_value() ::{rtmp_option(), any()}).
324-spec(get_options(State::rtmp_socket(), Key::rtmp_option()|[rtmp_option()]) -> term()).
325
326get_options(State, active) ->
327  {active, State#rtmp_socket.active};
328
329get_options(State, amf_version) ->
330  {amf_version, State#rtmp_socket.amf_version};
331
332get_options(State, chunk_size) ->
333  {chunk_size, State#rtmp_socket.server_chunk_size};
334
335get_options(State, window_size) ->
336  {window_size, State#rtmp_socket.window_size};
337
338get_options(State, client_buffer) ->
339  {client_buffer, State#rtmp_socket.client_buffer};
340
341get_options(State, debug) ->
342  {debug, State#rtmp_socket.debug};
343
344get_options(State, address) ->
345  {address, {State#rtmp_socket.address, State#rtmp_socket.port}};
346
347get_options(_State, []) ->
348  [];
349  
350get_options(State, [Key | Options]) ->
351  [get_options(State, Key) | get_options(State, Options)].
352
353get_stat(State) ->
354  get_stat(State, [recv_oct, send_oct]).
355
356get_stat(State, recv_oct) ->
357  {recv_oct, State#rtmp_socket.bytes_read};
358
359get_stat(State, send_oct) ->
360  {send_oct, State#rtmp_socket.bytes_sent};
361
362get_stat(_State, []) ->
363  [];
364
365get_stat(State, [Key | Options]) ->
366  [get_stat(State, Key) | get_stat(State, Options)].
367
368set_options(State, [{amf_version, Version} | Options]) ->
369  set_options(State#rtmp_socket{amf_version = Version}, Options);
370
371set_options(#rtmp_socket{socket = Socket, buffer = Data} = State, [{active, Active} | Options]) ->
372  State1 = flush_send(State#rtmp_socket{active = Active}),
373  State2 = case Active of
374    false -> 
375      inet:setopts(Socket, [{active, false}]),
376      State1;
377    true ->
378      inet:setopts(Socket, [{active, true}]),
379      State1;
380    once when size(Data) > 0 ->
381      handle_rtmp_data(State1);
382    once ->
383      activate_socket(Socket),
384      State1
385  end,
386  set_options(State2, Options);
387
388set_options(#rtmp_socket{} = State, [{debug, Debug} | Options]) ->
389  io:format("Set debug to ~p~n", [Debug]),
390  set_options(State#rtmp_socket{debug = Debug}, Options);
391
392set_options(#rtmp_socket{consumer = undefined} = State, [{consumer, Consumer} | Options]) ->
393  erlang:monitor(process, Consumer),
394  set_options(State#rtmp_socket{consumer = Consumer}, Options);
395
396set_options(State, [{chunk_size, ChunkSize} | Options]) ->
397  State1 = send_data(State, #rtmp_message{type = chunk_size, body = ChunkSize}),
398  set_options(State1#rtmp_socket{server_chunk_size = ChunkSize}, Options);
399
400set_options(State, []) -> State.
401  
402%%-------------------------------------------------------------------------
403%% Func: handle_event/3
404%% Returns: {next_state, NextStateName, NextStateData}          |
405%%          {next_state, NextStateName, NextStateData, Timeout} |
406%%          {stop, Reason, NewStateData}
407%% @private
408%%-------------------------------------------------------------------------
409handle_event({setopts, Options}, StateName, State) ->
410  NewState = set_options(State, Options),
411  {next_state, StateName, NewState, ?RTMP_TIMEOUT};
412
413handle_event(Event, StateName, StateData) ->
414  {stop, {StateName, undefined_event, Event}, StateData}.
415
416
417%%-------------------------------------------------------------------------
418%% Func: handle_sync_event/4
419%% Returns: {next_state, NextStateName, NextStateData}            |
420%%          {next_state, NextStateName, NextStateData, Timeout}   |
421%%          {reply, Reply, NextStateName, NextStateData}          |
422%%          {reply, Reply, NextStateName, NextStateData, Timeout} |
423%%          {stop, Reason, NewStateData}                          |
424%%          {stop, Reason, Reply, NewStateData}
425%% @private
426%%-------------------------------------------------------------------------
427
428handle_sync_event({getopts, Options}, _From, StateName, State) ->
429  {reply, get_options(State, Options), StateName, State, ?RTMP_TIMEOUT};
430
431handle_sync_event({getstat, Options}, _From, StateName, State) ->
432  {reply, get_stat(State, Options), StateName, State, ?RTMP_TIMEOUT};
433
434handle_sync_event(getstat, _From, StateName, State) ->
435  {reply, get_stat(State), StateName, State, ?RTMP_TIMEOUT};
436
437handle_sync_event(get_socket, _From, StateName, #rtmp_socket{socket = Socket} = State) ->
438  {reply, Socket, StateName, State, ?RTMP_TIMEOUT};
439
440handle_sync_event({setopts, Options}, _From, StateName, State) ->
441  NewState = set_options(State, Options),
442  {reply, ok, StateName, NewState, ?RTMP_TIMEOUT};
443
444handle_sync_event(Event, _From, StateName, StateData) ->
445  {stop, {StateName, undefined_event, Event}, StateData}.
446
447
448%%-------------------------------------------------------------------------
449%% Func: handle_info/3
450%% Returns: {next_state, NextStateName, NextStateData}          |
451%%          {next_state, NextStateName, NextStateData, Timeout} |
452%%          {stop, Reason, NewStateData}
453%% @private
454%%-------------------------------------------------------------------------
455handle_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 ->
456  activate_socket(Socket),
457  {next_state, handshake_c1, State#rtmp_socket{buffer = <<Buffer/binary, Data/binary>>, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
458  
459handle_info({tcp, Socket, Data}, handshake_c1, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead} = State) ->
460  activate_socket(Socket),
461  <<Handshake:(?HS_BODY_LEN+1)/binary, Rest/binary>> = <<Buffer/binary, Data/binary>>,
462  State1 = case rtmp_handshake:server(Handshake) of
463    {uncrypted, Reply} -> 
464  	  send_data(State, Reply);
465    {crypted, Reply, KeyIn, KeyOut} ->
466  	  NewState = send_data(State, Reply),
467  	  ?D({"Established RTMPE connection"}),
468      NewState#rtmp_socket{key_in = KeyIn, key_out = KeyOut}
469  end,
470	{next_state, 'handshake_c3', State1#rtmp_socket{buffer = Rest, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
471
472
473handle_info({tcp, Socket, Data}, handshake_c3, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead} = State) when size(Buffer) + size(Data) < ?HS_BODY_LEN ->
474  activate_socket(Socket),
475  {next_state, handshake_c3, State#rtmp_socket{buffer = <<Buffer/binary, Data/binary>>, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
476  
477handle_info({tcp, Socket, Data}, handshake_c3, #rtmp_socket{socket=Socket, consumer = Consumer, buffer = Buffer, bytes_read = BytesRead, active = Active, key_in = KeyIn} = State) ->
478  <<_HandShakeC3:?HS_BODY_LEN/binary, CryptedData/binary>> = <<Buffer/binary, Data/binary>>,
479  Consumer ! {rtmp, self(), connected},
480  case Active of
481    true -> inet:setopts(Socket, [{active, true}]);
482    _ -> ok
483  end,
484
485  {NewKeyIn, Rest} = case KeyIn of
486    undefined -> {undefined, CryptedData};
487    _ -> rtmpe:crypt(KeyIn, CryptedData)
488  end,
489  % ?D({decrypt, Rest, CryptedData == element(2, rtmpe:crypt(NewKeyIn, Rest))}),
490  
491  State1 = State#rtmp_socket{bytes_read = BytesRead + size(Data), key_in = NewKeyIn, buffer = Rest},
492  case Rest of
493    <<>> -> {next_state, loop, State1, ?RTMP_TIMEOUT};
494    _ -> {next_state, loop, handle_rtmp_data(State1), ?RTMP_TIMEOUT}
495  end;
496
497handle_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 ->
498  activate_socket(Socket),
499  {next_state, handshake_s1, State#rtmp_socket{buffer = <<Buffer/binary, Data/binary>>, bytes_read = BytesRead + size(Data)}, ?RTMP_TIMEOUT};
500
501handle_info({tcp, Socket, Data}, handshake_s1, #rtmp_socket{socket=Socket, consumer = Consumer, buffer = Buffer, bytes_read = BytesRead, active = Active} = State) ->
502  <<?HS_UNCRYPTED, S1:?HS_BODY_LEN/binary, _S2:?HS_BODY_LEN/binary, Rest/binary>> = <<Buffer/binary, Data/binary>>,
503  send_data(State, rtmp_handshake:c2(S1)),
504  Consumer ! {rtmp, self(), connected},
505  case Active of
506    true -> inet:setopts(Socket, [{active, true}]);
507    _ -> ok
508  end,
509  
510  {next_state, loop, handle_rtmp_data(State#rtmp_socket{bytes_read = BytesRead + size(Data), buffer = Rest}), ?RTMP_TIMEOUT};
511
512
513
514handle_info({tcp, Socket, CryptedData}, loop, #rtmp_socket{socket=Socket, buffer = Buffer, bytes_read = BytesRead, bytes_unack = BytesUnack, key_in = KeyIn} = State) ->
515  State1 = flush_send(State),
516  {NewKeyIn, Data} = case KeyIn of
517    undefined -> {undefined, CryptedData};
518    _ -> rtmpe:crypt(KeyIn, CryptedData)
519  end,
520  {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};
521
522handle_info({tcp_closed, Socket}, _StateName, #rtmp_socket{socket = Socket, consumer = Consumer} = StateData) ->
523  Consumer ! {rtmp, self(), disconnect, get_stat(StateData)},
524  {stop, normal, StateData};
525
526handle_info(#rtmp_message{} = Message, loop, State) ->
527  State1 = send_data(State, Message),
528  State2 = flush_send(State1),
529  {next_state, loop, State2, ?RTMP_TIMEOUT};
530
531handle_info({rtmpt, RTMPT, alive}, StateName, #rtmp_socket{socket = RTMPT} = State) ->
532  {next_state, StateName, State#rtmp_socket{pinged = false}, ?RTMP_TIMEOUT};
533
534handle_info({rtmpt, RTMPT, Data}, StateName, State) ->
535  handle_info({tcp, RTMPT, Data}, StateName, State);
536
537handle_info({'DOWN', _, process, _Client, _Reason}, _StateName, State) ->
538  {stop, normal, State};
539
540handle_info(_Info, StateName, StateData) ->
541  error_logger:error_msg("Unknown message to rtmp socket: ~p ~p ~p~n", [_Info, StateName, StateData]),
542  {next_state, StateName, StateData, ?RTMP_TIMEOUT}.
543
544% flush_send(State) -> flush_send([], State).
545
546flush_send(State) ->
547  receive
548    #rtmp_message{} = Message ->
549      NewState = send_data(State, Message),
550      % ?D({out, Message}),
551      % {NewState, Data} = rtmp:encode(State, Message),
552      % flush_send([Data | Packet], NewState)
553      flush_send(NewState)
554  after
555    0 -> State
556  end.
557  
558activate_socket(Socket) when is_port(Socket) ->
559  inet:setopts(Socket, [{active, once}]);
560activate_socket(Socket) when is_pid(Socket) ->
561  ok.
562  
563
564send_data(#rtmp_socket{sent_audio_notify = false} = Socket, 
565          #rtmp_message{type = audio, timestamp = DTS, stream_id = StreamId, body = Body} = Message) when size(Body) > 0 ->
566  State1 = send_data(Socket#rtmp_socket{sent_audio_notify = true}, rtmp_lib:empty_audio(StreamId, DTS)),
567  send_data(State1, Message#rtmp_message{ts_type = new});
568
569send_data(#rtmp_socket{sent_video_notify = false} = Socket, #rtmp_message{type = video, timestamp = DTS, stream_id = StreamId} = Message) ->
570  Msg = [
571    #rtmp_message{type = video, channel_id = rtmp_lib:channel_id(video, StreamId), timestamp = DTS, stream_id = StreamId, body = <<87,0>>, ts_type = new},
572    #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},
573    #rtmp_message{type = video, channel_id = rtmp_lib:channel_id(video, StreamId), timestamp = DTS, stream_id = StreamId, body = <<87,1>>, ts_type = delta}
574  ],
575  
576  State2 = lists:foldl(fun(M, State1) ->
577    send_data(State1, M)
578  end, Socket#rtmp_socket{sent_video_notify = true}, Msg),
579  send_data(State2, Message);
580  
581send_data(#rtmp_socket{sent_audio_notify = true} = Socket, #rtmp_message{type = stream_end} = Message) ->
582  send_data(Socket#rtmp_socket{sent_audio_notify = false}, Message);
583
584send_data(#rtmp_socket{sent_video_notify = true} = Socket, #rtmp_message{type = stream_end} = Message) ->
585  send_data(Socket#rtmp_socket{sent_video_notify = false}, Message);
586
587send_data(#rtmp_socket{socket = Socket, key_out = KeyOut} = State, Message) ->
588  case State#rtmp_socket.debug of
589    true -> print_rtmp_message(out, Message);
590    _ -> ok
591  end,
592  {NewState, Data} = case Message of
593    #rtmp_message{} ->
594      % ?D({Socket,Message#rtmp_message.type, Message#rtmp_message.timestamp, Message#rtmp_message.ts_type}),
595      rtmp:encode(State, Message);
596    _ -> 
597      {State, Message}
598  end,
599  {NewKeyOut, Crypt} = case KeyOut of
600    undefined -> {undefined, Data};
601    _ -> rtmpe:crypt(KeyOut, Data)
602  end,
603  % (catch rtmp_stat_collector:out_bytes(self(), iolist_size(Crypt))),
604  if
605    is_port(Socket) ->
606      gen_tcp:send(Socket, Crypt);
607    is_pid(Socket) ->
608      rtmpt:write(Socket, Crypt)
609  end,
610  NewState#rtmp_socket{key_out = NewKeyOut}.
611
612
613
614handle_rtmp_data(#rtmp_socket{buffer = <<>>} = State) ->
615  % ?D({empty_input, erlang:get_stacktrace()}),
616  State;
617
618handle_rtmp_data(#rtmp_socket{bytes_unack = Bytes, window_size = Window} = State) when Bytes >= Window ->
619  State1 = send_data(State, #rtmp_message{type = ack_read}),
620  handle_rtmp_data(State1);
621
622handle_rtmp_data(#rtmp_socket{buffer = Data} = State) ->
623  got_rtmp_message(rtmp:decode(State, Data)).
624
625
626print_rtmp_message(InOut, #rtmp_message{channel_id = Channel, ts_type = TSType, timestamp = TS, type = Type, stream_id = StreamId, body = Body}) ->
627  DecodedBody = case Type of
628    video when size(Body) > 10 -> erlang:setelement(5, flv:decode_video_tag(Body), size(Body));
629    audio when size(Body) > 0 -> erlang:setelement(7, flv:decode_audio_tag(Body), size(Body));
630    _ -> Body
631  end,
632  io:format("~p ~p ~p ~p ~p ~p ~p~n", [InOut, Channel, TSType, TS, Type, StreamId, DecodedBody]),
633  ok;
634
635print_rtmp_message(InOut, Msg) ->
636  io:format("~p ~s~n", [InOut, io_lib_pretty_limited:print(Msg, 20)]),
637  ok.
638  
639
640
641got_rtmp_message({#rtmp_socket{debug = true}, Msg, _} = Message) ->
642  print_rtmp_message(in, Msg),
643  handle_rtmp_message(Message);
644
645got_rtmp_message(Decoded) ->
646  handle_rtmp_message(Decoded).
647
648
649handle_rtmp_message({#rtmp_socket{consumer = Consumer} = State, #rtmp_message{type = window_size, body = Size} = Message, Rest}) ->
650  Consumer ! {rtmp, self(), Message},
651  rtmp_message_sent(State#rtmp_socket{window_size = Size, buffer = Rest});
652
653handle_rtmp_message({#rtmp_socket{consumer = Consumer, pinged = true} = State, #rtmp_message{type = pong} = Message, Rest}) ->
654  Consumer ! {rtmp, self(), Message},
655  rtmp_message_sent(State#rtmp_socket{pinged = false, buffer = Rest});
656
657handle_rtmp_message({#rtmp_socket{consumer = Consumer} = State, #rtmp_message{type = ping, body = Timestamp} = Message, Rest}) ->
658  Consumer ! {rtmp, self(), Message},
659  send_data(State, #rtmp_message{type = pong, body = Timestamp}),
660  rtmp_message_sent(State#rtmp_socket{buffer = Rest});
661
662handle_rtmp_message({#rtmp_socket{consumer = Consumer} = State, Message, Rest}) ->
663  Consumer ! {rtmp, self(), Message},
664  rtmp_message_sent(State#rtmp_socket{buffer = Rest});
665
666handle_rtmp_message({#rtmp_socket{socket=Socket} = State, Rest}) -> 
667  activate_socket(Socket),
668  State#rtmp_socket{buffer = Rest}.
669
670rtmp_message_sent(#rtmp_socket{active = true,socket = Socket} = State) ->
671  activate_socket(Socket),
672  handle_rtmp_data(State);
673
674rtmp_message_sent(#rtmp_socket{active = _} = State) ->
675  handle_rtmp_data(State);
676
677rtmp_message_sent(State) -> 
678  State.
679
680
681%%-------------------------------------------------------------------------
682%% Func: terminate/3
683%% Purpose: Shutdown the fsm
684%% Returns: any
685%% @private
686%%-------------------------------------------------------------------------
687terminate(_Reason, _StateName, #rtmp_socket{socket=Socket}) when is_pid(Socket) ->
688  gen_server:cast(Socket, close),
689  ok;
690  
691terminate(_Reason, _StateName, _State) ->
692  ok.
693
694
695%%-------------------------------------------------------------------------
696%% Func: code_change/4
697%% Purpose: Convert process state when code is changed
698%% Returns: {ok, NewState, NewStateData}
699%% @private
700%%-------------------------------------------------------------------------
701code_change(_OldVersion, _StateName, _State, _Extra) ->
702  ok.