/*
 * Return how busy a server is via a TCP connection.
 *
 * This software takes the instantaneous load average (as provided
 * by /usr/bin/uptime) and compares it to a command line provided
 * load. If the load average is higher, the software returns a '0' 
 * to anyone that connects to its TCP port. Otherwise it returns a '1'.
 *
 * Developers that need to change the location of uptime should change
 * the #define UPTIME line near the top of this program.
 *
 * Tested on FreeBSD 4.9 and RedHat Linux 7.3 (kernel 2.4.26)
 * 
 * This software is by Steve Shah (sshah@planetoid.org)
 * This software is in the public domain. Do as you wish.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED 
 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING 
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 * 
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/param.h>
#include <sys/uio.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/fcntl.h>
#include <sys/socket.h>
#include <sys/time.h>

#define NIPQUAD(addr) \
         ((unsigned char *)&addr)[0], \
         ((unsigned char *)&addr)[1], \
         ((unsigned char *)&addr)[2], \
         ((unsigned char *)&addr)[3]

#define BUFSIZE 80
#define DOWN 0
#define UP 1
#define UPTIME "/usr/bin/uptime"

int main (int argc, char *argv[])
{
    struct sockaddr_in incoming_connection;
    struct sockaddr_in bound_connection;
    int incoming_fd;
    int incoming_socket;
    unsigned int incoming_connection_len;
    struct protoent *tcp;
    int port;
    float loadlimit;
    int errcode;
    int i;
    char buf[BUFSIZE];
    FILE *loadavg_fd;
    float loadavg;
    int decision=0;
    int debug=0;

    /*
     * Parse command line parameters 
     *
     */

    if (argc < 3) {
	printf ("Usage: %s <port> <loadlimit> [-d]\n", argv[0]);
	printf ("The load limit should be expressed in an uptime ");
	printf ("style number\n");
	printf ("The -d parameter enables debugging output.\n");
	exit(0);
    }

    port = atoi(argv[1]);
    if (port < 1024 || port > 65534) {
	printf ("Port number is out of range.\n");
	exit(0);
    }

    strncpy(buf,argv[2],BUFSIZE);
    buf[BUFSIZE-1] = '\0';
    sscanf(buf,"%f",&loadlimit);
    if (loadlimit < 0.0 || loadlimit > 100.0) {
	printf ("loadlimit must be between 0 and 100\n");
	exit (0);
    }

    if (argc == 4) {
	debug = 1;
	printf ("Debug is enabled.\n");
    }
    
    /* 
     * Get TCP's protocol number (needed for the socket() system call)
     *
     */

    tcp = getprotobyname("tcp");
    if (tcp == NULL) {
	printf ("Can't get protocol numbers for TCP\n");
	exit(0);
    }


    /*
     * Create a new socket 
     *
     */

    incoming_socket = socket(AF_INET, SOCK_STREAM, tcp->p_proto);
    if (incoming_socket < 0) {
	printf ("Can't create a socket\n");
	exit(0);
    }


    /* 
     * Bind to a specific port using the new socket
     *
     */

    bound_connection.sin_family      = AF_INET;
    bound_connection.sin_addr.s_addr = htonl(INADDR_ANY);
    bound_connection.sin_port        = htons(port);

    /* Note that sockaddr_in is a internet interpretation of sockaddr */

    errcode = bind(incoming_socket, 
		   (struct sockaddr *)&bound_connection,
		   sizeof (struct sockaddr));
    if (errcode < 0) {
	printf ("Can't bind to port %d\n", port);
	perror("");
	exit(0);
    }


    /*
     * Tell us how many incoming connections we can at most queue
     *
     */

    errcode = listen(incoming_socket, 128);
    if (errcode == -1) {
	perror ("Can't listen to bound socket");
	exit(0);
    }


    /*
     * Make socket non-blocking
     *
     */

    if (fcntl(incoming_socket, F_SETFL, O_NDELAY) < 0) {
        perror("Can't set socket to non-blocking");
        exit(0);
    }


    /*
     * Enter the accept loop
     *
     */

    incoming_connection_len = sizeof(struct sockaddr_in);

    if (debug) {
	printf ("Waiting for a connection...\n");
    }
    
    do {
	incoming_fd = accept(incoming_socket, 
			     (struct sockaddr *)&incoming_connection,
			     &incoming_connection_len);
	if (incoming_fd != -1) {
	    if (debug) {
		printf ("Accepted connection from %d.%d.%d.%d:%d\n",
			NIPQUAD(incoming_connection.sin_addr.s_addr),
			ntohs(incoming_connection.sin_port));
	    }
            loadavg_fd = popen (UPTIME,"r");
	    if (loadavg_fd == NULL) { 
		    decision = DOWN;
		    goto return_decision;
	    } 
	    if (fgets(buf,BUFSIZE,loadavg_fd) == NULL) {
		    decision = DOWN;
		    goto return_decision;
	    } 

	    /*
	     * output of load average is:
	     * 1:08pm up 4 days, 6:00, 3 users, load average: 0.00, 0.00, 0.00
	     *
	     */
	    
	    /* start from the end of the line until the second comma */
	    
	    for (i=strlen(buf)-1;i >= 0 && buf[i] != ',';i--);
	    if (i == 0) {
		decision = DOWN;
		if (debug) {
		    printf ("parse error 0 in load average.\n");
		}
		goto return_decision;
	    }
	    for (i--;i >= 0 && buf[i] != ',';i--);
	    if (i == 0) {
		decision = DOWN;
		if (debug) {
		    printf ("parse error 1 in load average.\n");
		}
		goto return_decision;
	    }
	    buf[i] = '\0';

	    /* now find the ' ' (space) */
	    for (;i >= 0 && buf[i] != ' ';i--);
	    if (i == 0) {
		decision = DOWN;
		if (debug) {
		    printf ("parse error 2 in load average.\n");
		}
		goto return_decision;
	    }

	    /* extract the load average as a number */
	    sscanf(&buf[i],"%f",&loadavg);

	    if (debug) {
		printf ("uptime: %s\n",buf);
		printf ("I think the load average is: %f\n", loadavg);
	    }
	    
	    if (loadavg >= loadlimit) {
		decision = DOWN;
	    } else {
		decision = UP;
	    }

	return_decision:
	    if (debug) {
		printf ("Returning %s\n", decision ? "UP":"DOWN");
	    }

	    if (decision == UP) {
		write(incoming_fd,"1",1);
	    } else {
		write(incoming_fd,"0",1);
	    }
	    close (incoming_fd);
	    if (debug) {
		printf ("Waiting for a connection...\n");
	    }
	} else {
	    sleep(1);
	}
    } while (1);

    return 0;
}