Reuse X connection when adjusting gamma ramps.

Save gamma ramps when program starts and restore them on exit (only RANDR).
Try both RANDR and VidMode if nothing was specified, otherwise try only the
 specified method.

5 files changed, 395 insertions, 202 deletions

diff --git a/src/randr.c b/src/randr.c
index 99ac9ec..b67484c 100644
--- a/src/randr.c
+++ b/src/randr.c
@@ -19,169 +19,270 @@
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 #include <xcb/xcb.h>
 #include <xcb/randr.h>
 
+#include "randr.h"
 #include "colorramp.h"
 
 
+#define RANDR_VERSION_MAJOR  1
+#define RANDR_VERSION_MINOR  3
+
+
 int
-randr_check_extension()
+randr_init(randr_state_t *state, int screen_num)
 {
 	xcb_generic_error_t *error;
 
 	/* Open X server connection */
-	xcb_connection_t *conn = xcb_connect(NULL, NULL);
+	int preferred_screen;
+	state->conn = xcb_connect(NULL, &preferred_screen);
+
+	if (screen_num < 0) screen_num = preferred_screen;
 
 	/* Query RandR version */
 	xcb_randr_query_version_cookie_t ver_cookie =
-		xcb_randr_query_version(conn, 1, 3);
+		xcb_randr_query_version(state->conn, RANDR_VERSION_MAJOR,
+					RANDR_VERSION_MINOR);
 	xcb_randr_query_version_reply_t *ver_reply =
-		xcb_randr_query_version_reply(conn, ver_cookie, &error);
+		xcb_randr_query_version_reply(state->conn, ver_cookie, &error);
 
 	if (error) {
 		fprintf(stderr, "RANDR Query Version, error: %d\n",
 			error->error_code);
-		xcb_disconnect(conn);
+		xcb_disconnect(state->conn);
 		return -1;
 	}
 
-	if (ver_reply->major_version < 1 || ver_reply->minor_version < 3) {
+	if (ver_reply->major_version < RANDR_VERSION_MAJOR ||
+	    ver_reply->minor_version < RANDR_VERSION_MINOR) {
 		fprintf(stderr, "Unsupported RANDR version (%u.%u)\n",
 			ver_reply->major_version, ver_reply->minor_version);
 		free(ver_reply);
-		xcb_disconnect(conn);
+		xcb_disconnect(state->conn);
 		return -1;
 	}
 
 	free(ver_reply);
 
-	/* Close connection */
-	xcb_disconnect(conn);
+	/* Get screen */
+	const xcb_setup_t *setup = xcb_get_setup(state->conn);
+	xcb_screen_iterator_t iter = xcb_setup_roots_iterator(setup);
+	state->screen = NULL;
 
-	return 0;
-}
+	for (int i = 0; iter.rem > 0; i++) {
+		if (i == screen_num) {
+			state->screen = iter.data;
+			break;
+		}
+		xcb_screen_next(&iter);
+	}
 
-static int
-randr_crtc_set_temperature(xcb_connection_t *conn, xcb_randr_crtc_t crtc,
-			   int temp, float gamma[3])
-{
-	xcb_generic_error_t *error;
+	if (state->screen == NULL) {
+		fprintf(stderr, "Screen %i could not be found.\n",
+			screen_num);
+		xcb_disconnect(state->conn);
+		return -1;
+	}
 
-	/* Request size of gamma ramps */
-	xcb_randr_get_crtc_gamma_size_cookie_t gamma_size_cookie =
-		xcb_randr_get_crtc_gamma_size(conn, crtc);
-	xcb_randr_get_crtc_gamma_size_reply_t *gamma_size_reply =
-		xcb_randr_get_crtc_gamma_size_reply(conn, gamma_size_cookie,
-						    &error);
+	/* Get list of CRTCs for the screen */
+	xcb_randr_get_screen_resources_current_cookie_t res_cookie =
+		xcb_randr_get_screen_resources_current(state->conn,
+						       state->screen->root);
+	xcb_randr_get_screen_resources_current_reply_t *res_reply =
+		xcb_randr_get_screen_resources_current_reply(state->conn,
+							     res_cookie,
+							     &error);
 
 	if (error) {
-		fprintf(stderr, "RANDR Get CRTC Gamma Size, error: %d\n",
-			error->error_code);
+		fprintf(stderr, "RANDR Get Screen Resources Current,"
+			" error: %d\n", error->error_code);
+		xcb_disconnect(state->conn);
 		return -1;
 	}
 
-	int gamma_ramp_size = gamma_size_reply->size;
+	state->crtc_count = res_reply->num_crtcs;
+	state->crtcs = malloc(state->crtc_count * sizeof(randr_crtc_state_t));
+	if (state->crtcs == NULL) {
+		perror("malloc");
+		xcb_disconnect(state->conn);
+		return -1;
+	}
 
-	free(gamma_size_reply);
+	xcb_randr_crtc_t *crtcs =
+		xcb_randr_get_screen_resources_current_crtcs(res_reply);
 
-	if (gamma_ramp_size == 0) {
-		fprintf(stderr, "Gamma ramp size too small: %i\n",
-			gamma_ramp_size);
-		return -1;
+	/* Save CRTC identifier in state */
+	for (int i = 0; i < state->crtc_count; i++) {
+		state->crtcs[i].crtc = crtcs[i];
 	}
 
-	/* Create new gamma ramps */
-	uint16_t *gamma_ramps = malloc(3*gamma_ramp_size*sizeof(uint16_t));
-	if (gamma_ramps == NULL) abort();
+	free(res_reply);
 
-	uint16_t *gamma_r = &gamma_ramps[0*gamma_ramp_size];
-	uint16_t *gamma_g = &gamma_ramps[1*gamma_ramp_size];
-	uint16_t *gamma_b = &gamma_ramps[2*gamma_ramp_size];
+	/* Save size and gamma ramps of all CRTCs.
+	   Current gamma ramps are saved so we can restore them
+	   at program exit. */
+	for (int i = 0; i < state->crtc_count; i++) {
+		xcb_randr_crtc_t crtc = state->crtcs[i].crtc;
+
+		/* Request size of gamma ramps */
+		xcb_randr_get_crtc_gamma_size_cookie_t gamma_size_cookie =
+			xcb_randr_get_crtc_gamma_size(state->conn, crtc);
+		xcb_randr_get_crtc_gamma_size_reply_t *gamma_size_reply =
+			xcb_randr_get_crtc_gamma_size_reply(state->conn,
+							    gamma_size_cookie,
+							    &error);
+
+		if (error) {
+			fprintf(stderr, "RANDR Get CRTC Gamma Size,"
+				" error: %d\n", error->error_code);
+			xcb_disconnect(state->conn);
+			return -1;
+		}
 
-	colorramp_fill(gamma_r, gamma_g, gamma_b, gamma_ramp_size,
-		       temp, gamma);
+		unsigned int ramp_size = gamma_size_reply->size;
+		state->crtcs[i].ramp_size = ramp_size;
 
-	/* Set new gamma ramps */
-	xcb_void_cookie_t gamma_set_cookie =
-		xcb_randr_set_crtc_gamma_checked(conn, crtc, gamma_ramp_size,
-						 gamma_r, gamma_g, gamma_b);
-	error = xcb_request_check(conn, gamma_set_cookie);
+		free(gamma_size_reply);
 
-	if (error) {
-		fprintf(stderr, "RANDR Set CRTC Gamma, error: %d\n",
-			error->error_code);
-		free(gamma_ramps);
-		return -1;
-	}
+		if (ramp_size == 0) {
+			fprintf(stderr, "Gamma ramp size too small: %i\n",
+				ramp_size);
+			xcb_disconnect(state->conn);
+			return -1;
+		}
+
+		/* Request current gamma ramps */
+		xcb_randr_get_crtc_gamma_cookie_t gamma_get_cookie =
+			xcb_randr_get_crtc_gamma(state->conn, crtc);
+		xcb_randr_get_crtc_gamma_reply_t *gamma_get_reply =
+			xcb_randr_get_crtc_gamma_reply(state->conn,
+						       gamma_get_cookie,
+						       &error);
+
+		if (error) {
+			fprintf(stderr, "RANDR Get CRTC Gamma, error: %d\n",
+				error->error_code);
+			xcb_disconnect(state->conn);
+			return -1;
+		}
 
-	free(gamma_ramps);
+		uint16_t *gamma_r =
+			xcb_randr_get_crtc_gamma_red(gamma_get_reply);
+		uint16_t *gamma_g =
+			xcb_randr_get_crtc_gamma_green(gamma_get_reply);
+		uint16_t *gamma_b =
+			xcb_randr_get_crtc_gamma_blue(gamma_get_reply);
+
+		/* Allocate space for saved gamma ramps */
+		state->crtcs[i].saved_ramps =
+			malloc(3*ramp_size*sizeof(uint16_t));
+		if (state->crtcs[i].saved_ramps == NULL) {
+			perror("malloc");
+			free(gamma_get_reply);
+			xcb_disconnect(state->conn);
+			return -1;
+		}
+
+		/* Copy gamma ramps into CRTC state */
+		memcpy(&state->crtcs[i].saved_ramps[0*ramp_size], gamma_r,
+		       ramp_size*sizeof(uint16_t));
+		memcpy(&state->crtcs[i].saved_ramps[1*ramp_size], gamma_g,
+		       ramp_size*sizeof(uint16_t));
+		memcpy(&state->crtcs[i].saved_ramps[2*ramp_size], gamma_b,
+		       ramp_size*sizeof(uint16_t));
+
+		free(gamma_get_reply);
+	}
 
 	return 0;
 }
 
-int
-randr_set_temperature(int screen_num, int temp, float gamma[3])
+void
+randr_restore(randr_state_t *state)
 {
 	xcb_generic_error_t *error;
 
-	/* Open X server connection */
-	int preferred_screen;
-	xcb_connection_t *conn = xcb_connect(NULL, &preferred_screen);
-
-	if (screen_num < 0) screen_num = preferred_screen;
-
-	/* Get screen */
-	const xcb_setup_t *setup = xcb_get_setup(conn);
-	xcb_screen_iterator_t iter = xcb_setup_roots_iterator(setup);
-	xcb_screen_t *screen = NULL;
-
-	for (int i = 0; iter.rem > 0; i++) {
-		if (i == screen_num) {
-			screen = iter.data;
-			break;
+	/* Restore CRTC gamma ramps */
+	for (int i = 0; i < state->crtc_count; i++) {
+		xcb_randr_crtc_t crtc = state->crtcs[i].crtc;
+
+		unsigned int ramp_size = state->crtcs[i].ramp_size;
+		uint16_t *gamma_r = &state->crtcs[i].saved_ramps[0*ramp_size];
+		uint16_t *gamma_g = &state->crtcs[i].saved_ramps[1*ramp_size];
+		uint16_t *gamma_b = &state->crtcs[i].saved_ramps[2*ramp_size];
+
+		/* Set gamma ramps */
+		xcb_void_cookie_t gamma_set_cookie =
+			xcb_randr_set_crtc_gamma_checked(state->conn, crtc,
+							 ramp_size, gamma_r,
+							 gamma_g, gamma_b);
+		error = xcb_request_check(state->conn, gamma_set_cookie);
+
+		if (error) {
+			fprintf(stderr, "RANDR Set CRTC Gamma, error: %d\n"
+				"Unable to restore CRTC %i\n",
+				error->error_code, i);
 		}
-		xcb_screen_next(&iter);
 	}
+}
 
-	if (screen == NULL) {
-		fprintf(stderr, "Screen %i could not be found.\n",
-			screen_num);
-		xcb_disconnect(conn);
-		return -1;
+void
+randr_free(randr_state_t *state)
+{
+	/* Free CRTC state */
+	for (int i = 0; i < state->crtc_count; i++) {
+		free(state->crtcs[i].saved_ramps);
 	}
+	free(state->crtcs);
 
-	/* Get list of CRTCs for the screen */
-	xcb_randr_get_screen_resources_current_cookie_t res_cookie =
-		xcb_randr_get_screen_resources_current(conn, screen->root);
-	xcb_randr_get_screen_resources_current_reply_t *res_reply =
-		xcb_randr_get_screen_resources_current_reply(conn, res_cookie,
-							     &error);
+	/* Close connection */
+	xcb_disconnect(state->conn);
+}
 
-	if (error) {
-		fprintf(stderr, "RANDR Get Screen Resources Current,"
-			" error: %d\n", error->error_code);
-		xcb_disconnect(conn);
-		return -1;
-	}
+int
+randr_set_temperature(randr_state_t *state, int temp, float gamma[3])
+{
+	xcb_generic_error_t *error;
 
-	xcb_randr_crtc_t *crtcs =
-		xcb_randr_get_screen_resources_current_crtcs(res_reply);
+	/* Set temperature on all CRTCs */
+	for (int i = 0; i < state->crtc_count; i++) {
+		xcb_randr_crtc_t crtc = state->crtcs[i].crtc;
+		unsigned int ramp_size = state->crtcs[i].ramp_size;
 
-	for (int i = 0; i < res_reply->num_crtcs; i++) {
-		int r = randr_crtc_set_temperature(conn, crtcs[i],
-						   temp, gamma);
-		if (r < 0) {
-			fprintf(stderr, "WARNING: Unable to adjust"
-				" CRTC %i.\n", i);
+		/* Create new gamma ramps */
+		uint16_t *gamma_ramps = malloc(3*ramp_size*sizeof(uint16_t));
+		if (gamma_ramps == NULL) {
+			perror("malloc");
+			return -1;
 		}
-	}
 
-	free(res_reply);
+		uint16_t *gamma_r = &gamma_ramps[0*ramp_size];
+		uint16_t *gamma_g = &gamma_ramps[1*ramp_size];
+		uint16_t *gamma_b = &gamma_ramps[2*ramp_size];
+
+		colorramp_fill(gamma_r, gamma_g, gamma_b, ramp_size,
+			       temp, gamma);
+
+		/* Set new gamma ramps */
+		xcb_void_cookie_t gamma_set_cookie =
+			xcb_randr_set_crtc_gamma_checked(state->conn, crtc,
+							 ramp_size, gamma_r,
+							 gamma_g, gamma_b);
+		error = xcb_request_check(state->conn, gamma_set_cookie);
+
+		if (error) {
+			fprintf(stderr, "RANDR Set CRTC Gamma, error: %d\n",
+				error->error_code);
+			free(gamma_ramps);
+			return -1;
+		}
 
-	/* Close connection */
-	xcb_disconnect(conn);
+		free(gamma_ramps);
+	}
 
 	return 0;
 }
diff --git a/src/randr.h b/src/randr.h
index d03f766..404bc29 100644
--- a/src/randr.h
+++ b/src/randr.h
@@ -20,7 +20,26 @@
 #ifndef _REDSHIFT_RANDR_H
 #define _REDSHIFT_RANDR_H
 
-int randr_check_extension();
-int randr_set_temperature(int screen_num, int temp, float gamma[3]);
+#include <xcb/xcb.h>
+#include <xcb/randr.h>
+
+typedef struct {
+	xcb_randr_crtc_t crtc;
+	unsigned int ramp_size;
+	uint16_t *saved_ramps;
+} randr_crtc_state_t;
+
+typedef struct {
+	xcb_connection_t *conn;
+	xcb_screen_t *screen;
+	unsigned int crtc_count;
+	randr_crtc_state_t *crtcs;
+} randr_state_t;
+
+
+int randr_init(randr_state_t *state, int screen_num);
+void randr_free(randr_state_t *state);
+void randr_restore(randr_state_t *state);
+int randr_set_temperature(randr_state_t *state, int temp, float gamma[3]);
 
 #endif /* ! _REDSHIFT_RANDR_H */
diff --git a/src/redshift.c b/src/redshift.c
index c9904d2..d34009f 100644
--- a/src/redshift.c
+++ b/src/redshift.c
@@ -28,9 +28,14 @@
 #include <time.h>
 #include <math.h>
 #include <locale.h>
+#include <sys/signal.h>
 
 #include "solar.h"
 
+#if !(defined(ENABLE_RANDR) || defined(ENABLE_VIDMODE))
+# error "At least one of RANDR or VidMode must be enabled."
+#endif
+
 #ifdef ENABLE_RANDR
 # include "randr.h"
 #endif
@@ -39,9 +44,16 @@
 # include "vidmode.h"
 #endif
 
-#if !(defined(ENABLE_RANDR) || defined(ENABLE_VIDMODE))
-# error "Either RANDR or VidMode must be enabled."
+
+/* Union of state data for gamma adjustment methods */
+typedef union {
+#ifdef ENABLE_RANDR
+	randr_state_t randr;
 #endif
+#ifdef ENABLE_VIDMODE
+	vidmode_state_t vidmode;
+#endif
+} gamma_state_t;
 
 
 /* Bounds for parameters. */
@@ -88,6 +100,60 @@
 #define DEG_CHAR  0xb0
 
 
+static volatile sig_atomic_t exiting = 0;
+
+/* Signal handler for exit signals */
+static void
+sigexit(int signo)
+{
+	exiting = 1;
+}
+
+static void
+gamma_state_restore(gamma_state_t *state, int use_randr)
+{
+#ifdef ENABLE_RANDR
+	if (use_randr == 1) randr_restore(&state->randr);
+#endif
+
+#ifdef ENABLE_VIDMODE
+	if (use_randr == 0) vidmode_restore(&state->vidmode);
+#endif
+}
+
+static void
+gamma_state_free(gamma_state_t *state, int use_randr)
+{
+#ifdef ENABLE_RANDR
+	if (use_randr == 1) randr_free(&state->randr);
+#endif
+
+#ifdef ENABLE_VIDMODE
+	if (use_randr == 0) vidmode_free(&state->vidmode);
+#endif
+}
+
+static int
+gamma_state_set_temperature(gamma_state_t *state, int use_randr,
+			    int temp, float gamma[3])
+{
+#ifdef ENABLE_RANDR
+	if (use_randr == 1) {
+		return randr_set_temperature(&state->randr, temp, gamma);
+	}
+#endif
+
+#ifdef ENABLE_VIDMODE
+	if (use_randr == 0) {
+		return vidmode_set_temperature(&state->vidmode, temp, gamma);
+	}
+#endif
+
+	return -1;
+}
+
+
+/* Calculate color temperature for the specified solar elevation. */
 static int
 calculate_temp(double elevation, int temp_day, int temp_night,
 	       int verbose)
@@ -112,58 +178,6 @@ calculate_temp(double elevation, int temp_day, int temp_night,
 	return temp;
 }
 
-static int
-adjust_temperature(int screen_num, int use_randr, int temp, float gamma[3])
-{
-	/* Set color temperature */
-	int failed = 0;
-#ifdef ENABLE_RANDR
-	if (use_randr) {
-		/* Check RANDR extension. */
-		int r = randr_check_extension();
-		if (r < 0) {
-			fprintf(stderr, "RANDR 1.3 extension is"
-				" not available.\n");
-			failed = 1;
-		} else {
-			r = randr_set_temperature(screen_num, temp, gamma);
-			if (r < 0) {
-				fprintf(stderr, "Unable to set color"
-					" temperature with RANDR.\n");
-				failed = 1;
-			}
-		}
-	}
-#endif
-
-#ifdef ENABLE_VIDMODE
-	if (!use_randr || failed) {
-		failed = 0;
-		/* Check VidMode extension */
-		int r = vidmode_check_extension();
-		if (r < 0) {
-			fprintf(stderr, "VidMode extension is"
-				" not available.\n");
-			failed = 1;
-		} else {
-			r = vidmode_set_temperature(screen_num, temp, gamma);
-			if (r < 0) {
-				fprintf(stderr, "Unable to set color"
-					" temperature with VidMode.\n");
-				failed = 1;
-			}
-		}
-	}
-#endif
-
-	if (failed) {
-		fprintf(stderr, "Color temperature adjustment failed.\n");
-		return -1;
-	}
-
-	return 0;
-}
-
 
 int
 main(int argc, char *argv[])
@@ -183,18 +197,13 @@ main(int argc, char *argv[])
 	int temp_day = DEFAULT_DAY_TEMP;
 	int temp_night = DEFAULT_NIGHT_TEMP;
 	float gamma[3] = { DEFAULT_GAMMA, DEFAULT_GAMMA, DEFAULT_GAMMA };
-	int use_randr = 1;
+	int use_randr = -1;
 	int screen_num = -1;
 	int init_trans = 1;
 	int one_shot = 0;
 	int verbose = 0;
 	char *s;
 
-#ifndef ENABLE_RANDR
-	/* Don't use RANDR if it has been disabled. */
-	use_randr = 0;
-#endif
-
 	/* Parse arguments. */
 	int opt;
 	while ((opt = getopt(argc, argv, "g:hl:m:ors:t:v")) != -1) {
@@ -345,12 +354,46 @@ main(int argc, char *argv[])
 		       gamma[0], gamma[1], gamma[2]);
 	}
 
+	/* Initialize gamma adjustment method. If use_randr is negative
+	   try all methods until one that works is found. */
+	gamma_state_t state;
+#ifdef ENABLE_RANDR
+	if (use_randr < 0 || use_randr == 1) {
+		/* Initialize RANDR state */
+		r = randr_init(&state.randr, screen_num);
+		if (r < 0) {
+			fprintf(stderr, "Initialization of RANDR failed.\n");
+			if (use_randr < 0) {
+				fprintf(stderr, "Trying other method...\n");
+			} else {
+				exit(EXIT_FAILURE);
+			}
+		} else {
+			use_randr = 1;
+		}
+	}
+#endif
+
+#ifdef ENABLE_VIDMODE
+	if (use_randr < 0 || use_randr == 0) {
+		/* Initialize VidMode state */
+		r = vidmode_init(&state.vidmode, screen_num);
+		if (r < 0) {
+			fprintf(stderr, "Initialization of VidMode failed.\n");
+			exit(EXIT_FAILURE);
+		} else {
+			use_randr = 0;
+		}
+	}
+#endif
+
 	if (one_shot) {
 		/* Current angular elevation of the sun */
 		struct timespec now;
 		r = clock_gettime(CLOCK_REALTIME, &now);
 		if (r < 0) {
 			perror("clock_gettime");
+			gamma_state_free(&state, use_randr);
 			exit(EXIT_FAILURE);
 		}
 
@@ -368,9 +411,11 @@ main(int argc, char *argv[])
 		if (verbose) printf("Color temperature: %uK\n", temp);
 
 		/* Adjust temperature */
-		r = adjust_temperature(screen_num, use_randr, temp, gamma);
+		r = gamma_state_set_temperature(&state, use_randr,
+						temp, gamma);
 		if (r < 0) {
 			fprintf(stderr, "Temperature adjustment failed.\n");
+			gamma_state_free(&state, use_randr);
 			exit(EXIT_FAILURE);
 		}
 	} else {
@@ -379,18 +424,31 @@ main(int argc, char *argv[])
 		r = clock_gettime(CLOCK_REALTIME, &short_trans_end);
 		if (r < 0) {
 			perror("clock_gettime");
+			gamma_state_free(&state, use_randr);
 			exit(EXIT_FAILURE);
 		}
 
 		int short_trans_len = 10;
 		short_trans_end.tv_sec += short_trans_len;
 
-		while (1) {
+		/* Install signal handler for SIGINT */
+		struct sigaction sigact;
+		sigset_t sigset;
+
+		sigemptyset(&sigset);
+		sigact.sa_handler = sigexit;
+		sigact.sa_mask = sigset;
+		sigact.sa_flags = 0;
+		sigaction(SIGINT, &sigact, NULL);
+
+		/* Continously adjust color temperature */
+		while (!exiting) {
 			/* Current angular elevation of the sun */
 			struct timespec now;
 			r = clock_gettime(CLOCK_REALTIME, &now);
 			if (r < 0) {
 				perror("clock_gettime");
+				gamma_state_free(&state, use_randr);
 				exit(EXIT_FAILURE);
 			}
 
@@ -417,18 +475,23 @@ main(int argc, char *argv[])
 			}
 
 			/* Adjust temperature */
-			r = adjust_temperature(screen_num, use_randr,
-					       temp, gamma);
+			r = gamma_state_set_temperature(&state, use_randr,
+							temp, gamma);
 			if (r < 0) {
 				fprintf(stderr, "Temperature adjustment"
 					" failed.\n");
+				gamma_state_free(&state, use_randr);
 				exit(EXIT_FAILURE);
 			}
 
 			if (init_trans) usleep(100000);
 			else usleep(5000000);
 		}
+
+		gamma_state_restore(&state, use_randr);
 	}
 
+	gamma_state_free(&state, use_randr);
+
 	return EXIT_SUCCESS;
 }
diff --git a/src/vidmode.c b/src/vidmode.c
index 93f7bdb..7d7dda7 100644
--- a/src/vidmode.c
+++ b/src/vidmode.c
@@ -24,91 +24,92 @@
 #include <X11/Xlib.h>
 #include <X11/extensions/xf86vmode.h>
 
+#include "vidmode.h"
 #include "colorramp.h"
 
 
 int
-vidmode_check_extension()
+vidmode_init(vidmode_state_t *state, int screen_num)
 {
 	int r;
 
 	/* Open display */
-	Display *dpy = XOpenDisplay(NULL);
-	if (dpy == NULL) {
+	state->display = XOpenDisplay(NULL);
+	if (state->display == NULL) {
 		fprintf(stderr, "XOpenDisplay failed.\n");
 		return -1;
 	}
 
+	if (screen_num < 0) screen_num = DefaultScreen(state->display);
+	state->screen_num = screen_num;
+
 	/* Query extension version */
 	int major, minor;
-	r = XF86VidModeQueryVersion(dpy, &major, &minor);
+	r = XF86VidModeQueryVersion(state->display, &major, &minor);
 	if (!r) {
 		fprintf(stderr, "XF86VidModeQueryVersion failed.\n");
+		XCloseDisplay(state->display);
 		return -1;
 	}
 
-	/* TODO Check major/minor */
+	return 0;
+}
 
-	/* Close display */
-	XCloseDisplay(dpy);
+void
+vidmode_free(vidmode_state_t *state)
+{
+	/* Close display connection */
+	XCloseDisplay(state->display);
+}
 
-	return 0;
+void
+vidmode_restore(vidmode_state_t *state)
+{
 }
 
 int
-vidmode_set_temperature(int screen_num, int temp, float gamma[3])
+vidmode_set_temperature(vidmode_state_t *state, int temp, float gamma[3])
 {
 	int r;
 
-	/* Open display */
-	Display *dpy = XOpenDisplay(NULL);
-	if (dpy == NULL) {
-		fprintf(stderr, "XOpenDisplay failed.\n");
-		return -1;
-	}
-
-	if (screen_num < 0) screen_num = DefaultScreen(dpy);
-
 	/* Request size of gamma ramps */
-	int gamma_ramp_size;
-	r = XF86VidModeGetGammaRampSize(dpy, screen_num, &gamma_ramp_size);
+	int ramp_size;
+	r = XF86VidModeGetGammaRampSize(state->display, state->screen_num,
+					&ramp_size);
 	if (!r) {
 		fprintf(stderr, "XF86VidModeGetGammaRampSize failed.\n");
-		XCloseDisplay(dpy);
 		return -1;
 	}
 
-	if (gamma_ramp_size == 0) {
+	if (ramp_size == 0) {
 		fprintf(stderr, "Gamma ramp size too small: %i\n",
-			gamma_ramp_size);
-		XCloseDisplay(dpy);
+			ramp_size);
 		return -1;
 	}
 
 	/* Create new gamma ramps */
-	uint16_t *gamma_ramps = malloc(3*gamma_ramp_size*sizeof(uint16_t));
-	if (gamma_ramps == NULL) abort();
+	uint16_t *gamma_ramps = malloc(3*ramp_size*sizeof(uint16_t));
+	if (gamma_ramps == NULL) {
+		perror("malloc");
+		return -1;
+	}
 
-	uint16_t *gamma_r = &gamma_ramps[0*gamma_ramp_size];
-	uint16_t *gamma_g = &gamma_ramps[1*gamma_ramp_size];
-	uint16_t *gamma_b = &gamma_ramps[2*gamma_ramp_size];
+	uint16_t *gamma_r = &gamma_ramps[0*ramp_size];
+	uint16_t *gamma_g = &gamma_ramps[1*ramp_size];
+	uint16_t *gamma_b = &gamma_ramps[2*ramp_size];
 
-	colorramp_fill(gamma_r, gamma_g, gamma_b, gamma_ramp_size,
-		       temp, gamma);
+	colorramp_fill(gamma_r, gamma_g, gamma_b, ramp_size, temp, gamma);
 
 	/* Set new gamma ramps */
-	r = XF86VidModeSetGammaRamp(dpy, screen_num, gamma_ramp_size,
-				    gamma_r, gamma_g, gamma_b);
+	r = XF86VidModeSetGammaRamp(state->display, state->screen_num,
+				    ramp_size, gamma_r, gamma_g, gamma_b);
 	if (!r) {
 		fprintf(stderr, "XF86VidModeSetGammaRamp failed.\n");
-		XCloseDisplay(dpy);
+		free(gamma_ramps);
 		return -1;
 	}
 
 	free(gamma_ramps);
 
-	/* Close display */
-	XCloseDisplay(dpy);
-
 	return 0;
 }
diff --git a/src/vidmode.h b/src/vidmode.h
index cf5516c..b342254 100644
--- a/src/vidmode.h
+++ b/src/vidmode.h
@@ -20,7 +20,16 @@
 #ifndef _REDSHIFT_VIDMODE_H
 #define _REDSHIFT_VIDMODE_H
 
-int vidmode_check_extension();
-int vidmode_set_temperature(int screen_num, int temp, float gamma[3]);
+#include <X11/Xlib.h>
+
+typedef struct {
+	Display *display;
+	int screen_num;
+} vidmode_state_t;
+
+int vidmode_init(vidmode_state_t *state, int screen_num);
+void vidmode_free(vidmode_state_t *state);
+void vidmode_restore(vidmode_state_t *state);
+int vidmode_set_temperature(vidmode_state_t *state, int temp, float gamma[3]);
 
 #endif /* ! _REDSHIFT_VIDMODE_H */