This commit is contained in:
Barak Michener 2013-08-31 22:26:31 -04:00
parent 7b32ca1133
commit 5a36066325
2 changed files with 39 additions and 44 deletions

View file

@ -1,2 +1,2 @@
all: all:
gcc -o tpm tpm.c g++ -o tpm tpm.cc

81
tpm.cc
View file

@ -12,30 +12,28 @@
#include <time.h> #include <time.h>
#include <unistd.h> #include <unistd.h>
#include <string>
using std::string;
const int kDefaultLengthSecs = 25 * 60; const int kDefaultLengthSecs = 25 * 60;
const char* kDefaultDoneMessage = ""; const char* kDefaultDoneMessage = "";
char* MakeSocketName() { string SocketName() {
char* username = getenv("USER"); string socket_name;
char* socket_name = (char*) calloc(strlen(username) + strlen(P_tmpdir) + 40, sizeof(char)); socket_name.append(P_tmpdir);
strcat(socket_name, P_tmpdir); socket_name.append("/.tpm-");
if (P_tmpdir[strlen(P_tmpdir) - 1] != '/') { socket_name.append(getenv("USER"));
strcat(socket_name, "/");
}
strcat(socket_name, ".tpm-");
strcat(socket_name, username);
return socket_name; return socket_name;
} }
char* MakePostHookPath() { string MakePostHookPath() {
char* home = getenv("HOME"); string post_hook;
char* post_hook = (char*) calloc(strlen(home) + 40, sizeof(char)); post_hook = getenv("HOME");
strcat(post_hook, home); post_hook += "/.tpm-post.sh";
strcat(post_hook, "/.tpm-post.sh");
return post_hook; return post_hook;
} }
int ClientMain(char* done_message) { int ClientMain(string done_message) {
int sock_fd; int sock_fd;
struct sockaddr_un remote; struct sockaddr_un remote;
if ((sock_fd = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) { if ((sock_fd = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) {
@ -43,13 +41,12 @@ int ClientMain(char* done_message) {
exit(1); exit(1);
} }
remote.sun_family = AF_UNIX; remote.sun_family = AF_UNIX;
char * socket_name = MakeSocketName(); string socket_name = SocketName();
strcpy(remote.sun_path, socket_name); strcpy(remote.sun_path, socket_name.c_str());
free(socket_name);
int len = strlen(remote.sun_path) + sizeof(remote.sun_family) + 1; int len = strlen(remote.sun_path) + sizeof(remote.sun_family) + 1;
if (connect(sock_fd, (struct sockaddr *)&remote, len) == -1) { if (connect(sock_fd, (struct sockaddr *)&remote, len) == -1) {
printf("%s\n", done_message); printf("%s\n", done_message.c_str());
return 0; return 0;
} }
@ -71,8 +68,8 @@ int DaemonMain(int countdown_time) {
ts.tv_sec = 0; ts.tv_sec = 0;
ts.tv_nsec = 1000 /*micro*/ * 1000 /* milli */ * 50; ts.tv_nsec = 1000 /*micro*/ * 1000 /* milli */ * 50;
char * socket_name = MakeSocketName(); string socket_name = SocketName();
int sock_fd; int sock_fd;
struct sockaddr_un local, remote; struct sockaddr_un local, remote;
@ -81,9 +78,8 @@ int DaemonMain(int countdown_time) {
exit(1); exit(1);
} }
local.sun_family = AF_UNIX; local.sun_family = AF_UNIX;
strcpy(local.sun_path, socket_name); strcpy(local.sun_path, socket_name.c_str());
unlink(local.sun_path); unlink(local.sun_path);
free(socket_name);
int len = strlen(local.sun_path) + sizeof(local.sun_family) + 1; int len = strlen(local.sun_path) + sizeof(local.sun_family) + 1;
if (bind(sock_fd, (struct sockaddr *)&local, len) == -1) { if (bind(sock_fd, (struct sockaddr *)&local, len) == -1) {
perror("bind"); perror("bind");
@ -105,46 +101,45 @@ int DaemonMain(int countdown_time) {
/* Change the file mode mask */ /* Change the file mode mask */
umask(0); umask(0);
/* Open any logs here */ /* Open any logs here */
/* Create a new SID for the child process */ /* Create a new SID for the child process */
sid = setsid(); sid = setsid();
if (sid < 0) { if (sid < 0) {
/* Log the failure */ /* Log the failure */
exit(EXIT_FAILURE); exit(EXIT_FAILURE);
} }
/* Change the current working directory */ /* Change the current working directory */
if ((chdir("/")) < 0) { if ((chdir("/")) < 0) {
/* Log the failure */ /* Log the failure */
exit(EXIT_FAILURE); exit(EXIT_FAILURE);
} }
/* Close out the standard file descriptors */ /* Close out the standard file descriptors */
close(STDIN_FILENO); close(STDIN_FILENO);
close(STDOUT_FILENO); close(STDOUT_FILENO);
close(STDERR_FILENO); close(STDERR_FILENO);
if (listen(sock_fd, 5) == -1) { if (listen(sock_fd, 5) == -1) {
exit(1); exit(1);
} }
/* The Big Loop */
int wakeup = 0;
struct timeval start_time; struct timeval start_time;
ok = gettimeofday(&start_time, NULL); ok = gettimeofday(&start_time, NULL);
int remote_fd = -1; int remote_fd = -1;
bool wakeup = false;
while (!wakeup) { while (!wakeup) {
struct timeval current_time; struct timeval current_time;
struct timeval diff_time; struct timeval diff_time;
ok = gettimeofday(&current_time, NULL); ok = gettimeofday(&current_time, NULL);
timersub(&current_time, &start_time, &diff_time); timersub(&current_time, &start_time, &diff_time);
if (diff_time.tv_sec >= countdown_time) { if (diff_time.tv_sec >= countdown_time) {
wakeup = 1; wakeup = true;
continue; continue;
} }
long remaining_time = (long)countdown_time - diff_time.tv_sec; long remaining_time = (long)countdown_time - diff_time.tv_sec;
@ -166,20 +161,20 @@ int DaemonMain(int countdown_time) {
close(sock_fd); close(sock_fd);
printf("%s\n", local.sun_path); printf("%s\n", local.sun_path);
unlink(local.sun_path); unlink(local.sun_path);
char* post_hook = MakePostHookPath(); string post_hook = MakePostHookPath();
if (access(post_hook, X_OK) != -1) { char* argv[0];
execv(post_hook, NULL); if (access(post_hook.c_str(), X_OK) != -1) {
execv(post_hook.c_str(), argv);
} }
free(post_hook);
exit(EXIT_SUCCESS); exit(EXIT_SUCCESS);
} }
int main(int argc, char** argv) { int main(int argc, char** argv) {
int c; int c;
int countdown_time = kDefaultLengthSecs; int countdown_time = kDefaultLengthSecs;
char* done_message = (char*) kDefaultDoneMessage; string done_message = kDefaultDoneMessage;
int got_positional = 0; int got_positional = 0;
char* command = ""; string command = "";
while (1) { while (1) {
while ( (c = getopt(argc, argv, "bs:m:d:")) != -1) { while ( (c = getopt(argc, argv, "bs:m:d:")) != -1) {
switch (c) { switch (c) {
@ -196,7 +191,7 @@ int main(int argc, char** argv) {
} }
if (optind < argc && !got_positional) { if (optind < argc && !got_positional) {
command = argv[optind]; command = argv[optind];
got_positional = 1; got_positional = 1;
} else { } else {
break; break;
} }
@ -204,11 +199,11 @@ int main(int argc, char** argv) {
} }
if (strcmp(command, "start") == 0) { if (command == "start") {
DaemonMain(countdown_time); DaemonMain(countdown_time);
} }
else { else {
return ClientMain(done_message); return ClientMain(done_message);
} }
exit(0); exit(0);
} }