1/*
2 * Copyright (c) 2013 Apple Inc.
3 * All rights reserved.
4 */
5#include <sys/queue.h>
6#include <sys/socket.h>
7#include <dispatch/dispatch.h>
8#include <SystemConfiguration/SCPrivate.h>
9
10#include "fd_exchange.h"
11
12#define FD_EXCHANGE_MAX_DATA_SIZE	4096
13
14typedef void (^FDGenericHandler)(void);
15typedef void (^FDExchangeAckHandler)(void);
16
17enum {
18	kFDExchangeMessageTypeRequest,
19	kFDExchangeMessageTypeResponse,
20	kFDExchangeMessageTypeAck,
21};
22
23typedef int FDExchangeMessageType;
24
25struct fd_msg_handler {
26	FDExchangeMessageType			msg_type;
27	int								fd_type;
28	FDGenericHandler				cb;
29	STAILQ_ENTRY(fd_msg_handler)	link;
30};
31
32struct fd_exchange {
33	dispatch_queue_t				queue;
34	dispatch_source_t				socket_source;
35	CFRunLoopRef					callback_rl;
36	STAILQ_HEAD(, fd_msg_handler)	handler_queue;
37};
38
39static void
40fd_exchange_dealloc(void *context)
41{
42	CFAllocatorDeallocate(kCFAllocatorDefault, context);
43}
44
45static void
46fd_exchange_add_handler(struct fd_exchange *exchange, int fd_type, FDExchangeMessageType msg_type, FDGenericHandler handler)
47{
48	struct fd_msg_handler *new_handler = (struct fd_msg_handler *)CFAllocatorAllocate(kCFAllocatorDefault, sizeof(*new_handler), 0);
49	new_handler->fd_type = fd_type;
50	new_handler->msg_type = msg_type;
51	new_handler->cb = Block_copy(handler);
52	STAILQ_INSERT_TAIL(&exchange->handler_queue, new_handler, link);
53}
54
55static void
56fd_exchange_remove_handler(struct fd_exchange *exchange, struct fd_msg_handler *handler)
57{
58	STAILQ_REMOVE(&exchange->handler_queue, handler, fd_msg_handler, link);
59	Block_release(handler->cb);
60	CFAllocatorDeallocate(kCFAllocatorDefault, handler);
61}
62
63static bool
64fd_exchange_send_msg(int socket, int fd_type, FDExchangeMessageType msg_type, int fd, const void *data, size_t data_size)
65{
66	struct msghdr   msg;
67	struct iovec    iov[3];
68	int				iov_idx = 0;
69
70	if (data_size > FD_EXCHANGE_MAX_DATA_SIZE) {
71		SCLog(TRUE, LOG_ERR, CFSTR("fd_exchange_send_msg: data is to large (%u)"), data_size);
72		return false;
73	}
74
75	memset(&msg, 0, sizeof(msg));
76
77	iov[iov_idx].iov_base = (char *)&msg_type;
78	iov[iov_idx].iov_len = sizeof(msg_type);
79	iov_idx++;
80
81	iov[iov_idx].iov_base = (char *)&fd_type;
82	iov[iov_idx].iov_len = sizeof(fd_type);
83	iov_idx++;
84
85	if (data != NULL && data_size > 0) {
86		iov[iov_idx].iov_base = (char *)data;
87		iov[iov_idx].iov_len = data_size;
88		iov_idx++;
89	}
90
91	msg.msg_iov = iov;
92	msg.msg_iovlen = iov_idx;
93
94	if (fd >= 0) {
95		struct cmsghdr	*control_msg;
96		uint8_t         control_buffer[sizeof(*control_msg) + sizeof(fd)];
97		uint32_t        control_msg_len = sizeof(control_buffer);
98
99		memset(control_buffer, 0, control_msg_len);
100		control_msg = (struct cmsghdr *)((void *)control_buffer);
101		control_msg->cmsg_len = control_msg_len;
102		control_msg->cmsg_level = SOL_SOCKET;
103		control_msg->cmsg_type = SCM_RIGHTS;
104		memcpy(CMSG_DATA(control_msg), (void *)&fd, sizeof(fd));
105
106		msg.msg_control = control_msg;
107		msg.msg_controllen = control_msg_len;
108	}
109
110	if (sendmsg(socket, &msg, 0) < 0) {
111		SCLog(TRUE, LOG_ERR, CFSTR("sendmsg failed while sending an fd: %s"), strerror(errno));
112		return false;
113	}
114
115	return true;
116}
117
118static void
119fd_exchange_handle_message(struct fd_exchange *exchange, int msg_type, int fd_type, int received_fd, void *ancilliary_data, size_t ancilliary_data_size)
120{
121	__block FDGenericHandler callback = NULL;
122	dispatch_sync(exchange->queue, ^{
123		struct fd_msg_handler	*handler = NULL;
124		/* Find the next handler for this file descriptor type and message type */
125		STAILQ_FOREACH(handler, &exchange->handler_queue, link) {
126			if (handler->fd_type == fd_type && handler->msg_type == msg_type) {
127				break;
128			}
129		}
130
131		if (handler != NULL) {
132			callback = Block_copy(handler->cb);
133			/* Remove response and ack handlers. Request handlers are re-used */
134			if (msg_type == kFDExchangeMessageTypeResponse || msg_type == kFDExchangeMessageTypeAck) {
135				fd_exchange_remove_handler(exchange, handler);
136			}
137		}
138	});
139
140	if (callback != NULL) {
141		if (msg_type == kFDExchangeMessageTypeRequest || msg_type == kFDExchangeMessageTypeResponse) {
142			((FDExchangeMsgHandler)callback)(received_fd, ancilliary_data, ancilliary_data_size);
143		} else if (msg_type == kFDExchangeMessageTypeAck) {
144			((FDExchangeAckHandler)callback)();
145		}
146		Block_release(callback);
147	} else if (received_fd >= 0) {
148		SCLog(TRUE, LOG_WARNING, CFSTR("No handler available for message type %d and fd %d and fd type %d"),
149		      msg_type, received_fd, fd_type);
150		/* No handler for this message, close the received file descriptor */
151		close(received_fd);
152	}
153}
154
155struct fd_exchange *
156fd_exchange_create(int exchange_socket, CFRunLoopRef cb_runloop)
157{
158	struct fd_exchange *new_exchange = CFAllocatorAllocate(kCFAllocatorDefault, sizeof(*new_exchange), 0);
159
160	memset(new_exchange, 0, sizeof(*new_exchange));
161
162	new_exchange->callback_rl = (CFRunLoopRef)CFRetain(cb_runloop);
163	STAILQ_INIT(&new_exchange->handler_queue);
164
165	new_exchange->queue = dispatch_queue_create("fd_exchange_queue", NULL);
166	dispatch_set_context(new_exchange->queue, new_exchange);
167	dispatch_set_finalizer_f(new_exchange->queue, fd_exchange_dealloc);
168
169	new_exchange->socket_source = dispatch_source_create(DISPATCH_SOURCE_TYPE_READ, exchange_socket, 0, new_exchange->queue);
170	dispatch_source_set_cancel_handler(new_exchange->socket_source, ^{
171		close(exchange_socket);
172	});
173
174	dispatch_source_set_event_handler(new_exchange->socket_source, ^{
175		struct msghdr			msg;
176		struct cmsghdr			*control;
177		uint8_t					control_buffer[sizeof(*control) + sizeof(int)];
178		unsigned int			control_len = sizeof(control_buffer);
179		int						result;
180		char					data_buffer[FD_EXCHANGE_MAX_DATA_SIZE];
181		int						fd_type;
182		FDExchangeMessageType	msg_type;
183		struct iovec        	iov[3];
184		int						iov_idx = 0;
185
186		control = (struct cmsghdr *)((void *)control_buffer);
187
188		memset(&msg, 0, sizeof(msg));
189		memset(control_buffer, 0, control_len);
190
191		iov[iov_idx].iov_base = (char *)&msg_type;
192		iov[iov_idx].iov_len = sizeof(msg_type);
193		iov_idx++;
194
195		iov[iov_idx].iov_base = (char *)&fd_type;
196		iov[iov_idx].iov_len = sizeof(fd_type);
197		iov_idx++;
198
199		iov[iov_idx].iov_base = data_buffer;
200		iov[iov_idx].iov_len = sizeof(data_buffer);
201		iov_idx++;
202
203		msg.msg_iov = iov;
204		msg.msg_iovlen = iov_idx;
205
206		msg.msg_control = control;
207		msg.msg_controllen = control_len;
208
209		result = recvmsg(exchange_socket, &msg, 0);
210		if (result > 0) {
211			if (result >= sizeof(fd_type) + sizeof(msg_type)) {
212				size_t	data_size = result - sizeof(fd_type) - sizeof(msg_type);
213				int		received_fd = -1;
214				void	*ancilliary_data = NULL;
215
216				if (data_size > 0) {
217					ancilliary_data = CFAllocatorAllocate(kCFAllocatorDefault, data_size, 0);
218					memcpy(ancilliary_data, data_buffer, data_size);
219				}
220
221				if (control->cmsg_len == control_len &&
222				    control->cmsg_level == SOL_SOCKET &&
223				    control->cmsg_type == SCM_RIGHTS)
224				{
225					/* A file descriptor was received, copy it out */
226					memcpy(&received_fd, CMSG_DATA(control), sizeof(received_fd));
227				}
228
229				dispatch_retain(new_exchange->queue);
230				CFRunLoopPerformBlock(new_exchange->callback_rl, kCFRunLoopDefaultMode, ^{
231					fd_exchange_handle_message(new_exchange, msg_type, fd_type, received_fd, ancilliary_data, data_size);
232					dispatch_release(new_exchange->queue);
233					if (ancilliary_data != NULL) {
234						CFAllocatorDeallocate(kCFAllocatorDefault, ancilliary_data);
235					}
236				});
237				CFRunLoopWakeUp(new_exchange->callback_rl);
238
239				/* If a file descriptor was received in a response, send an ack */
240				if (msg_type == kFDExchangeMessageTypeResponse && received_fd >= 0) {
241					if (!fd_exchange_send_msg((int)dispatch_source_get_handle(new_exchange->socket_source),
242					                          fd_type, kFDExchangeMessageTypeAck, -1, NULL, 0))
243					{
244						SCLog(TRUE, LOG_WARNING, CFSTR("Failed to send an ACK for fd %d of type %d"), received_fd, fd_type);
245					}
246				}
247			} else {
248				SCLog(TRUE, LOG_ERR, CFSTR("Got a fd exchange message that is too short (%d)"), result);
249			}
250		} else if (result == 0) {
251			SCLog(TRUE, LOG_ERR, CFSTR("recvmsg returned EOF"));
252			dispatch_source_cancel(new_exchange->socket_source);
253		} else {
254			SCLog(TRUE, LOG_ERR, CFSTR("recvmsg returned an error: %s"), strerror(errno));
255			dispatch_source_cancel(new_exchange->socket_source);
256		}
257	});
258
259	dispatch_resume(new_exchange->socket_source);
260
261	return new_exchange;
262}
263
264void
265fd_exchange_destroy(struct fd_exchange *exchange)
266{
267	if (exchange != NULL) {
268		dispatch_queue_t queue = exchange->queue;
269
270		dispatch_retain(queue);
271		dispatch_sync(queue, ^{
272			/* Clear out all handlers. This will prevent any more callbacks from being called */
273			while (!STAILQ_EMPTY(&exchange->handler_queue)) {
274				struct fd_msg_handler *handler = STAILQ_FIRST(&exchange->handler_queue);
275				/* Call all ACK handlers to prevent file descriptor leaks */
276				if (handler->msg_type == kFDExchangeMessageTypeAck) {
277					((FDExchangeAckHandler)handler->cb)();
278				}
279				fd_exchange_remove_handler(exchange, handler);
280			}
281			CFRelease(exchange->callback_rl);
282			exchange->callback_rl = NULL;
283
284			dispatch_source_cancel(exchange->socket_source);
285			dispatch_release(exchange->socket_source);
286			exchange->socket_source = NULL;
287
288			/* The finalizer function for the queue will take care of de-allocating the exchange */
289			dispatch_release(exchange->queue);
290		});
291		dispatch_release(queue);
292	}
293}
294
295void
296fd_exchange_set_request_handler(struct fd_exchange *exchange, int type, FDExchangeMsgHandler handler)
297{
298	dispatch_sync(exchange->queue, ^{
299		fd_exchange_add_handler(exchange, type, kFDExchangeMessageTypeRequest, (FDGenericHandler)handler);
300	});
301}
302
303bool
304fd_exchange_send_response(struct fd_exchange *exchange, int type, int fd, const void *data, size_t data_size)
305{
306	__block bool success = true;
307
308	dispatch_sync(exchange->queue, ^{
309		if (fd_exchange_send_msg((int)dispatch_source_get_handle(exchange->socket_source), type, kFDExchangeMessageTypeResponse, fd, data, data_size)) {
310			if (fd >= 0) {
311				/* If a file descriptor was sent, close it when we receive the acknowledgment */
312				fd_exchange_add_handler(exchange, type, kFDExchangeMessageTypeAck, ^{
313					close(fd);
314				});
315			}
316		} else {
317			success = false;
318		}
319	});
320
321	return success;
322}
323
324bool
325fd_exchange_send_request(struct fd_exchange *exchange, int type, int fd, const void *data, size_t data_size, FDExchangeMsgHandler handler)
326{
327	__block bool success = true;
328
329	dispatch_sync(exchange->queue, ^{
330		if (fd_exchange_send_msg((int)dispatch_source_get_handle(exchange->socket_source), type, kFDExchangeMessageTypeRequest, fd, data, data_size))
331		{
332			/* The request was sent successfully, set up the response handler if one was specified */
333			if (handler != NULL) {
334				fd_exchange_add_handler(exchange, type, kFDExchangeMessageTypeResponse, (FDGenericHandler)handler);
335			}
336		} else {
337			success = false;
338		}
339	});
340
341	return success;
342}
343
344