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