From 20bb9193685b2fc0a305f464f9877ae3144fa5b3 Mon Sep 17 00:00:00 2001
From: Laszlo Toth <tothlac@gmail.com>
Date: Tue, 1 Dec 2015 13:50:29 +0100
Subject: [PATCH] (#123) Close the connection after receiving 204 responses

In case of receiving an http answer with answer code 204, shotgun will
return the result from the server, but the connection is automatically
closed after receiving this answer.
---
 src/shotgun.erl | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/shotgun.erl b/src/shotgun.erl
index a88973b..2137db0 100644
--- a/src/shotgun.erl
+++ b/src/shotgun.erl
@@ -503,12 +503,12 @@ wait_response({gun_response, _Pid, _StreamRef, fin, StatusCode, Headers},
                 gen_fsm:reply(From, {ok, Response}),
                 queue:in(Response, Responses)
         end,
-    {next_state, at_rest, State#{responses => NewResponses}, 0};
+    check_no_content_resp(StatusCode, State#{responses => NewResponses});
 wait_response({gun_response, _Pid, _StreamRef, nofin, StatusCode, Headers},
               #{from := From, stream := StreamRef, async := Async} = State) ->
     StateName =
       case lists:keyfind(<<"transfer-encoding">>, 1, Headers) of
-          {<<"transfer-encoding">>, <<"chunked">>} when Async == true->
+          {<<"transfer-encoding">>, <<"chunked">>} when Async =:= true->
               Result = {ok, StreamRef},
               gen_fsm:reply(From, Result),
               receive_chunk;
@@ -548,7 +548,7 @@ receive_data({gun_data, _Pid, _StreamRef, fin, Data},
                     body => NewData
                    }},
     gen_fsm:reply(From, Result),
-    {next_state, at_rest, State, 0};
+    check_no_content_resp(StatusCode, State);
 receive_data({gun_error, _Pid, StreamRef, _Reason},
              #{stream := StreamRef} = State) ->
     {next_state, at_rest, State, 0}.
@@ -558,11 +558,12 @@ receive_data({gun_error, _Pid, StreamRef, _Reason},
 -spec receive_chunk(term(), term()) -> term().
 receive_chunk({'DOWN', _, _, _, _Reason}, _State) ->
     error(incomplete);
-receive_chunk({gun_data, _Pid, StreamRef, IsFin, Data}, State) ->
+receive_chunk({gun_data, _Pid, StreamRef, IsFin, Data},
+              #{status_code := StatusCode} = State) ->
     NewState = manage_chunk(IsFin, StreamRef, Data, State),
     case IsFin of
         fin ->
-            {next_state, at_rest, NewState, 0};
+            check_no_content_resp(StatusCode, NewState);
         nofin ->
             {next_state, receive_chunk, NewState}
     end;
@@ -756,3 +757,12 @@ unexpected_event_warning(StateName, Event) ->
     error_logger:warning_msg( "Unexpected event in state '~p': ~p~n"
                             , [StateName, Event]
                             ).
+
+%% @private
+check_no_content_resp(StatusCode, State) ->
+    case StatusCode of
+        204 ->
+            {stop, no_content, State};
+        _ ->
+            {next_state, at_rest, State, 0}
+    end.