Current section

Files

Jump to
grpcbox src grpcbox_stream.erl
Raw

src/grpcbox_stream.erl

-module(grpcbox_stream).
-include_lib("chatterbox/include/http2.hrl").
-include("grpcbox.hrl").
-behaviour(h2_stream).
-export([send/2,
send/3,
add_headers/2,
add_trailers/2,
handle_streams/2,
handle_info/2]).
-export([init/3,
on_receive_request_headers/2,
on_send_push_promise/2,
on_receive_request_data/2,
on_request_end_stream/1]).
-export_type([t/0,
grpc_status/0,
grpc_status_message/0,
grpc_error/0,
grpc_error_response/0]).
-record(state, {handler :: pid(),
socket,
auth_fun,
ctx :: ctx:ctx(),
req_headers=[] :: list(),
input_ref :: reference(),
callback_pid :: pid(),
connection_pid :: pid(),
request_encoding :: gzip | undefined,
response_encoding :: gzip | undefined,
content_type :: proto | json,
resp_headers=[] :: list(),
resp_trailers=[] :: list(),
headers_sent=false :: boolean(),
stream_id :: stream_id(),
method :: #method{}}).
-opaque t() :: #state{}.
-type grpc_status_message() :: unicode:chardata().
-type grpc_status() :: 0..16.
-type grpc_error() :: {grpc_status(), grpc_status_message()}.
-type grpc_error_response() :: {grpc_error, grpc_error()}.
-define(THROW(Status, Message), throw({grpc_error, {Status, Message}})).
init(ConnPid, StreamId, [Socket, AuthFun]) ->
process_flag(trap_exit, true),
{ok, #state{connection_pid=ConnPid,
stream_id=StreamId,
auth_fun=AuthFun,
socket=Socket,
handler=self()}}.
on_receive_request_headers(Headers, State=#state{auth_fun=AuthFun,
socket=Socket}) ->
%% proplists:get_value(<<":method">>, Headers) =:= <<"POST">>,
case string:lexemes(proplists:get_value(<<":path">>, Headers), "/") of
[Service, Method] ->
case ets:lookup(?SERVICES_TAB, {Service, Method}) of
[M=#method{input={_, InputStreaming}}] ->
RequestEncoding = parse_options(<<"grpc-encoding">>, Headers),
ResponseEncoding = parse_options(<<"grpc-accept-encoding">>, Headers),
ContentType = parse_options(<<"content-type">>, Headers),
Metadata = lists:foldl(fun({K = <<"grpc-", _/binary>>, V}, Acc) ->
maps:put(K, V, Acc);
(_, Acc) ->
Acc
end, #{}, Headers),
Ctx = case parse_options(<<"grpc-timeout">>, Headers) of
infinity ->
ctx:with_values(Metadata);
D ->
ctx:with_deadline(ctx:with_values(Metadata), D, nanosecond)
end,
RespHeaders = [{<<":status">>, <<"200">>},
{<<"user-agent">>, <<"grpc-erlang/0.1.0">>},
{<<"content-type">>, content_type(ContentType)}
| response_encoding(ResponseEncoding)],
State1 = State#state{resp_headers=RespHeaders,
req_headers=Headers,
request_encoding=RequestEncoding,
response_encoding=ResponseEncoding,
content_type=ContentType,
ctx=Ctx,
method=M},
case authenticate(sock:peercert(Socket), AuthFun) of
{true, _Identity} ->
case InputStreaming of
true ->
Ref = make_ref(),
Pid = proc_lib:spawn_link(?MODULE, handle_streams,
[Ref, State1#state{handler=self()}]),
{ok, State1#state{input_ref=Ref,
callback_pid=Pid}};
_ ->
{ok, State1}
end;
_ ->
end_stream(?GRPC_STATUS_UNAUTHENTICATED, <<"">>, State1)
end;
_ ->
end_stream(?GRPC_STATUS_UNIMPLEMENTED, <<"Method not found on server">>, State)
end;
_ ->
end_stream(?GRPC_STATUS_UNIMPLEMENTED, <<"failed parsing path">>, State#state{resp_headers=[{<<":status">>, <<"200">>},
{<<"user-agent">>, <<"grpc-erlang/0.1.0">>}]}),
{ok, State#state{resp_headers=[{<<":status">>, <<"200">>},
{<<"user-agent">>, <<"grpc-erlang/0.1.0">>}]}}
end.
authenticate(_, undefined) ->
{true, undefined};
authenticate({ok, Cert}, Fun) ->
Fun(Cert);
authenticate(_, _) ->
false.
handle_streams(Ref, State=#state{method=#method{module=Module,
function=Function,
output={_, false}}}) ->
{ok, Response, State1} = Module:Function(Ref, State),
send(false, Response, State1);
handle_streams(Ref, State=#state{method=#method{module=Module,
function=Function,
output={_, true}}}) ->
Module:Function(Ref, State).
on_send_push_promise(_, State) ->
{ok, State}.
on_receive_request_data(Bin, State=#state{request_encoding=Encoding,
input_ref=Ref,
callback_pid=Pid,
ctx=Ctx,
method=#method{module=Module,
function=Function,
proto=Proto,
input={Input, InputStream},
output={_Output, OutputStream}}})->
try
Messages = split_frame(Bin, Encoding),
State1 = lists:foldl(fun(EncodedMessage, StateAcc) ->
try Proto:decode_msg(EncodedMessage, Input) of
Message ->
case {InputStream, OutputStream} of
{true, _} ->
Pid ! {Ref, Message},
StateAcc;
{false, true} ->
_ = proc_lib:spawn_link(?MODULE, handle_streams,
[Message, StateAcc#state{handler=self()}]),
StateAcc;
{false, false} ->
{ok, Response, Ctx1} = Module:Function(Ctx, Message),
send(false, Response, StateAcc#state{ctx=Ctx1})
end
catch
error:{gpb_error, _} ->
?THROW(?GRPC_STATUS_INTERNAL, <<"Error parsing request protocol buffer">>)
end
end, State, Messages),
{ok, State1}
catch
throw:{grpc_error, {Status, Message}} ->
end_stream(Status, Message, State);
_C:_T ->
end_stream(?GRPC_STATUS_UNKNOWN, <<>>, State)
end.
on_request_end_stream(State=#state{input_ref=Ref,
callback_pid=Pid,
method=#method{input={_Input, true},
output={_Output, false}}}) ->
Pid ! {Ref, eos},
{ok, State};
on_request_end_stream(State=#state{input_ref=Ref,
callback_pid=Pid,
method=#method{input={_Input, true},
output={_Output, true}}}) ->
Pid ! {Ref, eos},
{ok, State};
on_request_end_stream(State=#state{input_ref=_Ref,
callback_pid=_Pid,
method=#method{input={_Input, false},
output={_Output, true}}}) ->
{ok, State};
on_request_end_stream(State=#state{method=#method{output={_Output, false}}}) ->
end_stream(State),
{ok, State};
on_request_end_stream(State) ->
end_stream(State),
{ok, State}.
%% Internal
end_stream(State) ->
end_stream(?GRPC_STATUS_OK, <<>>, State).
end_stream(Status, Message, State=#state{headers_sent=false}) ->
end_stream(Status, Message, send_headers(State));
end_stream(Status, Message, #state{connection_pid=ConnPid,
stream_id=StreamId,
resp_trailers=Trailers}) ->
h2_connection:send_trailers(ConnPid, StreamId, [{<<"grpc-status">>, Status},
{<<"grpc-message">>, Message} | Trailers],
[{send_end_stream, true}]).
send_headers(State=#state{headers_sent=true}) ->
State;
send_headers(State=#state{connection_pid=ConnPid,
stream_id=StreamId,
resp_headers=Headers,
headers_sent=false}) ->
h2_connection:send_headers(ConnPid, StreamId, Headers, [{send_end_stream, false}]),
State#state{headers_sent=true}.
handle_info({add_headers, Headers}, State) ->
update_headers(Headers, State);
handle_info({add_trailers, Trailers}, State) ->
update_trailers(Trailers, State);
handle_info({send_proto, Message}, State) ->
send(false, Message, State);
handle_info({'EXIT', _, _}, State) ->
end_stream(State),
State.
add_headers(Headers, #state{handler=Pid}) ->
Pid ! {add_headers, Headers}.
add_trailers(Headers, #state{handler=Pid}) ->
Pid ! {add_trailers, Headers}.
update_headers(Headers, State=#state{resp_headers=RespHeaders}) ->
State#state{resp_headers=RespHeaders++Headers}.
update_trailers(Trailers, State=#state{resp_trailers=RespTrailers}) ->
State#state{resp_trailers=RespTrailers++Trailers}.
send(Message, #state{handler=Pid}) ->
Pid ! {send_proto, Message}.
send(End, Message, State=#state{headers_sent=false}) ->
State1 = send_headers(State),
send(End, Message, State1);
send(End, Message, State=#state{connection_pid=ConnPid,
stream_id=StreamId,
response_encoding=Encoding,
method=#method{proto=Proto,
input={_Input, _},
output={Output, _}}}) ->
BodyToSend = Proto:encode_msg(Message, Output),
OutFrame = encode_frame(Encoding, BodyToSend),
ok = h2_connection:send_body(ConnPid, StreamId, OutFrame, [{send_end_stream, End}]),
State.
encode_frame(gzip, Bin) ->
CompressedBin = zlib:gzip(Bin),
Length = byte_size(CompressedBin),
<<1, Length:32, CompressedBin/binary>>;
encode_frame(_, Bin) ->
Length = byte_size(Bin),
<<0, Length:32, Bin/binary>>.
split_frame(Frame, Encoding) ->
split_frame(Frame, Encoding, []).
split_frame(<<>>, _Encoding, Acc) ->
lists:reverse(Acc);
split_frame(<<0, Length:32, Encoded:Length/binary, Rest/binary>>, Encoding, Acc) ->
split_frame(Rest, Encoding, [Encoded | Acc]);
split_frame(<<1, Length:32, Compressed:Length/binary, Rest/binary>>, Encoding, Acc) ->
Encoded = case Encoding of
<<"gzip">> ->
try zlib:gunzip(Compressed)
catch
error:data_error ->
?THROW(?GRPC_STATUS_INTERNAL,
<<"Could not decompress but compression algorithm supported">>)
end;
_ ->
?THROW(?GRPC_STATUS_UNIMPLEMENTED,
<<"Compression mechanism used by client not supported by the server">>)
end,
split_frame(Rest, Encoding, [Encoded | Acc]).
response_encoding(gzip) ->
[{<<"grpc-encoding">>, <<"gzip">>}];
response_encoding(snappy) ->
[{<<"grpc-encoding">>, <<"snappy">>}];
response_encoding(deflate) ->
[{<<"grpc-encoding">>, <<"deflate">>}];
response_encoding(undefined) ->
[].
content_type(json) ->
<<"application/grpc+json">>;
content_type(_) ->
<<"application/grpc+proto">>.
timeout_to_duration(T, "H") ->
erlang:convert_time_unit(timer:hours(T), millisecond, nanosecond);
timeout_to_duration(T, "M") ->
erlang:convert_time_unit(timer:minutes(T), millisecond, nanosecond);
timeout_to_duration(T, "S") ->
erlang:convert_time_unit(T, second, nanosecond);
timeout_to_duration(T, "m") ->
erlang:convert_time_unit(T, millisecond, nanosecond);
timeout_to_duration(T, "u") ->
erlang:convert_time_unit(T, microsecond, nanosecond);
timeout_to_duration(T, "n") ->
timer:seconds(T).
parse_options(<<"grpc-timeout">>, Headers) ->
case proplists:get_value(<<"grpc-timeout">>, Headers, infinity) of
infinity ->
infinity;
T ->
{I, U} = string:to_integer(T),
timeout_to_duration(I, U)
end;
parse_options(<<"content-type">>, Headers) ->
case proplists:get_value(<<"content-type">>, Headers, undefined) of
undefined ->
proto;
<<"application/grpc">> ->
proto;
<<"application/grpc+proto">> ->
proto;
<<"application/grpc+json">> ->
json;
<<"application/grpc+", _>> ->
?THROW(?GRPC_STATUS_UNIMPLEMENTED, <<"unknown content-type">>)
end;
parse_options(<<"grpc-encoding">>, Headers) ->
parse_encoding(<<"grpc-encoding">>, Headers);
parse_options(<<"grpc-accept-encoding">>, Headers) ->
parse_encoding(<<"grpc-accept-encoding">>, Headers).
parse_encoding(EncodingType, Headers) ->
case proplists:get_value(EncodingType, Headers, undefined) of
undefined ->
undefined;
<<"gzip">> ->
gzip;
<<"snappy">> ->
snappy;
<<"deflate">> ->
deflate;
_ ->
?THROW(?GRPC_STATUS_UNIMPLEMENTED, <<"unknown encoding">>)
end.