Merge pull request #228 from qzhuyan/dev/william/cache-stream-id

emqx · Nov 13, 2023 · a3a008a · a3a008a
2 parents c63fd66 + 3423ecd
commit a3a008a
Show file tree

Hide file tree

Showing 7 changed files with 51 additions and 1 deletion.
diff --git a/c_src/quicer_config.c b/c_src/quicer_config.c
@@ -781,6 +781,11 @@ getopt3(ErlNifEnv *env,
     }
   else if (enif_get_resource(env, ctx, ctx_stream_t, &q_ctx))
     {
+      if (ATOM_QUIC_PARAM_STREAM_ID == eopt
+          && ((QuicerStreamCTX *)q_ctx)->StreamID != UNSET_STREAMID)
+        {
+          return SUCCESS(ETERM_UINT_64(((QuicerStreamCTX *)q_ctx)->StreamID));
+        }
       if (!get_stream_handle(q_ctx))
         {
           goto Exit;

diff --git a/c_src/quicer_connection.c b/c_src/quicer_connection.c
@@ -1359,6 +1359,7 @@ handle_connection_event_peer_stream_started(QuicerConnCTX *c_ctx,
 
   s_ctx->owner = acc;
   s_ctx->is_closed = FALSE;
+  cache_stream_id(s_ctx);
 
   ERL_NIF_TERM props_name[] = { ATOM_FLAGS, ATOM_IS_ORPHAN };
   ERL_NIF_TERM props_value[]

diff --git a/c_src/quicer_ctx.c b/c_src/quicer_ctx.c
@@ -233,6 +233,7 @@ init_s_ctx()
     }
   CxPlatZeroMemory(s_ctx, sizeof(QuicerStreamCTX));
   s_ctx->magic = 0xefefefef; // 4025479151
+  s_ctx->StreamID = UNSET_STREAMID;
   s_ctx->env = enif_alloc_env();
   s_ctx->imm_env = enif_alloc_env();
   s_ctx->lock = enif_mutex_create("quicer:s_ctx");
@@ -372,3 +373,14 @@ get_reg_handle(QuicerRegistrationCTX *r_ctx)
 {
   return CxPlatRefIncrementNonZero(&r_ctx->ref_count, 1);
 }
+
+void
+cache_stream_id(QuicerStreamCTX *s_ctx)
+{
+  uint32_t bufferlen = sizeof(s_ctx->StreamID);
+  if (QUIC_FAILED(MsQuic->GetParam(
+          s_ctx->Stream, QUIC_PARAM_STREAM_ID, &bufferlen, &s_ctx->StreamID)))
+    {
+      s_ctx->StreamID = UNSET_STREAMID;
+    }
+}
diff --git a/c_src/quicer_ctx.h b/c_src/quicer_ctx.h
@@ -112,6 +112,7 @@ typedef struct QuicerStreamCTX
   uint32_t magic;
   QuicerConnCTX *c_ctx;
   HQUIC Stream;
+  uint64_t StreamID;
   ACCEPTOR *owner;
   ErlNifMonitor owner_mon;
   ErlNifEnv *env;
@@ -187,4 +188,6 @@ BOOLEAN get_listener_handle(QuicerListenerCTX *l_ctx);
 void put_reg_handle(QuicerRegistrationCTX *r_ctx);
 BOOLEAN get_reg_handle(QuicerRegistrationCTX *r_ctx);
 
+void cache_stream_id(QuicerStreamCTX *s_ctx);
+
 #endif // __QUICER_CTX_H_
diff --git a/c_src/quicer_stream.c b/c_src/quicer_stream.c
@@ -362,6 +362,7 @@ async_start_stream2(ErlNifEnv *env,
     }
   HQUIC Stream = s_ctx->Stream;
   Status = MsQuic->StreamStart(Stream, start_flag);
+  cache_stream_id(s_ctx);
   put_stream_handle(s_ctx);
 
   if (QUIC_FAILED(Status))
@@ -1048,7 +1049,7 @@ handle_stream_event_start_complete(QuicerStreamCTX *s_ctx,
           = { atom_status(env, Event->START_COMPLETE.Status),
               enif_make_uint64(env, Event->START_COMPLETE.ID),
               ATOM_BOOLEAN(Event->START_COMPLETE.PeerAccepted) };
-
+      cache_stream_id(s_ctx);
       report = make_event_with_props(env,
                                      ATOM_START_COMPLETE,
                                      enif_make_copy(env, s_ctx->eHandle),

diff --git a/c_src/quicer_stream.h b/c_src/quicer_stream.h
@@ -21,6 +21,8 @@ limitations under the License.
 #include "quicer_internal.h"
 #include "quicer_nif.h"
 
+#define UNSET_STREAMID 0xFFFFFFFFFFFFFFF
+
 typedef enum QUICER_SEND_FLAGS
 {
   QUICER_SEND_FLAGS_SYNC = 0x1000

diff --git a/test/quicer_SUITE.erl b/test/quicer_SUITE.erl
@@ -97,6 +97,7 @@
         , tc_strm_opt_active_badarg/1
         , tc_conn_opt_sslkeylogfile/1
         , tc_get_stream_id/1
+        , tc_get_stream_id_after_close/1
         , tc_getstat/1
         , tc_getstat_closed/1
         , tc_peername_v4/1
@@ -1013,6 +1014,29 @@ tc_get_stream_id(Config) ->
       ct:fail("listener_timeout")
   end.
 
+tc_get_stream_id_after_close(Config) ->
+  Port = select_port(),
+  Owner = self(),
+  {SPid, Ref} = spawn_monitor(fun() -> echo_server(Owner, Config, Port) end),
+  receive
+    listener_ready ->
+      {ok, Conn} = quicer:connect("localhost", Port, default_conn_opts(), 5000),
+      {ok, Stm} = quicer:start_stream(Conn, []),
+      {ok, Stm2} = quicer:start_stream(Conn, []),
+      {ok, 4} = quicer:send(Stm, <<"ping">>),
+      {ok, 4} = quicer:send(Stm2, <<"ping">>),
+      ok = quicer:close_stream(Stm),
+      {ok, 0} = quicer:get_stream_id(Stm),
+      {ok, 4} = quicer:get_stream_id(Stm2),
+      ok = quicer:close_connection(Conn),
+      SPid ! done,
+      {ok, 0} = quicer:get_stream_id(Stm),
+      {ok, 4} = quicer:get_stream_id(Stm2),
+      ensure_server_exit_normal(Ref)
+  after 5000 ->
+      ct:fail("listener_timeout")
+  end.
+
 tc_get_stream_0rtt_length(Config) ->
   Port = select_port(),
   Owner = self(),
@@ -2696,6 +2720,8 @@ ping_pong_server_stm_loop(L, Conn, Stm) ->
     {quic, peer_send_shutdown, Stm, undefined} ->
       ct:pal("closing stream"),
       quicer:close_stream(Stm),
+      ?assertNotEqual({error, closed},
+                      quicer:get_stream_id(Stm)),
       ping_pong_server_stm_loop(L, Conn, Stm);
     {quic, shutdown, Conn, ErrorCode} ->
       ct:pal("closing conn: ~p", [ErrorCode]),