client.c (3466B)
1 #define _POSIX_C_SOURCE 200809L 2 3 #include <sys/socket.h> 4 5 #include <fcntl.h> 6 #include <stdarg.h> 7 #include <stdio.h> 8 #include <stdlib.h> 9 #include <string.h> 10 #include <unistd.h> 11 12 #include <ev.h> 13 #include <tls.h> 14 15 #include "client.h" 16 #include "common.h" 17 #include "task.h" 18 19 static void child_cb(EV_P_ ev_io *w, int revents) 20 { 21 struct client *c = PTR_FROM_FIELD(struct client, watcher, w); 22 23 if (c->task > TASK_READ) { 24 if (!client_flush(c)) { 25 client_close(EV_A_ c); 26 return; 27 } 28 29 if (c->buffer_used) 30 return; 31 } 32 33 tasks[c->task].update(EV_A_ c, revents); 34 } 35 36 static void child_timeout(EV_P_ ev_timer *w, int revents) 37 { 38 struct client *c = PTR_FROM_FIELD(struct client, timeout, w); 39 40 (void)revents; 41 42 client_close(EV_A_ c); 43 } 44 45 struct client *client_new(EV_P_ int fd, struct sockaddr *addr, socklen_t addrlen 46 #ifdef USE_TLS 47 , struct tls *tlsctx 48 #endif 49 ) 50 { 51 struct client *c = malloc(sizeof(*c)); 52 53 if (!c) 54 return NULL; 55 56 c->fd = fd; 57 memcpy(&c->addr, addr, (c->addrlen = addrlen)); 58 59 fcntl(c->fd, F_SETFL, O_NONBLOCK); 60 fcntl(c->fd, F_SETFD, FD_CLOEXEC); 61 62 c->buffer_used = 0; 63 c->task = TASK_READ; 64 #ifdef USE_TLS 65 c->tlsstate = UNKNOWN; 66 c->tlsctx = tlsctx; 67 #endif 68 69 ev_timer_init(&c->timeout, child_timeout, 60.0, 0); 70 ev_timer_start(EV_A_ &c->timeout); 71 72 ev_io_init(&c->watcher, child_cb, c->fd, EV_READ); 73 ev_io_start(EV_A_ &c->watcher); 74 75 return c; 76 } 77 78 bool client_printf(struct client *c, const char *fmt, ...) 79 { 80 int n = 0; 81 va_list args; 82 83 if (c->broken_client) { 84 n += snprintf(c->buffer + c->buffer_used, sizeof(c->buffer) - c->buffer_used, "+INFO: "); 85 if (c->buffer_used + n >= sizeof(c->buffer)) 86 return false; 87 } 88 89 va_start(args, fmt); 90 n += vsnprintf(c->buffer + c->buffer_used + n, sizeof(c->buffer) - c->buffer_used - n, fmt, args); 91 va_end(args); 92 93 if (n < 0 || n + c->buffer_used >= sizeof(c->buffer)) 94 return false; 95 96 c->buffer_used += n; 97 return true; 98 } 99 100 void client_error(EV_P_ struct client *c, const char *fmt, ...) 101 { 102 va_list args; 103 int n; 104 105 if (tasks[c->task].finish) 106 tasks[c->task].finish(EV_A_ c); 107 c->task = TASK_ERROR; 108 109 client_printf(c, "3"); 110 111 va_start(args, fmt); 112 n = vsnprintf(c->buffer + c->buffer_used, sizeof(c->buffer) - c->buffer_used, fmt, args); 113 va_end(args); 114 115 if (c->buffer_used + n >= sizeof(c->buffer)) 116 return; 117 118 c->buffer_used += n; 119 c->buffer_used += snprintf(c->buffer + c->buffer_used, sizeof(c->buffer) - c->buffer_used, "\t.\t.\t.\r\n.\r\n"); 120 } 121 122 bool client_eos(struct client *c) 123 { 124 const char eos[] = ".\r\n"; 125 126 if (c->buffer_used + sizeof(eos) - 1 > sizeof(c->buffer)) 127 return false; 128 129 memcpy(c->buffer + c->buffer_used, eos, sizeof(eos) - 1); 130 c->buffer_used += sizeof(eos) - 1; 131 return true; 132 } 133 134 int client_write(struct client *c, void *buffer, size_t n) 135 { 136 #ifdef USE_TLS 137 if (c->tlsstate == READY) 138 return tls_write(c->tlsctx, buffer, n); 139 #endif 140 return write(c->fd, buffer, n); 141 } 142 143 bool client_flush(struct client *c) 144 { 145 int w; 146 147 if (!c->buffer_used) 148 return true; 149 150 w = client_write(c, c->buffer, c->buffer_used); 151 152 if (w <= 0) 153 return false; 154 155 if ((size_t)w < c->buffer_used) { 156 memmove(c->buffer, c->buffer + w, c->buffer_used - w); 157 c->buffer_used -= w; 158 } else { 159 c->buffer_used = 0; 160 } 161 162 return true; 163 } 164 165 void client_close(EV_P_ struct client *c) 166 { 167 tasks[c->task].finish(EV_A_ c); 168 #ifdef USE_TLS 169 if (c->tlsstate > PLAIN) { 170 tls_close(c->tlsctx); 171 tls_free(c->tlsctx); 172 } 173 #endif 174 ev_timer_stop(EV_A_ &c->timeout); 175 ev_io_stop(EV_A_ &c->watcher); 176 close(c->fd); 177 free(c); 178 }