r11627: give the caller much more control over the stream to packet process,

allowing it to specify the initial read size (thus preventing
over-reading) and to stop the recv process when needed. This is used
by the dcerpc socket code, which relies on not getting packets when it
isn't ready for them
(This used to be commit f869fd674ec4b148dc9a264e94d19ce79d35131d)

2 files changed, 46 insertions, 2 deletions

diff --git a/source4/lib/stream/packet.c b/source4/lib/stream/packet.c
index f367368def..7aed43a943 100644
--- a/source4/lib/stream/packet.c
+++ b/source4/lib/stream/packet.c
@@ -36,6 +36,7 @@ struct packet_context {
 	DATA_BLOB partial;
 	uint32_t initial_read_size;
 	uint32_t num_read;
+	uint32_t initial_read;
 	struct tls_context *tls;
 	struct socket_context *sock;
 	struct event_context *ev;
@@ -44,6 +45,7 @@ struct packet_context {
 	struct fd_event *fde;
 	BOOL serialise;
 	BOOL processing;
+	BOOL recv_disable;
 
 	struct send_element {
 		struct send_element *next, *prev;
@@ -134,6 +136,15 @@ void packet_set_serialise(struct packet_context *pc, struct fd_event *fde)
 	pc->fde = fde;
 }
 
+/*
+  tell the packet layer how much to read when starting a new packet
+  this ensures it doesn't overread
+*/
+void packet_set_initial_read(struct packet_context *pc, uint32_t initial_read)
+{
+	pc->initial_read = initial_read;
+}
+
 
 /*
   tell the caller we have an error
@@ -172,7 +183,9 @@ static void packet_next_event(struct event_context *ev, struct timed_event *te,
 			      struct timeval t, void *private)
 {
 	struct packet_context *pc = talloc_get_type(private, struct packet_context);
-	packet_recv(pc);
+	if (pc->num_read != 0 && pc->packet_size >= pc->num_read) {
+		packet_recv(pc);
+	}
 }
 
 /*
@@ -190,6 +203,11 @@ void packet_recv(struct packet_context *pc)
 		return;
 	}
 
+	if (pc->recv_disable) {
+		EVENT_FD_NOT_READABLE(pc->fde);
+		return;
+	}
+
 	if (pc->packet_size != 0 && pc->num_read >= pc->packet_size) {
 		goto next_partial;
 	}
@@ -198,6 +216,8 @@ void packet_recv(struct packet_context *pc)
 		/* we've already worked out how long this next packet is, so skip the
 		   socket_pending() call */
 		npending = pc->packet_size - pc->num_read;
+	} else if (pc->initial_read != 0) {
+		npending = pc->initial_read - pc->num_read;
 	} else {
 		if (pc->tls) {
 			status = tls_socket_pending(pc->tls, &npending);
@@ -306,7 +326,7 @@ next_partial:
 	if (pc->partial.length == 0) {
 		return;
 	}
-	
+
 	/* we got multiple packets in one tcp read */
 	if (pc->ev == NULL) {
 		goto next_partial;
@@ -330,6 +350,27 @@ next_partial:
 
 
 /*
+  temporarily disable receiving 
+*/
+void packet_recv_disable(struct packet_context *pc)
+{
+	EVENT_FD_NOT_READABLE(pc->fde);
+	pc->recv_disable = True;
+}
+
+/*
+  re-enable receiving 
+*/
+void packet_recv_enable(struct packet_context *pc)
+{
+	EVENT_FD_READABLE(pc->fde);
+	pc->recv_disable = False;
+	if (pc->num_read != 0 && pc->packet_size >= pc->num_read) {
+		event_add_timed(pc->ev, pc, timeval_zero(), packet_next_event, pc);
+	}
+}
+
+/*
   trigger a run of the send queue
 */
 void packet_queue_run(struct packet_context *pc)
diff --git a/source4/lib/stream/packet.h b/source4/lib/stream/packet.h
index bba8a1940f..a8db89853c 100644
--- a/source4/lib/stream/packet.h
+++ b/source4/lib/stream/packet.h
@@ -39,7 +39,10 @@ void packet_set_tls(struct packet_context *pc, struct tls_context *tls);
 void packet_set_socket(struct packet_context *pc, struct socket_context *sock);
 void packet_set_event_context(struct packet_context *pc, struct event_context *ev);
 void packet_set_serialise(struct packet_context *pc, struct fd_event *fde);
+void packet_set_initial_read(struct packet_context *pc, uint32_t initial_read);
 void packet_recv(struct packet_context *pc);
+void packet_recv_disable(struct packet_context *pc);
+void packet_recv_enable(struct packet_context *pc);
 NTSTATUS packet_send(struct packet_context *pc, DATA_BLOB blob);
 void packet_queue_run(struct packet_context *pc);