ustd/ustd.c

   1 /* Copyright (C) 2009  Pierre-Marc Fournier
   2  *               2010  Alexis Halle
   3  *
   4  * This library is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * This library is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with this library; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
  17  */
  18
  19 #define _GNU_SOURCE
  20
  21 #include <sys/types.h>
  22 #include <sys/stat.h>
  23 #include <sys/shm.h>
  24 #include <fcntl.h>
  25 #include <unistd.h>
  26 #include <signal.h>
  27
  28 #include <stdlib.h>
  29 #include <stdio.h>
  30 #include <string.h>
  31 #include <errno.h>
  32 #include <assert.h>
  33 #include <getopt.h>
  34
  35 #include "libustd.h"
  36 #include "usterr.h"
  37
  38 char *sock_path=NULL;
  39 char *trace_path=NULL;
  40 int daemon_mode = 0;
  41 char *pidfile = NULL;
  42
  43 struct libustd_instance *instance;
  44
  45 struct buffer_info_local {
  46         /* output file */
  47         int file_fd;
  48         /* the offset we must truncate to, to unput the last subbuffer */
  49         off_t previous_offset;
  50 };
  51
  52 static int write_pidfile(const char *file_name, pid_t pid)
  53 {
  54         FILE *pidfp;
  55
  56         pidfp = fopen(file_name, "w");
  57         if(!pidfp) {
  58                 PERROR("fopen (%s)", file_name);
  59                 WARN("killing child process");
  60                 return -1;
  61         }
  62
  63         fprintf(pidfp, "%d\n", pid);
  64
  65         fclose(pidfp);
  66
  67         return 0;
  68 }
  69
  70 int create_dir_if_needed(char *dir)
  71 {
  72         int result;
  73         result = mkdir(dir, 0777);
  74         if(result == -1) {
  75                 if(errno != EEXIST) {
  76                         PERROR("mkdir");
  77                         return -1;
  78                 }
  79         }
  80
  81         return 0;
  82 }
  83
  84 int unwrite_last_subbuffer(struct buffer_info *buf)
  85 {
  86         int result;
  87         struct buffer_info_local *buf_local = buf->user_data;
  88
  89         result = ftruncate(buf_local->file_fd, buf_local->previous_offset);
  90         if(result == -1) {
  91                 PERROR("ftruncate");
  92                 return -1;
  93         }
  94
  95         result = lseek(buf_local->file_fd, buf_local->previous_offset, SEEK_SET);
  96         if(result == (int)(off_t)-1) {
  97                 PERROR("lseek");
  98                 return -1;
  99         }
 100
 101         return 0;
 102 }
 103
 104 int write_current_subbuffer(struct buffer_info *buf)
 105 {
 106         int result;
 107         struct buffer_info_local *buf_local = buf->user_data;
 108
 109         void *subbuf_mem = buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1));
 110
 111         size_t cur_sb_size = subbuffer_data_size(subbuf_mem);
 112
 113         off_t cur_offset = lseek(buf_local->file_fd, 0, SEEK_CUR);
 114         if(cur_offset == (off_t)-1) {
 115                 PERROR("lseek");
 116                 return -1;
 117         }
 118
 119         buf_local->previous_offset = cur_offset;
 120         DBG("previous_offset: %ld", cur_offset);
 121
 122         result = patient_write(buf_local->file_fd, subbuf_mem, cur_sb_size);
 123         if(result == -1) {
 124                 PERROR("write");
 125                 return -1;
 126         }
 127
 128         return 0;
 129 }
 130
 131 int on_read_subbuffer(struct libustd_callbacks *data, struct buffer_info *buf)
 132 {
 133         return write_current_subbuffer(buf);
 134 }
 135
 136 int on_read_partial_subbuffer(struct libustd_callbacks *data, struct buffer_info *buf,
 137                                 long subbuf_index, unsigned long valid_length)
 138 {
 139         struct buffer_info_local *buf_local = buf->user_data;
 140         char *tmp;
 141         int result;
 142         unsigned long pad_size;
 143
 144         result = patient_write(buf_local->file_fd, buf->mem + subbuf_index * buf->subbuf_size, valid_length);
 145         if(result == -1) {
 146                 ERR("Error writing to buffer file");
 147                 return;
 148         }
 149
 150         /* pad with empty bytes */
 151         pad_size = PAGE_ALIGN(valid_length)-valid_length;
 152         if(pad_size) {
 153                 tmp = malloc(pad_size);
 154                 memset(tmp, 0, pad_size);
 155                 result = patient_write(buf_local->file_fd, tmp, pad_size);
 156                 if(result == -1) {
 157                         ERR("Error writing to buffer file");
 158                         return;
 159                 }
 160                 free(tmp);
 161         }
 162
 163 }
 164
 165 int on_open_buffer(struct libustd_callbacks *data, struct buffer_info *buf)
 166 {
 167         char *tmp;
 168         int result;
 169         int fd;
 170         struct buffer_info_local *buf_local =
 171                 malloc(sizeof(struct buffer_info_local));
 172
 173         if(!buf_local) {
 174                 ERR("could not allocate buffer_info_local struct");
 175                 return 1;
 176         }
 177
 178         buf->user_data = buf_local;
 179
 180         /* open file for output */
 181         if(!trace_path) {
 182                 /* Only create the directory if using the default path, because
 183                  * of the risk of typo when using trace path override. We don't
 184                  * want to risk creating plenty of useless directories in that case.
 185                  */
 186                 result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
 187                 if(result == -1) {
 188                         ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
 189                         return 1;
 190                 }
 191
 192                 trace_path = USTD_DEFAULT_TRACE_PATH;
 193         }
 194
 195         asprintf(&tmp, "%s/%u_%lld", trace_path, buf->pid, buf->pidunique);
 196         result = create_dir_if_needed(tmp);
 197         if(result == -1) {
 198                 ERR("could not create directory %s", tmp);
 199                 free(tmp);
 200                 return 1;
 201         }
 202         free(tmp);
 203
 204         asprintf(&tmp, "%s/%u_%lld/%s", trace_path, buf->pid, buf->pidunique, buf->name);
 205         result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC | O_EXCL, 00600);
 206         if(result == -1) {
 207                 PERROR("open");
 208                 ERR("failed opening trace file %s", tmp);
 209                 return 1;
 210         }
 211         buf_local->file_fd = fd;
 212         free(tmp);
 213
 214         return 0;
 215 }
 216
 217 int on_close_buffer(struct libustd_callbacks *data, struct buffer_info *buf)
 218 {
 219         struct buffer_info_local *buf_local = buf->user_data;
 220         int result = close(buf_local->file_fd);
 221         free(buf_local);
 222         if(result == -1) {
 223                 PERROR("close");
 224         }
 225         return 0;
 226 }
 227
 228 int on_put_error(struct libustd_callbacks *data, struct buffer_info *buf)
 229 {
 230         unwrite_last_subbuffer(buf);
 231 }
 232
 233 struct libustd_callbacks *new_callbacks()
 234 {
 235         struct libustd_callbacks *callbacks =
 236                 malloc(sizeof(struct libustd_callbacks));
 237
 238         if(!callbacks)
 239                 return NULL;
 240
 241         callbacks->on_open_buffer = on_open_buffer;
 242         callbacks->on_close_buffer = on_close_buffer;
 243         callbacks->on_read_subbuffer = on_read_subbuffer;
 244         callbacks->on_read_partial_subbuffer = on_read_partial_subbuffer;
 245         callbacks->on_put_error = on_put_error;
 246         callbacks->on_new_thread = NULL;
 247         callbacks->on_close_thread = NULL;
 248         callbacks->on_trace_end = NULL;
 249
 250         return callbacks;
 251
 252 }
 253
 254 int is_directory(const char *dir)
 255 {
 256         int result;
 257         struct stat st;
 258
 259         result = stat(dir, &st);
 260         if(result == -1) {
 261                 PERROR("stat");
 262                 return 0;
 263         }
 264
 265         if(!S_ISDIR(st.st_mode)) {
 266                 return 0;
 267         }
 268
 269         return 1;
 270 }
 271
 272 void usage(void)
 273 {
 274         fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
 275                         "\t-h\t\tDisplay this usage.\n"
 276                         "\t-o DIR\t\tSpecify the directory where to output the traces.\n"
 277                         "\t-s PATH\t\tSpecify the path to use for the daemon socket.\n"
 278                         "\t-d\t\tStart as a daemon.\n"
 279                         "\t--pidfile FILE\tWrite the PID in this file (when using -d).\n");
 280 }
 281
 282 int parse_args(int argc, char **argv)
 283 {
 284         int c;
 285
 286         while (1) {
 287                 int option_index = 0;
 288                 static struct option long_options[] = {
 289                         {"pidfile", 1, 0, 'p'},
 290                         {"help", 0, 0, 'h'},
 291                         {"version", 0, 0, 'V'},
 292                         {0, 0, 0, 0}
 293                 };
 294
 295                 c = getopt_long(argc, argv, "hs:o:d", long_options, &option_index);
 296                 if (c == -1)
 297                         break;
 298
 299                 switch (c) {
 300                 case 0:
 301                         printf("option %s", long_options[option_index].name);
 302                         if (optarg)
 303                                 printf(" with arg %s", optarg);
 304                         printf("\n");
 305                         break;
 306                 case 's':
 307                         sock_path = optarg;
 308                         break;
 309                 case 'o':
 310                         trace_path = optarg;
 311                         if(!is_directory(trace_path)) {
 312                                 ERR("Not a valid directory. (%s)", trace_path);
 313                                 return -1;
 314                         }
 315                         break;
 316                 case 'd':
 317                         daemon_mode = 1;
 318                         break;
 319                 case 'p':
 320                         pidfile = strdup(optarg);
 321                         break;
 322                 case 'h':
 323                         usage();
 324                         exit(0);
 325                 case 'V':
 326                         printf("Version 0.0\n");
 327                         break;
 328
 329                 default:
 330                         /* unknown option or other error; error is
 331                         printed by getopt, just return */
 332                         return -1;
 333                 }
 334         }
 335
 336         return 0;
 337 }
 338
 339 void sigterm_handler(int sig)
 340 {
 341         libustd_stop_instance(instance, 0);
 342 }
 343
 344 int start_ustd(int fd)
 345 {
 346         struct ustcomm_ustd ustd;
 347         int result;
 348         sigset_t sigset;
 349         struct sigaction sa;
 350         int timeout = -1;
 351
 352         struct libustd_callbacks *callbacks = new_callbacks();
 353         if(!callbacks) {
 354                 PERROR("new_callbacks");
 355                 return 1;
 356         }
 357
 358         result = sigemptyset(&sigset);
 359         if(result == -1) {
 360                 PERROR("sigemptyset");
 361                 return 1;
 362         }
 363         sa.sa_handler = sigterm_handler;
 364         sa.sa_mask = sigset;
 365         sa.sa_flags = 0;
 366         result = sigaction(SIGTERM, &sa, NULL);
 367         if(result == -1) {
 368                 PERROR("sigaction");
 369                 return 1;
 370         }
 371         result = sigaction(SIGINT, &sa, NULL);
 372         if(result == -1) {
 373                 PERROR("sigaction");
 374                 return 1;
 375         }
 376
 377         instance = libustd_new_instance(callbacks, sock_path);
 378         if(!instance) {
 379                 ERR("failed to create libustd instance");
 380                 return 1;
 381         }
 382
 383         result = libustd_init_instance(instance);
 384         if(result) {
 385                 ERR("failed to initialize libustd instance");
 386                 return 1;
 387         }
 388
 389         /* setup handler for SIGPIPE */
 390         result = sigemptyset(&sigset);
 391         if(result == -1) {
 392                 PERROR("sigemptyset");
 393                 return 1;
 394         }
 395         result = sigaddset(&sigset, SIGPIPE);
 396         if(result == -1) {
 397                 PERROR("sigaddset");
 398                 return 1;
 399         }
 400         result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 401         if(result == -1) {
 402                 PERROR("sigprocmask");
 403                 return 1;
 404         }
 405
 406         /* Write pidfile */
 407         if(pidfile) {
 408                 result = write_pidfile(pidfile, getpid());
 409                 if(result == -1) {
 410                         ERR("failed to write pidfile");
 411                         return 1;
 412                 }
 413         }
 414
 415         /* Notify parent that we are successfully started. */
 416         if(fd != -1) {
 417                 /* write any one character */
 418                 result = write(fd, "!", 1);
 419                 if(result == -1) {
 420                         PERROR("write");
 421                         return -1;
 422                 }
 423                 if(result != 1) {
 424                         ERR("Problem sending confirmation of daemon start to parent");
 425                         return -1;
 426                 }
 427                 result = close(fd);
 428                 if(result == -1) {
 429                         PERROR("close");
 430                 }
 431         }
 432
 433         libustd_start_instance(instance);
 434
 435         free(callbacks);
 436
 437         return 0;
 438 }
 439
 440 int start_ustd_daemon()
 441 {
 442         int result;
 443         int fd[2];
 444         pid_t child_pid;
 445
 446         result = pipe(fd);
 447
 448         result = child_pid = fork();
 449         if(result == -1) {
 450                 PERROR("fork");
 451                 return -1;
 452         }
 453         else if(result == 0) {
 454                 return start_ustd(fd[1]);
 455         }
 456         else {
 457                 char buf;
 458
 459                 result = read(fd[0], &buf, 1);
 460                 if(result == -1) {
 461                         PERROR("read");
 462                         return -1;
 463                 }
 464                 if(result != 1) {
 465                         ERR("did not receive valid confirmation that the daemon is started");
 466                         return -1;
 467                 }
 468
 469                 result = close(fd[0]);
 470                 if(result == -1) {
 471                         PERROR("close");
 472                 }
 473
 474                 DBG("The daemon is now successfully started");
 475         }
 476
 477         /* Wait for confirmation that the server is ready. */
 478
 479
 480         return 0;
 481 }
 482
 483 int main(int argc, char **argv)
 484 {
 485         int result;
 486
 487         result = parse_args(argc, argv);
 488         if(result == -1) {
 489                 exit(1);
 490         }
 491
 492         if(daemon_mode) {
 493                 result = start_ustd_daemon();
 494         }
 495         else {
 496                 result = start_ustd(-1);
 497         }
 498
 499         return result;
 500 }