STEP01: librpc/rpc/dcerpc_connection.c use_trans

diff --git a/librpc/rpc/dcerpc_connection.c b/librpc/rpc/dcerpc_connection.c
index 02124f21d4d3c84229357e4eb9a1e5f94295aded..3ee878855b1f5cba9af9d5e997f566a3723d0225 100644 (file)
--- a/librpc/rpc/dcerpc_connection.c
+++ b/librpc/rpc/dcerpc_connection.c
@@ -49,6 +49,7 @@ struct dcerpc_connection {
 
        struct {
                struct tstream_context *stream;
+               dcerpc_connection_use_trans_fn use_trans_fn;
                struct tevent_queue *write_queue;
        } transport;
 
@@ -154,6 +155,12 @@ struct dcerpc_connection *dcerpc_connection_create(TALLOC_CTX *mem_ctx,
        return conn;
 }
 
+void dcerpc_connection_set_use_trans_fn(struct dcerpc_connection *conn,
+                                       dcerpc_connection_use_trans_fn fn)
+{
+       conn->transport.use_trans_fn = fn;
+}
+
 struct dcerpc_security *dcerpc_security_allocate(
                                        TALLOC_CTX *mem_ctx,
                                        struct dcerpc_connection *conn,
@@ -798,10 +805,14 @@ struct dcerpc_do_bind_state {
 };
 
 struct dcerpc_do_bind_out_frag {
+       struct tevent_context *ev;
+       struct dcerpc_connection *conn;
        struct tevent_req *req;
        enum dcerpc_pkt_type ptype;
        DATA_BLOB blob;
        struct iovec vector;
+       struct tevent_req *subreq_wait1;
+       struct tevent_req *subreq_wait2;
 };
 
 static void dcerpc_do_bind_cleanup(struct tevent_req *req,
@@ -927,7 +938,9 @@ static void dcerpc_do_bind_sec_next(struct tevent_req *subreq)
        }
 }
 
+static void dcerpc_do_bind_out_frag_trans_wait1(struct tevent_req *subreq);
 static void dcerpc_do_bind_out_frag_done(struct tevent_req *subreq);
+static void dcerpc_do_bind_out_frag_trans_wait2(struct tevent_req *subreq);
 
 static void dcerpc_do_bind_out_frag_next(struct tevent_req *req,
                                         void *private_data)
@@ -942,6 +955,7 @@ static void dcerpc_do_bind_out_frag_next(struct tevent_req *req,
        union dcerpc_payload u;
        struct tevent_req *subreq;
        uint32_t i;
+       bool use_trans = true;
 
        /*
         * the fragment belongs to the connection instead of the request
@@ -951,7 +965,8 @@ static void dcerpc_do_bind_out_frag_next(struct tevent_req *req,
        if (tevent_req_nomem(frag, req)) {
                return;
        }
-
+       frag->ev = state->ev;
+       frag->conn = state->conn;
        frag->req = req;
        state->out_frag = frag;
 
@@ -1048,17 +1063,52 @@ static void dcerpc_do_bind_out_frag_next(struct tevent_req *req,
                return;
        }
 
+       if (frag->ptype == DCERPC_PKT_AUTH3) {
+               use_trans = false;
+       }
+
+       if (frag->conn->transport.use_trans_fn == NULL) {
+               use_trans = false;
+       }
+
+       if (frag->conn->loop.subreq != NULL) {
+               use_trans = false;
+       }
+
+       if (frag->conn->features.concurrent_multiplex) {
+               use_trans = false;
+       }
+
+       if (tevent_queue_length(frag->conn->calls.out_queue) > 1) {
+               use_trans = false;
+       }
+
+       if (use_trans) {
+               frag->subreq_wait1 = tevent_queue_wait_send(frag,
+                                               frag->ev,
+                                               frag->conn->transport.write_queue);
+               if (tevent_req_nomem(req, frag->subreq_wait1)) {
+                       return;
+               }
+               tevent_req_set_callback(frag->subreq_wait1,
+                                       dcerpc_do_bind_out_frag_trans_wait1,
+                                       frag);
+               /*
+                * we need to block reads until our write is
+                * the next in the write queue.
+                */
+               state->conn->loop.subreq = frag->subreq_wait1;
+       }
+
        /*
-        * TODO: add smb_trans handling for ncacn_np
-        *
         * We need to add a dcerpc_write_fragment_queue_send/recv()
         */
 
        frag->vector.iov_base = frag->blob.data;
        frag->vector.iov_len = frag->blob.length;
-       subreq = tstream_writev_queue_send(frag, state->ev,
-                                          state->conn->transport.stream,
-                                          state->conn->transport.write_queue,
+       subreq = tstream_writev_queue_send(frag, frag->ev,
+                                          frag->conn->transport.stream,
+                                          frag->conn->transport.write_queue,
                                           &frag->vector, 1);
        if (tevent_req_nomem(subreq, req)) {
                return;
@@ -1067,12 +1117,95 @@ static void dcerpc_do_bind_out_frag_next(struct tevent_req *req,
                                dcerpc_do_bind_out_frag_done,
                                frag);
 
-       status = dcerpc_connection_loop_restart(state->conn, state->ev);
+       if (use_trans) {
+               frag->subreq_wait2 = tevent_queue_wait_send(frag,
+                                               frag->ev,
+                                               frag->conn->transport.write_queue);
+               if (tevent_req_nomem(req, frag->subreq_wait2)) {
+                       return;
+               }
+               tevent_req_set_callback(frag->subreq_wait2,
+                                       dcerpc_do_bind_out_frag_trans_wait2,
+                                       frag);
+       }
+
+       status = dcerpc_connection_loop_restart(frag->conn, frag->ev);
        if (tevent_req_nterror(req, status)) {
                return;
        }
 }
 
+static void dcerpc_do_bind_out_frag_trans_wait1(struct tevent_req *subreq)
+{
+       struct dcerpc_do_bind_out_frag *frag =
+               tevent_req_callback_data(subreq,
+               struct dcerpc_do_bind_out_frag);
+       struct tevent_req *req = frag->req;
+       NTSTATUS status;
+       bool ok;
+
+       /*
+        * TODO; what if the caller has been free'ed?
+        */
+
+       frag->subreq_wait1 = NULL;
+       frag->conn->loop.subreq = NULL;
+
+       ok = tevent_queue_wait_recv(subreq);
+       if (!ok) {
+               status = NT_STATUS_INTERNAL_ERROR;
+               TALLOC_FREE(frag);
+               if (req) {
+                       tevent_req_nterror(req, status);
+               }
+               //dcerpc_transport_dead(p, NT_STATUS_NO_MEMORY);
+               return;
+       }
+
+       if (tevent_queue_length(frag->conn->transport.write_queue) > 3) {
+               /*
+                * We added 3 entries into the queue,
+                * wait1, writev and wait2.
+                *
+                * There's more to write, we should not block
+                * further writev calls for a trans call.
+                *
+                * The wait2 stage will trigger the read.
+                */
+               TALLOC_FREE(subreq);
+               return;
+       }
+
+       /*
+        * we don't need wait2 anymore, we're sure that
+        * we'll do a trans call.
+        */
+       TALLOC_FREE(frag->subreq_wait2);
+
+       status = frag->conn->transport.use_trans_fn(frag->conn->transport.stream);
+       if (!NT_STATUS_IS_OK(status)) {
+               TALLOC_FREE(frag);
+               if (req) {
+                       tevent_req_nterror(req, status);
+               }
+               //dcerpc_transport_dead(p, NT_STATUS_NO_MEMORY);
+               return;
+       }
+
+       /* we free subreq after tstream_cli_np_use_trans */
+       TALLOC_FREE(subreq);
+
+       status = dcerpc_connection_loop_restart(frag->conn, frag->ev);
+       if (!NT_STATUS_IS_OK(status)) {
+               TALLOC_FREE(frag);
+               if (req) {
+                       tevent_req_nterror(req, status);
+               }
+               //dcerpc_transport_dead(p, NT_STATUS_NO_MEMORY);
+               return;
+       }
+}
+
 static void dcerpc_do_bind_out_frag_done(struct tevent_req *subreq)
 {
        struct dcerpc_do_bind_out_frag *frag =
@@ -1114,6 +1247,50 @@ static void dcerpc_do_bind_out_frag_done(struct tevent_req *subreq)
                tevent_req_done(req);
                return;
        }
+
+       if (frag->subreq_wait2 != NULL) {
+               return;
+       }
+
+       TALLOC_FREE(frag);
+
+       /* we need to wait for incoming pdus */
+}
+
+static void dcerpc_do_bind_out_frag_trans_wait2(struct tevent_req *subreq)
+{
+       struct dcerpc_do_bind_out_frag *frag =
+               tevent_req_callback_data(subreq,
+               struct dcerpc_do_bind_out_frag);
+       struct tevent_req *req = frag->req;
+       NTSTATUS status;
+       bool ok;
+
+       frag->subreq_wait2 = NULL;
+
+       ok = tevent_queue_wait_recv(subreq);
+       if (!ok) {
+               status = NT_STATUS_INTERNAL_ERROR;
+               TALLOC_FREE(frag);
+               if (req) {
+                       tevent_req_nterror(req, status);
+               }
+               //dcerpc_transport_dead(p, NT_STATUS_NO_MEMORY);
+               return;
+       }
+
+       TALLOC_FREE(subreq);
+
+       status = dcerpc_connection_loop_restart(frag->conn, frag->ev);
+       if (!NT_STATUS_IS_OK(status)) {
+               TALLOC_FREE(frag);
+               if (req) {
+                       tevent_req_nterror(req, status);
+               }
+               //dcerpc_transport_dead(p, NT_STATUS_NO_MEMORY);
+               return;
+       }
+
        TALLOC_FREE(frag);
 
        /* we need to wait for incoming pdus */

diff --git a/librpc/rpc/dcerpc_connection.h b/librpc/rpc/dcerpc_connection.h
index 555f3e132d167dc2ed6b23892b258308fbe01f8b..1e8b096b8959a97943570df5d24b48b71ed1e633 100644 (file)
--- a/librpc/rpc/dcerpc_connection.h
+++ b/librpc/rpc/dcerpc_connection.h
@@ -37,6 +37,9 @@ struct dcerpc_association *dcerpc_association_create(TALLOC_CTX *mem_ctx,
 struct dcerpc_connection *dcerpc_connection_create(TALLOC_CTX *mem_ctx,
                                        struct dcerpc_association *assoc,
                                        struct tstream_context **stream);
+typedef NTSTATUS (*dcerpc_connection_use_trans_fn)(struct tstream_context *stream);
+void dcerpc_connection_set_use_trans_fn(struct dcerpc_connection *conn,
+                                       dcerpc_connection_use_trans_fn fn);
 struct dcerpc_security *dcerpc_security_allocate(
                                        TALLOC_CTX *mem_ctx,
                                        struct dcerpc_connection *conn,