xref: /openbsd-src/usr.sbin/vmd/atomicio.c (revision d12ef5f337ab466012c0d0342965e310661b355a)
1 /* $OpenBSD: atomicio.c,v 1.2 2024/11/21 13:25:30 claudio Exp $ */
2 /*
3  * Copyright (c) 2006 Damien Miller. All rights reserved.
4  * Copyright (c) 2005 Anil Madhavapeddy. All rights reserved.
5  * Copyright (c) 1995,1999 Theo de Raadt.  All rights reserved.
6  * All rights reserved.
7  *
8  * Redistribution and use in source and binary forms, with or without
9  * modification, are permitted provided that the following conditions
10  * are met:
11  * 1. Redistributions of source code must retain the above copyright
12  *    notice, this list of conditions and the following disclaimer.
13  * 2. Redistributions in binary form must reproduce the above copyright
14  *    notice, this list of conditions and the following disclaimer in the
15  *    documentation and/or other materials provided with the distribution.
16  *
17  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
18  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
19  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
20  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
21  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
22  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
23  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
24  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
25  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
26  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
27  */
28 
29 #include <sys/uio.h>
30 
31 #include <errno.h>
32 #include <poll.h>
33 #include <string.h>
34 #include <unistd.h>
35 #include <limits.h>
36 
37 #include <sys/queue.h>
38 #include <imsg.h>
39 
40 #include "atomicio.h"
41 
42 /*
43  * ensure all of data on socket comes through. f==read || f==vwrite
44  */
45 size_t
46 atomicio6(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n,
47     int (*cb)(void *, size_t), void *cb_arg)
48 {
49 	char *s = _s;
50 	size_t pos = 0;
51 	ssize_t res;
52 	struct pollfd pfd;
53 
54 	pfd.fd = fd;
55 	pfd.events = f == read ? POLLIN : POLLOUT;
56 	while (n > pos) {
57 		res = (f) (fd, s + pos, n - pos);
58 		switch (res) {
59 		case -1:
60 			if (errno == EINTR)
61 				continue;
62 			if (errno == EAGAIN) {
63 				(void)poll(&pfd, 1, -1);
64 				continue;
65 			}
66 			return 0;
67 		case 0:
68 			errno = EPIPE;
69 			return pos;
70 		default:
71 			pos += (size_t)res;
72 			if (cb != NULL && cb(cb_arg, (size_t)res) == -1) {
73 				errno = EINTR;
74 				return pos;
75 			}
76 		}
77 	}
78 	return pos;
79 }
80 
81 size_t
82 atomicio(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n)
83 {
84 	return atomicio6(f, fd, _s, n, NULL, NULL);
85 }
86 
87 /*
88  * ensure all of data on socket comes through. f==readv || f==writev
89  */
90 size_t
91 atomiciov6(ssize_t (*f) (int, const struct iovec *, int), int fd,
92     const struct iovec *_iov, int iovcnt,
93     int (*cb)(void *, size_t), void *cb_arg)
94 {
95 	size_t pos = 0, rem;
96 	ssize_t res;
97 	struct iovec iov_array[IOV_MAX], *iov = iov_array;
98 	struct pollfd pfd;
99 
100 	if (iovcnt < 0 || iovcnt > IOV_MAX) {
101 		errno = EINVAL;
102 		return 0;
103 	}
104 	/* Make a copy of the iov array because we may modify it below */
105 	memcpy(iov, _iov, (size_t)iovcnt * sizeof(*_iov));
106 
107 	pfd.fd = fd;
108 	pfd.events = f == readv ? POLLIN : POLLOUT;
109 	for (; iovcnt > 0 && iov[0].iov_len > 0;) {
110 		res = (f) (fd, iov, iovcnt);
111 		switch (res) {
112 		case -1:
113 			if (errno == EINTR)
114 				continue;
115 			if (errno == EAGAIN) {
116 				(void)poll(&pfd, 1, -1);
117 				continue;
118 			}
119 			return 0;
120 		case 0:
121 			errno = EPIPE;
122 			return pos;
123 		default:
124 			rem = (size_t)res;
125 			pos += rem;
126 			/* skip completed iov entries */
127 			while (iovcnt > 0 && rem >= iov[0].iov_len) {
128 				rem -= iov[0].iov_len;
129 				iov++;
130 				iovcnt--;
131 			}
132 			/* This shouldn't happen... */
133 			if (rem > 0 && (iovcnt <= 0 || rem > iov[0].iov_len)) {
134 				errno = EFAULT;
135 				return 0;
136 			}
137 			if (iovcnt == 0)
138 				break;
139 			/* update pointer in partially complete iov */
140 			iov[0].iov_base = ((char *)iov[0].iov_base) + rem;
141 			iov[0].iov_len -= rem;
142 		}
143 		if (cb != NULL && cb(cb_arg, (size_t)res) == -1) {
144 			errno = EINTR;
145 			return pos;
146 		}
147 	}
148 	return pos;
149 }
150 
151 size_t
152 atomiciov(ssize_t (*f) (int, const struct iovec *, int), int fd,
153     const struct iovec *_iov, int iovcnt)
154 {
155 	return atomiciov6(f, fd, _iov, iovcnt, NULL, NULL);
156 }
157 
158 int
159 imsgbuf_read_one(struct imsgbuf *imsgbuf, struct imsg *imsg)
160 {
161 	struct pollfd pfd;
162 	int dopoll = 0;
163 
164 	pfd.fd = imsgbuf->fd;
165 	pfd.events = POLLIN;
166 	while (1) {
167 		switch (imsg_get(imsgbuf, imsg)) {
168 		case -1:
169 			return (-1);
170 		case 0:
171 			if (dopoll)
172 				(void)poll(&pfd, 1, -1);
173 			break;
174 		default:
175 			return (1);
176 		}
177 		dopoll = 1;
178 
179 		switch (imsgbuf_read(imsgbuf)) {
180 		case -1:
181 			return (-1);
182 		case 0:
183 			return (0);
184 		}
185 	}
186 }
187