libsocket: Add "mem_ctx" to socket_create()
[samba.git] / source4 / lib / socket / connect_multi.c
1 /* 
2    Unix SMB/CIFS implementation.
3
4    Fire connect requests to a host and a number of ports, with a timeout
5    between the connect request. Return if the first connect comes back
6    successfully or return the last error.
7
8    Copyright (C) Volker Lendecke 2005
9    
10    This program is free software; you can redistribute it and/or modify
11    it under the terms of the GNU General Public License as published by
12    the Free Software Foundation; either version 3 of the License, or
13    (at your option) any later version.
14    
15    This program is distributed in the hope that it will be useful,
16    but WITHOUT ANY WARRANTY; without even the implied warranty of
17    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
18    GNU General Public License for more details.
19    
20    You should have received a copy of the GNU General Public License
21    along with this program.  If not, see <http://www.gnu.org/licenses/>.
22 */
23
24 #include "includes.h"
25 #include "lib/socket/socket.h"
26 #include "lib/events/events.h"
27 #include "libcli/composite/composite.h"
28 #include "libcli/resolve/resolve.h"
29
30 #define MULTI_PORT_DELAY 2000 /* microseconds */
31
32 /*
33   overall state
34 */
35 struct connect_multi_state {
36         struct socket_address **server_address;
37         unsigned num_address, current_address, current_port;
38         int num_ports;
39         uint16_t *ports;
40
41         struct socket_context *sock;
42         uint16_t result_port;
43
44         int num_connects_sent, num_connects_recv;
45
46         struct socket_connect_multi_ex *ex;
47 };
48
49 /*
50   state of an individual socket_connect_send() call
51 */
52 struct connect_one_state {
53         struct composite_context *result;
54         struct socket_context *sock;
55         struct socket_address *addr;
56 };
57
58 static void continue_resolve_name(struct composite_context *creq);
59 static void connect_multi_timer(struct tevent_context *ev,
60                                     struct tevent_timer *te,
61                                     struct timeval tv, void *p);
62 static void connect_multi_next_socket(struct composite_context *result);
63 static void continue_one(struct composite_context *creq);
64 static void continue_one_ex(struct tevent_req *subreq);
65
66 /*
67   setup an async socket_connect, with multiple ports
68 */
69 _PUBLIC_ struct composite_context *socket_connect_multi_ex_send(
70                                                     TALLOC_CTX *mem_ctx,
71                                                     const char *server_name,
72                                                     int num_server_ports,
73                                                     uint16_t *server_ports,
74                                                     struct resolve_context *resolve_ctx,
75                                                     struct tevent_context *event_ctx,
76                                                     struct socket_connect_multi_ex *ex)
77 {
78         struct composite_context *result;
79         struct connect_multi_state *multi;
80         int i;
81
82         struct nbt_name name;
83         struct composite_context *creq;
84                 
85         result = talloc_zero(mem_ctx, struct composite_context);
86         if (result == NULL) return NULL;
87         result->state = COMPOSITE_STATE_IN_PROGRESS;
88         result->event_ctx = event_ctx;
89
90         multi = talloc_zero(result, struct connect_multi_state);
91         if (composite_nomem(multi, result)) goto failed;
92         result->private_data = multi;
93
94         multi->num_ports = num_server_ports;
95         multi->ports = talloc_array(multi, uint16_t, multi->num_ports);
96         if (composite_nomem(multi->ports, result)) goto failed;
97
98         for (i=0; i<multi->num_ports; i++) {
99                 multi->ports[i] = server_ports[i];
100         }
101
102         multi->ex = ex;
103
104         /*  
105             we don't want to do the name resolution separately
106                     for each port, so start it now, then only start on
107                     the real sockets once we have an IP
108         */
109         make_nbt_name_server(&name, server_name);
110
111         creq = resolve_name_all_send(resolve_ctx, multi, 0, multi->ports[0], &name, result->event_ctx);
112         if (composite_nomem(creq, result)) goto failed;
113
114         composite_continue(result, creq, continue_resolve_name, result);
115
116         return result;
117
118
119  failed:
120         composite_error(result, result->status);
121         return result;
122 }
123
124 /*
125   start connecting to the next socket/port in the list
126 */
127 static void connect_multi_next_socket(struct composite_context *result)
128 {
129         struct connect_multi_state *multi = talloc_get_type(result->private_data, 
130                                                             struct connect_multi_state);
131         struct connect_one_state *state;
132         struct composite_context *creq;
133         int next = multi->num_connects_sent;
134
135         if (next == multi->num_address * multi->num_ports) {
136                 /* don't do anything, just wait for the existing ones to finish */
137                 return;
138         }
139
140         if (multi->current_address == multi->num_address) {
141                 multi->current_address = 0;
142                 multi->current_port += 1;
143         }
144         multi->num_connects_sent += 1;
145
146         if (multi->server_address == NULL || multi->server_address[multi->current_address] == NULL) {
147                 composite_error(result, NT_STATUS_OBJECT_NAME_NOT_FOUND);
148                 return;
149         }
150
151         state = talloc(multi, struct connect_one_state);
152         if (composite_nomem(state, result)) return;
153
154         state->result = result;
155         result->status = socket_create(
156                 state, multi->server_address[multi->current_address]->family,
157                 SOCKET_TYPE_STREAM, &state->sock, 0);
158         if (!composite_is_ok(result)) return;
159
160         state->addr = socket_address_copy(state, multi->server_address[multi->current_address]);
161         if (composite_nomem(state->addr, result)) return;
162
163         socket_address_set_port(state->addr, multi->ports[multi->current_port]);
164
165         creq = socket_connect_send(state->sock, NULL, 
166                                    state->addr, 0,
167                                    result->event_ctx);
168         if (composite_nomem(creq, result)) return;
169         talloc_steal(state, creq);
170
171         multi->current_address++;
172         composite_continue(result, creq, continue_one, state);
173
174         /* if there are more ports / addresses to go then setup a timer to fire when we have waited
175            for a couple of milli-seconds, when that goes off we try the next port regardless
176            of whether this port has completed */
177         if (multi->num_ports * multi->num_address > multi->num_connects_sent) {
178                 /* note that this timer is a child of the single
179                    connect attempt state, so it will go away when this
180                    request completes */
181                 tevent_add_timer(result->event_ctx, state,
182                                 timeval_current_ofs_usec(MULTI_PORT_DELAY),
183                                 connect_multi_timer, result);
184         }
185 }
186
187 /*
188   a timer has gone off telling us that we should try the next port
189 */
190 static void connect_multi_timer(struct tevent_context *ev,
191                                 struct tevent_timer *te,
192                                 struct timeval tv, void *p)
193 {
194         struct composite_context *result = talloc_get_type(p, struct composite_context);
195         connect_multi_next_socket(result);
196 }
197
198
199 /*
200   recv name resolution reply then send the next connect
201 */
202 static void continue_resolve_name(struct composite_context *creq)
203 {
204         struct composite_context *result = talloc_get_type(creq->async.private_data, 
205                                                            struct composite_context);
206         struct connect_multi_state *multi = talloc_get_type(result->private_data, 
207                                                             struct connect_multi_state);
208         struct socket_address **addr;
209         unsigned i;
210
211         result->status = resolve_name_all_recv(creq, multi, &addr, NULL);
212         if (!composite_is_ok(result)) return;
213
214         for(i=0; addr[i]; i++);
215         multi->num_address = i;
216         multi->server_address = talloc_steal(multi, addr);
217
218         connect_multi_next_socket(result);
219 }
220
221 /*
222   one of our socket_connect_send() calls hash finished. If it got a
223   connection or there are none left then we are done
224 */
225 static void continue_one(struct composite_context *creq)
226 {
227         struct connect_one_state *state = talloc_get_type(creq->async.private_data, 
228                                                           struct connect_one_state);
229         struct composite_context *result = state->result;
230         struct connect_multi_state *multi = talloc_get_type(result->private_data, 
231                                                             struct connect_multi_state);
232         NTSTATUS status;
233
234         status = socket_connect_recv(creq);
235
236         if (multi->ex) {
237                 struct tevent_req *subreq;
238
239                 subreq = multi->ex->establish_send(state,
240                                                    result->event_ctx,
241                                                    state->sock,
242                                                    state->addr,
243                                                    multi->ex->private_data);
244                 if (composite_nomem(subreq, result)) return;
245                 tevent_req_set_callback(subreq, continue_one_ex, state);
246                 return;
247         }
248
249         multi->num_connects_recv++;
250
251         if (NT_STATUS_IS_OK(status)) {
252                 multi->sock = talloc_steal(multi, state->sock);
253                 multi->result_port = state->addr->port;
254         }
255
256         talloc_free(state);
257
258         if (NT_STATUS_IS_OK(status) ||
259             multi->num_connects_recv == (multi->num_address * multi->num_ports)) {
260                 result->status = status;
261                 composite_done(result);
262                 return;
263         }
264
265         /* try the next port */
266         connect_multi_next_socket(result);
267 }
268
269 /*
270   one of our multi->ex->establish_send() calls hash finished. If it got a
271   connection or there are none left then we are done
272 */
273 static void continue_one_ex(struct tevent_req *subreq)
274 {
275         struct connect_one_state *state =
276                 tevent_req_callback_data(subreq,
277                 struct connect_one_state);
278         struct composite_context *result = state->result;
279         struct connect_multi_state *multi =
280                 talloc_get_type_abort(result->private_data,
281                 struct connect_multi_state);
282         NTSTATUS status;
283         multi->num_connects_recv++;
284
285         status = multi->ex->establish_recv(subreq);
286         TALLOC_FREE(subreq);
287
288         if (NT_STATUS_IS_OK(status)) {
289                 multi->sock = talloc_steal(multi, state->sock);
290                 multi->result_port = state->addr->port;
291         }
292
293         talloc_free(state);
294
295         if (NT_STATUS_IS_OK(status) ||
296             multi->num_connects_recv == (multi->num_address * multi->num_ports)) {
297                 result->status = status;
298                 composite_done(result);
299                 return;
300         }
301
302         /* try the next port */
303         connect_multi_next_socket(result);
304 }
305
306 /*
307   async recv routine for socket_connect_multi()
308  */
309 _PUBLIC_ NTSTATUS socket_connect_multi_ex_recv(struct composite_context *ctx,
310                                    TALLOC_CTX *mem_ctx,
311                                    struct socket_context **sock,
312                                    uint16_t *port)
313 {
314         NTSTATUS status = composite_wait(ctx);
315         if (NT_STATUS_IS_OK(status)) {
316                 struct connect_multi_state *multi =
317                         talloc_get_type(ctx->private_data,
318                                         struct connect_multi_state);
319                 *sock = talloc_steal(mem_ctx, multi->sock);
320                 *port = multi->result_port;
321         }
322         talloc_free(ctx);
323         return status;
324 }
325
326 NTSTATUS socket_connect_multi_ex(TALLOC_CTX *mem_ctx,
327                                  const char *server_address,
328                                  int num_server_ports, uint16_t *server_ports,
329                                  struct resolve_context *resolve_ctx,
330                                  struct tevent_context *event_ctx,
331                                  struct socket_connect_multi_ex *ex,
332                                  struct socket_context **result,
333                                  uint16_t *result_port)
334 {
335         struct composite_context *ctx =
336                 socket_connect_multi_ex_send(mem_ctx, server_address,
337                                              num_server_ports, server_ports,
338                                              resolve_ctx,
339                                              event_ctx,
340                                              ex);
341         return socket_connect_multi_ex_recv(ctx, mem_ctx, result, result_port);
342 }
343
344 /*
345   setup an async socket_connect, with multiple ports
346 */
347 _PUBLIC_ struct composite_context *socket_connect_multi_send(
348                                                     TALLOC_CTX *mem_ctx,
349                                                     const char *server_name,
350                                                     int num_server_ports,
351                                                     uint16_t *server_ports,
352                                                     struct resolve_context *resolve_ctx,
353                                                     struct tevent_context *event_ctx)
354 {
355         return socket_connect_multi_ex_send(mem_ctx,
356                                             server_name,
357                                             num_server_ports,
358                                             server_ports,
359                                             resolve_ctx,
360                                             event_ctx,
361                                             NULL); /* ex */
362 }
363
364 /*
365   async recv routine for socket_connect_multi()
366  */
367 _PUBLIC_ NTSTATUS socket_connect_multi_recv(struct composite_context *ctx,
368                                    TALLOC_CTX *mem_ctx,
369                                    struct socket_context **sock,
370                                    uint16_t *port)
371 {
372         return socket_connect_multi_ex_recv(ctx, mem_ctx, sock, port);
373 }
374
375 NTSTATUS socket_connect_multi(TALLOC_CTX *mem_ctx,
376                               const char *server_address,
377                               int num_server_ports, uint16_t *server_ports,
378                               struct resolve_context *resolve_ctx,
379                               struct tevent_context *event_ctx,
380                               struct socket_context **result,
381                               uint16_t *result_port)
382 {
383         return socket_connect_multi_ex(mem_ctx,
384                                        server_address,
385                                        num_server_ports,
386                                        server_ports,
387                                        resolve_ctx,
388                                        event_ctx,
389                                        NULL, /* ex */
390                                        result,
391                                        result_port);
392 }