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