From 706e10820d490e0865b73f66a8665951aa6a71e3 Mon Sep 17 00:00:00 2001 From: Stefan Metzmacher Date: Tue, 29 Nov 2011 16:25:03 +0100 Subject: s4:lib/socket: add socket_connect_multi_ex_* This add an async establish hook, that gets called before a connection success is detected. This can be used to do a NBT session request and it makes sure that we don't cancel the connection on port 445, just because the tcp connect on port 139 worked. metze --- source4/lib/socket/connect_multi.c | 129 ++++++++++++++++++++++++++++++++++--- source4/lib/socket/socket.h | 28 ++++++++ 2 files changed, 147 insertions(+), 10 deletions(-) diff --git a/source4/lib/socket/connect_multi.c b/source4/lib/socket/connect_multi.c index 2b926c8bd9..c8231b3cff 100644 --- a/source4/lib/socket/connect_multi.c +++ b/source4/lib/socket/connect_multi.c @@ -42,6 +42,8 @@ struct connect_multi_state { uint16_t result_port; int num_connects_sent, num_connects_recv; + + struct socket_connect_multi_ex *ex; }; /* @@ -59,17 +61,19 @@ static void connect_multi_timer(struct tevent_context *ev, struct timeval tv, void *p); static void connect_multi_next_socket(struct composite_context *result); static void continue_one(struct composite_context *creq); +static void continue_one_ex(struct tevent_req *subreq); /* setup an async socket_connect, with multiple ports */ -_PUBLIC_ struct composite_context *socket_connect_multi_send( +_PUBLIC_ struct composite_context *socket_connect_multi_ex_send( TALLOC_CTX *mem_ctx, const char *server_name, int num_server_ports, uint16_t *server_ports, struct resolve_context *resolve_ctx, - struct tevent_context *event_ctx) + struct tevent_context *event_ctx, + struct socket_connect_multi_ex *ex) { struct composite_context *result; struct connect_multi_state *multi; @@ -95,6 +99,8 @@ _PUBLIC_ struct composite_context *socket_connect_multi_send( multi->ports[i] = server_ports[i]; } + multi->ex = ex; + /* we don't want to do the name resolution separately for each port, so start it now, then only start on @@ -225,10 +231,61 @@ static void continue_one(struct composite_context *creq) struct connect_multi_state *multi = talloc_get_type(result->private_data, struct connect_multi_state); NTSTATUS status; - multi->num_connects_recv++; status = socket_connect_recv(creq); + if (multi->ex) { + struct tevent_req *subreq; + + subreq = multi->ex->establish_send(state, + result->event_ctx, + state->sock, + state->addr, + multi->ex->private_data); + if (composite_nomem(subreq, result)) return; + tevent_req_set_callback(subreq, continue_one_ex, state); + return; + } + + multi->num_connects_recv++; + + if (NT_STATUS_IS_OK(status)) { + multi->sock = talloc_steal(multi, state->sock); + multi->result_port = state->addr->port; + } + + talloc_free(state); + + if (NT_STATUS_IS_OK(status) || + multi->num_connects_recv == (multi->num_address * multi->num_ports)) { + result->status = status; + composite_done(result); + return; + } + + /* try the next port */ + connect_multi_next_socket(result); +} + +/* + one of our multi->ex->establish_send() calls hash finished. If it got a + connection or there are none left then we are done +*/ +static void continue_one_ex(struct tevent_req *subreq) +{ + struct connect_one_state *state = + tevent_req_callback_data(subreq, + struct connect_one_state); + struct composite_context *result = state->result; + struct connect_multi_state *multi = + talloc_get_type_abort(result->private_data, + struct connect_multi_state); + NTSTATUS status; + multi->num_connects_recv++; + + status = multi->ex->establish_recv(subreq); + TALLOC_FREE(subreq); + if (NT_STATUS_IS_OK(status)) { multi->sock = talloc_steal(multi, state->sock); multi->result_port = state->addr->port; @@ -250,7 +307,7 @@ static void continue_one(struct composite_context *creq) /* async recv routine for socket_connect_multi() */ -_PUBLIC_ NTSTATUS socket_connect_multi_recv(struct composite_context *ctx, +_PUBLIC_ NTSTATUS socket_connect_multi_ex_recv(struct composite_context *ctx, TALLOC_CTX *mem_ctx, struct socket_context **sock, uint16_t *port) @@ -267,6 +324,55 @@ _PUBLIC_ NTSTATUS socket_connect_multi_recv(struct composite_context *ctx, return status; } +NTSTATUS socket_connect_multi_ex(TALLOC_CTX *mem_ctx, + const char *server_address, + int num_server_ports, uint16_t *server_ports, + struct resolve_context *resolve_ctx, + struct tevent_context *event_ctx, + struct socket_connect_multi_ex *ex, + struct socket_context **result, + uint16_t *result_port) +{ + struct composite_context *ctx = + socket_connect_multi_ex_send(mem_ctx, server_address, + num_server_ports, server_ports, + resolve_ctx, + event_ctx, + ex); + return socket_connect_multi_ex_recv(ctx, mem_ctx, result, result_port); +} + +/* + setup an async socket_connect, with multiple ports +*/ +_PUBLIC_ struct composite_context *socket_connect_multi_send( + TALLOC_CTX *mem_ctx, + const char *server_name, + int num_server_ports, + uint16_t *server_ports, + struct resolve_context *resolve_ctx, + struct tevent_context *event_ctx) +{ + return socket_connect_multi_ex_send(mem_ctx, + server_name, + num_server_ports, + server_ports, + resolve_ctx, + event_ctx, + NULL); /* ex */ +} + +/* + async recv routine for socket_connect_multi() + */ +_PUBLIC_ NTSTATUS socket_connect_multi_recv(struct composite_context *ctx, + TALLOC_CTX *mem_ctx, + struct socket_context **sock, + uint16_t *port) +{ + return socket_connect_multi_ex_recv(ctx, mem_ctx, sock, port); +} + NTSTATUS socket_connect_multi(TALLOC_CTX *mem_ctx, const char *server_address, int num_server_ports, uint16_t *server_ports, @@ -275,10 +381,13 @@ NTSTATUS socket_connect_multi(TALLOC_CTX *mem_ctx, struct socket_context **result, uint16_t *result_port) { - struct composite_context *ctx = - socket_connect_multi_send(mem_ctx, server_address, - num_server_ports, server_ports, - resolve_ctx, - event_ctx); - return socket_connect_multi_recv(ctx, mem_ctx, result, result_port); + return socket_connect_multi_ex(mem_ctx, + server_address, + num_server_ports, + server_ports, + resolve_ctx, + event_ctx, + NULL, /* ex */ + result, + result_port); } diff --git a/source4/lib/socket/socket.h b/source4/lib/socket/socket.h index e00b61ba6c..403a723edb 100644 --- a/source4/lib/socket/socket.h +++ b/source4/lib/socket/socket.h @@ -202,6 +202,34 @@ NTSTATUS socket_connect_ev(struct socket_context *sock, uint32_t flags, struct tevent_context *ev); +struct socket_connect_multi_ex { + void *private_data; + struct tevent_req *(*establish_send)(TALLOC_CTX *mem_ctx, + struct tevent_context *ev, + struct socket_context *sock, + struct socket_address *addr, + void *private_data); + NTSTATUS (*establish_recv)(struct tevent_req *req); +}; +struct composite_context *socket_connect_multi_ex_send(TALLOC_CTX *mem_ctx, + const char *server_address, + int num_server_ports, + uint16_t *server_ports, + struct resolve_context *resolve_ctx, + struct tevent_context *event_ctx, + struct socket_connect_multi_ex *ex); +NTSTATUS socket_connect_multi_ex_recv(struct composite_context *ctx, + TALLOC_CTX *mem_ctx, + struct socket_context **result, + uint16_t *port); +NTSTATUS socket_connect_multi_ex(TALLOC_CTX *mem_ctx, const char *server_address, + int num_server_ports, uint16_t *server_ports, + struct resolve_context *resolve_ctx, + struct tevent_context *event_ctx, + struct socket_connect_multi_ex *ex, + struct socket_context **result, + uint16_t *port); + struct composite_context *socket_connect_multi_send(TALLOC_CTX *mem_ctx, const char *server_address, int num_server_ports, -- cgit