1/* 2 * Copyright (c) 2010 Apple Inc. All rights reserved. 3 * 4 * @APPLE_LICENSE_HEADER_START@ 5 * 6 * This file contains Original Code and/or Modifications of Original Code 7 * as defined in and that are subject to the Apple Public Source License 8 * Version 2.0 (the 'License'). You may not use this file except in 9 * compliance with the License. Please obtain a copy of the License at 10 * http://www.opensource.apple.com/apsl/ and read it before using this 11 * file. 12 * 13 * The Original Code and all software distributed under the License are 14 * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER 15 * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES, 16 * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY, 17 * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT. 18 * Please see the License for the specific language governing rights and 19 * limitations under the License. 20 * 21 * @APPLE_LICENSE_HEADER_END@ 22 */ 23 24#include <errno.h> 25#include <stdio.h> 26#include <stdlib.h> 27#include <string.h> 28#include <unistd.h> 29#include <getopt.h> 30 31#include <dispatch/dispatch.h> 32#include <CoreFoundation/CFNumber.h> 33 34#include <IOKit/pwr_mgt/IOPMLib.h> 35#include <IOKit/pwr_mgt/IOPMLibPrivate.h> 36 37typedef enum { 38 kDefaultAssertionFlag = 0, 39 kIdleAssertionFlag = (1 << 0), 40 kDisplayAssertionFlag = (1 << 1), 41 kSystemAssertionFlag = (1 << 2), 42 kUserActiveAssertionFlag = (1 << 3), 43 kDiskAssertionFlag = (1 << 4) 44} AssertionFlag; 45 46typedef struct { 47 AssertionFlag assertionFlag; 48 CFStringRef assertionType; 49} AssertionMapEntry; 50 51 52AssertionMapEntry assertionMap[] = { 53 { kIdleAssertionFlag, kIOPMAssertionTypePreventUserIdleSystemSleep }, 54 { kDisplayAssertionFlag, kIOPMAssertionTypePreventUserIdleDisplaySleep }, 55 { kSystemAssertionFlag, kIOPMAssertionTypePreventSystemSleep}, 56 { kUserActiveAssertionFlag, kIOPMAssertionUserIsActive}, 57 { kDiskAssertionFlag, kIOPMAssertPreventDiskIdle}}; 58 59 60static CFStringRef kHumanReadableReason = CFSTR("THE CAFFEINATE TOOL IS PREVENTING SLEEP."); 61static CFStringRef kLocalizationBundlePath = CFSTR("/System/Library/CoreServices/powerd.bundle"); 62 63#define kAssertionNameString "caffeinate command-line tool" 64 65int createAssertions(const char *progname, AssertionFlag flags, long timeout); 66void forkChild(char *argv[]); 67void usage(void); 68 69pid_t waitforpid; 70 71 72int 73main(int argc, char *argv[]) 74{ 75 AssertionFlag flags = kDefaultAssertionFlag; 76 char ch; 77 unsigned long timeout = 0; 78 dispatch_source_t disp_src; 79 80 errno = 0; 81 while ((ch = getopt(argc, argv, "mdhisut:w:")) != -1) { 82 switch((char)ch) { 83 case 'm': 84 flags |= kDiskAssertionFlag; 85 break; 86 case 'd': 87 flags |= kDisplayAssertionFlag; 88 break; 89 case 'h': 90 usage(); 91 exit(EXIT_SUCCESS); 92 case 'i': 93 flags |= kIdleAssertionFlag; 94 break; 95 case 's': 96 flags |= kSystemAssertionFlag; 97 break; 98 case 'u': 99 flags |= kUserActiveAssertionFlag; 100 break; 101 case 'w': 102 waitforpid = strtol(optarg, NULL, 0); 103 if (waitforpid == 0 && errno != 0) { 104 usage(); 105 exit(EXIT_FAILURE); 106 } 107 break; 108 109 case 't': 110 timeout = strtol(optarg, NULL, 0); 111 if (timeout == 0 && errno != 0) { 112 usage(); 113 exit(EXIT_FAILURE); 114 } 115 break; 116 117 case '?': 118 default: 119 usage(); 120 exit(EXIT_FAILURE); 121 } 122 } 123 124 if (flags == kDefaultAssertionFlag) { 125 flags = kIdleAssertionFlag; 126 } 127 128 /* Spwan early, otherwise libraries might open resources that don't survive fork or exec */ 129 if (argc - optind) { 130 argv += optind; 131 forkChild(argv); 132 if (createAssertions(*argv, flags, timeout)) { 133 exit(EXIT_FAILURE); 134 } 135 } else { 136 if (timeout) { 137 dispatch_time_t d_timeout = dispatch_time(DISPATCH_TIME_NOW, timeout * NSEC_PER_SEC); 138 dispatch_after(d_timeout, dispatch_get_main_queue(), ^{ 139 exit(EXIT_SUCCESS); 140 }); 141 } 142 if (waitforpid) { 143 disp_src = dispatch_source_create(DISPATCH_SOURCE_TYPE_PROC, waitforpid, 144 DISPATCH_PROC_EXIT, dispatch_get_main_queue()); 145 146 dispatch_source_set_event_handler(disp_src, ^{ 147 exit(EXIT_SUCCESS); 148 }); 149 150 dispatch_source_set_cancel_handler(disp_src, ^{ 151 dispatch_release(disp_src); 152 }); 153 154 dispatch_resume(disp_src); 155 156 } 157 if (createAssertions(NULL, flags, timeout)) { 158 exit(EXIT_FAILURE); 159 } 160 } 161 162 dispatch_main(); 163} 164 165int 166createAssertions(const char *progname, AssertionFlag flags, long timeout) 167{ 168 IOReturn result = 1; 169 char assertionDetails[128]; 170 CFStringRef assertionDetailsString = NULL; 171 IOPMAssertionID assertionID = 0; 172 u_int i = 0; 173 174 assertionDetails[0] = 0; 175 if (progname) { 176 (void)snprintf(assertionDetails, sizeof(assertionDetails), 177 "caffeinate asserting on behalf of \'%s\' (pid %d)", progname, getppid()); 178 } else if (waitforpid) { 179 if (timeout) { 180 snprintf(assertionDetails, sizeof(assertionDetails), 181 "caffeinate asserting on behalf of Process ID %d for %ld secs", 182 waitforpid, timeout); 183 } 184 else { 185 snprintf(assertionDetails, sizeof(assertionDetails), 186 "caffeinate asserting on behalf of Process ID %d", waitforpid); 187 } 188 } else if (timeout) { 189 (void)snprintf(assertionDetails, sizeof(assertionDetails), 190 "caffeinate asserting for %ld secs", timeout); 191 } else { 192 (void)snprintf(assertionDetails, sizeof(assertionDetails), 193 "caffeinate asserting forever"); 194 } 195 196 assertionDetailsString = CFStringCreateWithCString(kCFAllocatorDefault, 197 assertionDetails, kCFStringEncodingMacRoman); 198 if (!assertionDetailsString) { 199 fprintf(stderr, "Failed to create assertion name %s\n", progname); 200 goto finish; 201 } 202 203 for (i = 0; i < sizeof(assertionMap)/sizeof(AssertionMapEntry); ++i) 204 { 205 AssertionMapEntry *entry = assertionMap + i; 206 207 if (!(flags & entry->assertionFlag)) continue; 208 if ( (entry->assertionFlag == kUserActiveAssertionFlag) && (timeout == 0)) 209 timeout = 5; /* Force a 5sec timeout on user active assertions */ 210 211 result = IOPMAssertionCreateWithDescription(entry->assertionType, 212 CFSTR(kAssertionNameString), assertionDetailsString, 213 kHumanReadableReason, kLocalizationBundlePath, 214 (CFTimeInterval)timeout, kIOPMAssertionTimeoutActionRelease, 215 &assertionID); 216 217 if (result != kIOReturnSuccess) 218 { 219 fprintf(stderr, "Failed to create %s assertion\n", 220 CFStringGetCStringPtr(entry->assertionType, kCFStringEncodingMacRoman)); 221 goto finish; 222 } 223 if (!progname && waitforpid) { 224 CFNumberRef waitforpidRef = CFNumberCreate(0, kCFNumberIntType, &waitforpid); 225 if (waitforpidRef) { 226 IOPMAssertionSetProperty(assertionID, kIOPMAssertionOnBehalfOfPID, waitforpidRef); 227 CFRelease(waitforpidRef); 228 } 229 } 230 } 231 232 result = kIOReturnSuccess; 233finish: 234 if (assertionDetailsString) CFRelease(assertionDetailsString); 235 236 return result; 237} 238 239void 240forkChild(char *argv[]) 241{ 242 pid_t pid; 243 dispatch_source_t source; 244 int fd, max_fd; 245 246 /* Our parent might care about the total life cycle of this process, 247 * therefore rather than propagate exit status, Unix signals, Mach 248 * exceptions, etc, we just flip the normal parent/child relationship and 249 * have the parent exec() and the child monitor the parent for death rather 250 * than the other way around. 251 */ 252 switch(fork()) { 253 case 0: /* child */ 254 break; 255 case -1: /* error */ 256 perror(""); 257 exit(EXIT_SUCCESS); 258 /* NOTREACHED */ 259 default: /* parent */ 260 execvp(*argv, argv); 261 int saved_errno = errno; 262 perror(*argv); 263 _exit((saved_errno == ENOENT) ? 127 : 126); 264 /* NOTREACHED */ 265 } 266 267 pid = getppid(); 268 if (pid == 1) { 269 exit(EXIT_SUCCESS); 270 } 271 272 /* Don't leak inherited FDs from our grandparent. */ 273 max_fd = getdtablesize(); 274 for (fd = STDERR_FILENO + 1; fd < max_fd; fd++) { 275 close(fd); 276 } 277 278 (void)signal(SIGINT, SIG_IGN); 279 (void)signal(SIGQUIT, SIG_IGN); 280 281 source = dispatch_source_create(DISPATCH_SOURCE_TYPE_PROC, pid, 282 DISPATCH_PROC_EXIT, dispatch_get_main_queue()); 283 dispatch_source_set_event_handler(source, ^{ 284 /* we do not need to waitpid() because we are actually the child */ 285 exit(EXIT_SUCCESS); 286 }); 287 dispatch_resume(source); 288 289 return; 290} 291 292void 293usage(void) 294{ 295 fprintf(stderr, "usage: caffeinate [-disu] [-t timeout] [-w Process ID] [command arguments...]\n"); 296 return; 297} 298