xref: /openbsd-src/usr.sbin/vmd/atomicio.c (revision d12ef5f337ab466012c0d0342965e310661b355a)
1*d12ef5f3Sclaudio /* $OpenBSD: atomicio.c,v 1.2 2024/11/21 13:25:30 claudio Exp $ */
2149417b6Sreyk /*
3149417b6Sreyk  * Copyright (c) 2006 Damien Miller. All rights reserved.
4149417b6Sreyk  * Copyright (c) 2005 Anil Madhavapeddy. All rights reserved.
5149417b6Sreyk  * Copyright (c) 1995,1999 Theo de Raadt.  All rights reserved.
6149417b6Sreyk  * All rights reserved.
7149417b6Sreyk  *
8149417b6Sreyk  * Redistribution and use in source and binary forms, with or without
9149417b6Sreyk  * modification, are permitted provided that the following conditions
10149417b6Sreyk  * are met:
11149417b6Sreyk  * 1. Redistributions of source code must retain the above copyright
12149417b6Sreyk  *    notice, this list of conditions and the following disclaimer.
13149417b6Sreyk  * 2. Redistributions in binary form must reproduce the above copyright
14149417b6Sreyk  *    notice, this list of conditions and the following disclaimer in the
15149417b6Sreyk  *    documentation and/or other materials provided with the distribution.
16149417b6Sreyk  *
17149417b6Sreyk  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
18149417b6Sreyk  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
19149417b6Sreyk  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
20149417b6Sreyk  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
21149417b6Sreyk  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
22149417b6Sreyk  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
23149417b6Sreyk  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
24149417b6Sreyk  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
25149417b6Sreyk  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
26149417b6Sreyk  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27149417b6Sreyk  */
28149417b6Sreyk 
29149417b6Sreyk #include <sys/uio.h>
30149417b6Sreyk 
31149417b6Sreyk #include <errno.h>
32149417b6Sreyk #include <poll.h>
33149417b6Sreyk #include <string.h>
34149417b6Sreyk #include <unistd.h>
35149417b6Sreyk #include <limits.h>
36149417b6Sreyk 
37*d12ef5f3Sclaudio #include <sys/queue.h>
38*d12ef5f3Sclaudio #include <imsg.h>
39*d12ef5f3Sclaudio 
40149417b6Sreyk #include "atomicio.h"
41149417b6Sreyk 
42149417b6Sreyk /*
43149417b6Sreyk  * ensure all of data on socket comes through. f==read || f==vwrite
44149417b6Sreyk  */
45149417b6Sreyk size_t
46149417b6Sreyk atomicio6(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n,
47149417b6Sreyk     int (*cb)(void *, size_t), void *cb_arg)
48149417b6Sreyk {
49149417b6Sreyk 	char *s = _s;
50149417b6Sreyk 	size_t pos = 0;
51149417b6Sreyk 	ssize_t res;
52149417b6Sreyk 	struct pollfd pfd;
53149417b6Sreyk 
54149417b6Sreyk 	pfd.fd = fd;
55149417b6Sreyk 	pfd.events = f == read ? POLLIN : POLLOUT;
56149417b6Sreyk 	while (n > pos) {
57149417b6Sreyk 		res = (f) (fd, s + pos, n - pos);
58149417b6Sreyk 		switch (res) {
59149417b6Sreyk 		case -1:
60149417b6Sreyk 			if (errno == EINTR)
61149417b6Sreyk 				continue;
62149417b6Sreyk 			if (errno == EAGAIN) {
63149417b6Sreyk 				(void)poll(&pfd, 1, -1);
64149417b6Sreyk 				continue;
65149417b6Sreyk 			}
66149417b6Sreyk 			return 0;
67149417b6Sreyk 		case 0:
68149417b6Sreyk 			errno = EPIPE;
69149417b6Sreyk 			return pos;
70149417b6Sreyk 		default:
71149417b6Sreyk 			pos += (size_t)res;
72149417b6Sreyk 			if (cb != NULL && cb(cb_arg, (size_t)res) == -1) {
73149417b6Sreyk 				errno = EINTR;
74149417b6Sreyk 				return pos;
75149417b6Sreyk 			}
76149417b6Sreyk 		}
77149417b6Sreyk 	}
78149417b6Sreyk 	return pos;
79149417b6Sreyk }
80149417b6Sreyk 
81149417b6Sreyk size_t
82149417b6Sreyk atomicio(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n)
83149417b6Sreyk {
84149417b6Sreyk 	return atomicio6(f, fd, _s, n, NULL, NULL);
85149417b6Sreyk }
86149417b6Sreyk 
87149417b6Sreyk /*
88149417b6Sreyk  * ensure all of data on socket comes through. f==readv || f==writev
89149417b6Sreyk  */
90149417b6Sreyk size_t
91149417b6Sreyk atomiciov6(ssize_t (*f) (int, const struct iovec *, int), int fd,
92149417b6Sreyk     const struct iovec *_iov, int iovcnt,
93149417b6Sreyk     int (*cb)(void *, size_t), void *cb_arg)
94149417b6Sreyk {
95149417b6Sreyk 	size_t pos = 0, rem;
96149417b6Sreyk 	ssize_t res;
97149417b6Sreyk 	struct iovec iov_array[IOV_MAX], *iov = iov_array;
98149417b6Sreyk 	struct pollfd pfd;
99149417b6Sreyk 
100149417b6Sreyk 	if (iovcnt < 0 || iovcnt > IOV_MAX) {
101149417b6Sreyk 		errno = EINVAL;
102149417b6Sreyk 		return 0;
103149417b6Sreyk 	}
104149417b6Sreyk 	/* Make a copy of the iov array because we may modify it below */
105149417b6Sreyk 	memcpy(iov, _iov, (size_t)iovcnt * sizeof(*_iov));
106149417b6Sreyk 
107149417b6Sreyk 	pfd.fd = fd;
108149417b6Sreyk 	pfd.events = f == readv ? POLLIN : POLLOUT;
109149417b6Sreyk 	for (; iovcnt > 0 && iov[0].iov_len > 0;) {
110149417b6Sreyk 		res = (f) (fd, iov, iovcnt);
111149417b6Sreyk 		switch (res) {
112149417b6Sreyk 		case -1:
113149417b6Sreyk 			if (errno == EINTR)
114149417b6Sreyk 				continue;
115149417b6Sreyk 			if (errno == EAGAIN) {
116149417b6Sreyk 				(void)poll(&pfd, 1, -1);
117149417b6Sreyk 				continue;
118149417b6Sreyk 			}
119149417b6Sreyk 			return 0;
120149417b6Sreyk 		case 0:
121149417b6Sreyk 			errno = EPIPE;
122149417b6Sreyk 			return pos;
123149417b6Sreyk 		default:
124149417b6Sreyk 			rem = (size_t)res;
125149417b6Sreyk 			pos += rem;
126149417b6Sreyk 			/* skip completed iov entries */
127149417b6Sreyk 			while (iovcnt > 0 && rem >= iov[0].iov_len) {
128149417b6Sreyk 				rem -= iov[0].iov_len;
129149417b6Sreyk 				iov++;
130149417b6Sreyk 				iovcnt--;
131149417b6Sreyk 			}
132149417b6Sreyk 			/* This shouldn't happen... */
133149417b6Sreyk 			if (rem > 0 && (iovcnt <= 0 || rem > iov[0].iov_len)) {
134149417b6Sreyk 				errno = EFAULT;
135149417b6Sreyk 				return 0;
136149417b6Sreyk 			}
137149417b6Sreyk 			if (iovcnt == 0)
138149417b6Sreyk 				break;
139149417b6Sreyk 			/* update pointer in partially complete iov */
140149417b6Sreyk 			iov[0].iov_base = ((char *)iov[0].iov_base) + rem;
141149417b6Sreyk 			iov[0].iov_len -= rem;
142149417b6Sreyk 		}
143149417b6Sreyk 		if (cb != NULL && cb(cb_arg, (size_t)res) == -1) {
144149417b6Sreyk 			errno = EINTR;
145149417b6Sreyk 			return pos;
146149417b6Sreyk 		}
147149417b6Sreyk 	}
148149417b6Sreyk 	return pos;
149149417b6Sreyk }
150149417b6Sreyk 
151149417b6Sreyk size_t
152149417b6Sreyk atomiciov(ssize_t (*f) (int, const struct iovec *, int), int fd,
153149417b6Sreyk     const struct iovec *_iov, int iovcnt)
154149417b6Sreyk {
155149417b6Sreyk 	return atomiciov6(f, fd, _iov, iovcnt, NULL, NULL);
156149417b6Sreyk }
157*d12ef5f3Sclaudio 
158*d12ef5f3Sclaudio int
159*d12ef5f3Sclaudio imsgbuf_read_one(struct imsgbuf *imsgbuf, struct imsg *imsg)
160*d12ef5f3Sclaudio {
161*d12ef5f3Sclaudio 	struct pollfd pfd;
162*d12ef5f3Sclaudio 	int dopoll = 0;
163*d12ef5f3Sclaudio 
164*d12ef5f3Sclaudio 	pfd.fd = imsgbuf->fd;
165*d12ef5f3Sclaudio 	pfd.events = POLLIN;
166*d12ef5f3Sclaudio 	while (1) {
167*d12ef5f3Sclaudio 		switch (imsg_get(imsgbuf, imsg)) {
168*d12ef5f3Sclaudio 		case -1:
169*d12ef5f3Sclaudio 			return (-1);
170*d12ef5f3Sclaudio 		case 0:
171*d12ef5f3Sclaudio 			if (dopoll)
172*d12ef5f3Sclaudio 				(void)poll(&pfd, 1, -1);
173*d12ef5f3Sclaudio 			break;
174*d12ef5f3Sclaudio 		default:
175*d12ef5f3Sclaudio 			return (1);
176*d12ef5f3Sclaudio 		}
177*d12ef5f3Sclaudio 		dopoll = 1;
178*d12ef5f3Sclaudio 
179*d12ef5f3Sclaudio 		switch (imsgbuf_read(imsgbuf)) {
180*d12ef5f3Sclaudio 		case -1:
181*d12ef5f3Sclaudio 			return (-1);
182*d12ef5f3Sclaudio 		case 0:
183*d12ef5f3Sclaudio 			return (0);
184*d12ef5f3Sclaudio 		}
185*d12ef5f3Sclaudio 	}
186*d12ef5f3Sclaudio }
187