fix message buffering
[project/ubus.git] / ubusd.c
1 #include <sys/socket.h>
2 #include <sys/uio.h>
3 #include <signal.h>
4 #include <stdio.h>
5 #include <unistd.h>
6 #include <fcntl.h>
7
8 #include <libubox/blob.h>
9 #include <libubox/uloop.h>
10 #include <libubox/usock.h>
11 #include <libubox/list.h>
12
13 #include "ubusd.h"
14
15 static struct avl_tree clients;
16
17 static struct ubus_msg_buf *ubus_msg_unshare(struct ubus_msg_buf *ub)
18 {
19         ub = realloc(ub, sizeof(*ub) + ub->len);
20         if (!ub)
21                 return NULL;
22
23         ub->refcount = 1;
24         memcpy(ub + 1, ub->data, ub->len);
25         ub->data = (void *) (ub + 1);
26         return ub;
27 }
28
29 struct ubus_msg_buf *ubus_msg_ref(struct ubus_msg_buf *ub)
30 {
31         if (ub->refcount == ~0)
32                 return ubus_msg_unshare(ub);
33
34         ub->refcount++;
35         return ub;
36 }
37
38 struct ubus_msg_buf *ubus_msg_new(void *data, int len, bool shared)
39 {
40         struct ubus_msg_buf *ub;
41         int buflen = sizeof(*ub);
42
43         if (!shared)
44                 buflen += len;
45
46         ub = calloc(1, buflen);
47         if (!ub)
48                 return NULL;
49
50         if (shared) {
51                 ub->refcount = ~0;
52                 ub->data = data;
53         } else {
54                 ub->refcount = 1;
55                 ub->data = (void *) (ub + 1);
56                 if (data)
57                         memcpy(ub + 1, data, len);
58         }
59
60         ub->len = len;
61         return ub;
62 }
63
64 void ubus_msg_free(struct ubus_msg_buf *ub)
65 {
66         switch (ub->refcount) {
67         case 1:
68         case ~0:
69                 free(ub);
70                 break;
71         default:
72                 ub->refcount--;
73                 break;
74         }
75 }
76
77 static int ubus_msg_writev(int fd, struct ubus_msg_buf *ub, int offset)
78 {
79         struct iovec iov[2];
80
81         if (offset < sizeof(ub->hdr)) {
82                 iov[0].iov_base = ((char *) &ub->hdr) + offset;
83                 iov[0].iov_len = sizeof(ub->hdr) - offset;
84                 iov[1].iov_base = (char *) ub->data;
85                 iov[1].iov_len = ub->len;
86                 return writev(fd, iov, 2);
87         } else {
88                 offset -= sizeof(ub->hdr);
89                 return write(fd, ((char *) ub->data) + offset, ub->len - offset);
90         }
91 }
92
93 static void ubus_msg_enqueue(struct ubus_client *cl, struct ubus_msg_buf *ub)
94 {
95         if (cl->tx_queue[cl->txq_tail])
96                 return;
97
98         cl->tx_queue[cl->txq_tail] = ub;
99         cl->txq_tail = (cl->txq_tail + 1) % ARRAY_SIZE(cl->tx_queue);
100 }
101
102 /* takes the msgbuf reference */
103 void ubus_msg_send(struct ubus_client *cl, struct ubus_msg_buf *ub, bool free)
104 {
105         int written;
106
107         if (!cl->tx_queue[cl->txq_cur]) {
108                 written = ubus_msg_writev(cl->sock.fd, ub, 0);
109                 if (written >= ub->len + sizeof(ub->hdr))
110                         goto out;
111
112                 if (written < 0)
113                         written = 0;
114
115                 cl->txq_ofs = written;
116
117                 /* get an event once we can write to the socket again */
118                 uloop_fd_add(&cl->sock, ULOOP_READ | ULOOP_WRITE | ULOOP_EDGE_TRIGGER);
119         }
120         ubus_msg_enqueue(cl, ub);
121
122 out:
123         if (free)
124                 ubus_msg_free(ub);
125 }
126
127 static struct ubus_msg_buf *ubus_msg_head(struct ubus_client *cl)
128 {
129         return cl->tx_queue[cl->txq_cur];
130 }
131
132 static void ubus_msg_dequeue(struct ubus_client *cl)
133 {
134         struct ubus_msg_buf *ub = ubus_msg_head(cl);
135
136         if (!ub)
137                 return;
138
139         ubus_msg_free(ub);
140         cl->txq_ofs = 0;
141         cl->tx_queue[cl->txq_cur] = NULL;
142         cl->txq_cur = (cl->txq_cur + 1) % ARRAY_SIZE(cl->tx_queue);
143 }
144
145 static void handle_client_disconnect(struct ubus_client *cl)
146 {
147         struct ubus_object *obj;
148
149         while (!list_empty(&cl->objects)) {
150                 obj = list_first_entry(&cl->objects, struct ubus_object, list);
151                 ubusd_free_object(obj);
152         }
153
154         while (ubus_msg_head(cl))
155                 ubus_msg_dequeue(cl);
156
157         ubus_free_id(&clients, &cl->id);
158         uloop_fd_delete(&cl->sock);
159         close(cl->sock.fd);
160         free(cl);
161 }
162
163 static void client_cb(struct uloop_fd *sock, unsigned int events)
164 {
165         struct ubus_client *cl = container_of(sock, struct ubus_client, sock);
166         struct ubus_msg_buf *ub;
167
168         /* first try to tx more pending data */
169         while ((ub = ubus_msg_head(cl))) {
170                 int written;
171
172                 written = ubus_msg_writev(sock->fd, ub, cl->txq_ofs);
173                 if (written < 0) {
174                         switch(errno) {
175                         case EINTR:
176                         case EAGAIN:
177                                 break;
178                         default:
179                                 goto disconnect;
180                         }
181                         break;
182                 }
183                 if (written == 0)
184                         break;
185
186                 cl->txq_ofs += written;
187                 if (cl->txq_ofs < ub->len + sizeof(ub->hdr))
188                         break;
189
190                 ubus_msg_dequeue(cl);
191         }
192
193         /* prevent further ULOOP_WRITE events if we don't have data
194          * to send anymore */
195         if (!ubus_msg_head(cl) && (events & ULOOP_WRITE))
196                 uloop_fd_add(sock, ULOOP_READ | ULOOP_EDGE_TRIGGER);
197
198 retry:
199         if (!sock->eof && cl->pending_msg_offset < sizeof(cl->hdrbuf)) {
200                 int offset = cl->pending_msg_offset;
201                 int bytes;
202
203                 bytes = read(sock->fd, (char *)&cl->hdrbuf + offset, sizeof(cl->hdrbuf) - offset);
204                 if (bytes < 0)
205                         goto out;
206
207                 cl->pending_msg_offset += bytes;
208                 if (cl->pending_msg_offset < sizeof(cl->hdrbuf))
209                         goto out;
210
211                 if (blob_len(&cl->hdrbuf.data) + sizeof(cl->hdrbuf) > UBUS_MAX_MSGLEN)
212                         goto disconnect;
213
214                 cl->pending_msg = ubus_msg_new(NULL, blob_raw_len(&cl->hdrbuf.data), false);
215                 if (!cl->pending_msg)
216                         goto disconnect;
217
218                 memcpy(&cl->pending_msg->hdr, &cl->hdrbuf.hdr, sizeof(cl->hdrbuf.hdr));
219                 memcpy(cl->pending_msg->data, &cl->hdrbuf.data, sizeof(cl->hdrbuf.data));
220         }
221
222         ub = cl->pending_msg;
223         if (ub) {
224                 int offset = cl->pending_msg_offset - sizeof(ub->hdr);
225                 int len = blob_raw_len(ub->data) - offset;
226                 int bytes = 0;
227
228                 if (len > 0) {
229                         bytes = read(sock->fd, (char *) ub->data + offset, len);
230                         if (bytes <= 0)
231                                 goto out;
232                 }
233
234                 if (bytes < len) {
235                         cl->pending_msg_offset += bytes;
236                         goto out;
237                 }
238
239                 /* accept message */
240                 cl->pending_msg_offset = 0;
241                 cl->pending_msg = NULL;
242                 ubusd_receive_message(cl, ub);
243                 goto retry;
244         }
245
246 out:
247         if (!sock->eof || ubus_msg_head(cl))
248                 return;
249
250 disconnect:
251         handle_client_disconnect(cl);
252 }
253
254 struct ubus_client *ubusd_get_client_by_id(uint32_t id)
255 {
256         struct ubus_id *clid;
257
258         clid = ubus_find_id(&clients, id);
259         if (!clid)
260                 return NULL;
261
262         return container_of(clid, struct ubus_client, id);
263 }
264
265 static bool get_next_connection(int fd)
266 {
267         struct ubus_client *cl;
268         int client_fd;
269
270         client_fd = accept(fd, NULL, 0);
271         if (client_fd < 0) {
272                 switch (errno) {
273                 case ECONNABORTED:
274                 case EINTR:
275                         return true;
276                 default:
277                         return false;
278                 }
279         }
280
281         cl = calloc(1, sizeof(*cl));
282         cl->sock.fd = client_fd;
283
284         INIT_LIST_HEAD(&cl->objects);
285         if (!ubus_alloc_id(&clients, &cl->id, 0))
286                 goto error;
287
288         cl->sock.cb = client_cb;
289         uloop_fd_add(&cl->sock, ULOOP_READ | ULOOP_EDGE_TRIGGER);
290         if (!ubusd_send_hello(cl))
291                 goto error_free;
292
293         return true;
294
295 error_free:
296         ubus_free_id(&clients, &cl->id);
297 error:
298         close(cl->sock.fd);
299         free(cl);
300         return true;
301 }
302
303 static void server_cb(struct uloop_fd *fd, unsigned int events)
304 {
305         bool next;
306
307         do {
308                 next = get_next_connection(fd->fd);
309         } while (next);
310 }
311
312 static struct uloop_fd server_fd = {
313         .cb = server_cb,
314 };
315
316 int main(int argc, char **argv)
317 {
318         int ret = 0;
319
320         signal(SIGPIPE, SIG_IGN);
321
322         ubus_init_id_tree(&clients);
323
324         uloop_init();
325
326         unlink(UBUS_UNIX_SOCKET);
327         server_fd.fd = usock(USOCK_UNIX | USOCK_SERVER | USOCK_NONBLOCK, UBUS_UNIX_SOCKET, NULL);
328         if (server_fd.fd < 0) {
329                 perror("usock");
330                 ret = -1;
331                 goto out;
332         }
333         uloop_fd_add(&server_fd, ULOOP_READ | ULOOP_EDGE_TRIGGER);
334
335         uloop_run();
336
337 out:
338         uloop_done();
339         return ret;
340 }