sproxy.c

#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <arpa/inet.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>
#include <stdbool.h>
#include <sys/time.h>
#include <time.h>
#include "LinkedList.c"

// Used for allowing debug prints
#define DEBUG false

#if DEBUG
    #define DEBUG_PRINT printf("[Debug] "); printf
#else
    #define DEBUG_PRINT(...)
#endif

// The maximum length of the string we can send (in characters)
#define MAX_LEN 20000

// timeout and heartbeat send rate in seconds
int delay = 1;

// seqn
int seqn = 0;
int last_ackn = -1;

// Both are defaults
char* IP = "127.0.0.1";
int PORT = 9002;

// For testing (Telnet is on port 23)
int outPort = 23;

// Counts the number of times the timeout happens in a row
int counter = 0;

// The sockets
int localSocket = -1;
int remoteSocket = -1;

// The linked list
LinkedList list;

// Force the program to fail (Testing purposes only)
bool forceFail = false;
bool AllowForceFail = false;    // If true, the program will force a disconnect every 5 messages

// Define the heartbeat header
typedef struct h_header {
    char type;
    unsigned long id;
    int ackn;
} h_header;

// Define the payload header
typedef struct p_header {
    char type;
    int seqn;
    int ackn;
    int size;
} p_header;

// Avoid implicit declaration
void ResendAllFromList(int ackn);

/* Creates a heartbeat message and returns it.
 * id: the session id
 * ackn: the ack number
 * returns: the heartbeat message
 */
char *create_heartbeat(unsigned long id, int ackn) {
    char *buffer = malloc(sizeof(h_header));
    h_header *header = (h_header *) buffer;
    header->type = 'h';
    header->id = id;
    header->ackn = ackn;
    return buffer;
}

/* Creates a payload message and returns it.
 * message: the message to be sent
 * n: the length of the message
 * seqn: the sequence number
 * ackn: the ack number
 * returns: the payload message
 */
char *create_payload(char *message, int n, int seqn, int ackn) {
    char *buffer = malloc(sizeof(p_header)+n+1);
    p_header *header = (p_header *)  buffer;
    header->type = 'p';
    header->seqn = seqn;
    header->ackn = ackn;
    header->size = n;
    memcpy(buffer+sizeof(p_header), message, n+1);
    return buffer;
}

// global values for unpack_message
int partial = 0;
int p_index = 0;
char p_buffer[MAX_LEN+1];
char **messages = NULL;
int msg_count = 0;

/* Unpacks the buffer into messages.
 * buffer: the buffer to be unpacked
 * len: the length of the buffer
 */
void unpack(char *buffer, int len) {
    // buffer recv'd from recv() and len of buffer requested
    msg_count = 0;
    if(messages != NULL) free(messages);
    messages = (char **) malloc(sizeof(char *) * ((int)(MAX_LEN/sizeof(h_header))+1));

    int i = 0;
    if(partial) {
        partial = 0;

        messages[msg_count++] = p_buffer;
        if(p_buffer[0] == 'h') {
            // copy rest of heartbeat into p_buffer
            i = sizeof(h_header) - p_index;
            memcpy(p_buffer + p_index, buffer, i);

            // execute heartbeat
            h_header *header = (h_header *) p_buffer;
            DEBUG_PRINT("[P] RECV'd hearbeat message:\tID:%lu\tACKN:%d\n", header->id, header->ackn);

            // Remove the messages that have been ack'd
            ClearSeqn(&list, header->ackn);
        }
        else if(p_buffer[0] == 'p') {
            if(p_index < sizeof(p_header)) {
                // copy rest of payload header first
                i = sizeof(p_header) - p_index;
                memcpy(p_buffer + p_index, buffer, i);

                // read payload header to see length of message
                p_header *header = (p_header *) p_buffer;
                int n = header->size;

                // copy n bytes of message
                memcpy(p_buffer + sizeof(p_header), buffer + i, n);
                i += n;
            }
            else {
                // read payload header to see length of message
                p_header *header = (p_header *) p_buffer;
                int n = header->size;
                int x = p_index - sizeof(p_header); // size of partial message

                // copy remaining bytes of message (n-x)
                memcpy(p_buffer + p_index, buffer, n-x);
                i = n-x;
            }

            // execute payload
            p_header *header = (p_header *) p_buffer;
            int n = header->size;

            char message[n];
            memcpy(message, p_buffer+sizeof(p_header), n);
            // DEBUG_PRINT("[P] RECV'd telnet command message:\n");
            // DEBUG_PRINT("\tSEQN: %d\n", header->seqn);
            // DEBUG_PRINT("\tACKN: %d\n", header->ackn);
            // DEBUG_PRINT("\tSIZE: %d\n", header->size);
            //DEBUG_PRINT("\tMESSAGE: %s\n", message);
            

            // Remove the messages that have been ack'd
            ClearSeqn(&list, header->ackn);
        }
        else {
            DEBUG_PRINT("Unexpected value at index[0] of partial buffer\n");
            exit(1);
        }
    }

    // read all messages in buffer
    while(i < len) {
        if(buffer[i] == 'h') {
            if(i+sizeof(h_header) <= len) {
                // read heartbeat
                messages[msg_count++] = buffer+i;
                h_header *header = (h_header *) (buffer+i);
                // DEBUG_PRINT("[P] RECV'd hearbeat message:\tID:%lu\tACKN:%d\n", header->id, header->ackn);
                i += sizeof(h_header);

                // Remove the messages that have been ack'd
                ClearSeqn(&list, header->ackn);
                continue;
            }
            else {
                // store partial heartbeat
                memcpy(p_buffer, buffer+i, len-i);
                partial = 1;
                p_index = len-i;
                return;
            }
        }
        else if(buffer[i] == 'p') {
            if(i+sizeof(p_header) <= len) {
                // read payload header
                p_header *header = (p_header *) (buffer+i);
                int n = header->size; // message length

                if(i+sizeof(p_header)+n <= len) {
                    // the entire payload is on buffer
                    char message[n];
                    messages[msg_count++] = buffer+i;
                    i += sizeof(p_header);
                    memcpy(message, buffer+i, n);
                    // DEBUG_PRINT("[P] RECV'd telnet command message:\n");
                    // DEBUG_PRINT("\tSEQN: %d\n", header->seqn);
                    // DEBUG_PRINT("\tACKN: %d\n", header->ackn);
                    // DEBUG_PRINT("\tSIZE: %d\n", header->size);
                    //DEBUG_PRINT("\tMESSAGE: %s\n", message);

                    // Remove the messages that have been ack'd
                    ClearSeqn(&list, header->ackn);
                    
                    i += n;
                    continue;
                }
                else {
                    // store partial message along with header
                    memcpy(p_buffer, buffer+i, len-i);
                    partial = 1;
                    p_index = len-i;
                    return;
                }
            }
            else {
                // store partial payload header
                memcpy(p_buffer, buffer+i, len-i);
                partial = 1;
                p_index = len-i;
                return;
            }
        }
        else {
            DEBUG_PRINT("End of message\n");
            return;
        }
    }
}

/* Send heartbeat to Client
 * socket: socket to send heartbeat to
 * return: true if successful, false otherwise
 */
bool SendHeartbeat(int socket) {
    char *heartbeat = create_heartbeat(999999, last_ackn); // ID, ACKN
    int n = send(socket, heartbeat, sizeof(h_header), 0);
    free(heartbeat);
    if(n == -1) {
        // failed to send heartbeat to server
        DEBUG_PRINT("Couldn't send heartbeat to server...\n");
        return false;
    }
    return true;
}

/* Recv message from local telnet daemon and process
 * the message with proper header and send to client
 * from: socket to recv message from
 * to: socket to send message to
 * return: true if successful, false otherwise
 */
bool SendTelnetOutput(int from, int to){
    
    // Receive message from telnet daemon
    char buf[MAX_LEN + 1];
    int len = recv(from, buf, MAX_LEN, 0);
    if (len == -1 || len == 0) {
        // There was an error (maybe connection closed)
        return false;
    }

    int curSeqn = ++seqn;

    // pack message with proper header
    char *payload = create_payload(buf, len, curSeqn, last_ackn); // SEQN, ACKN
    
    // Store the message in the Linked List (Will be freed there as well)
    AddNode(&list, payload, len+sizeof(p_header), curSeqn);

    // Send message to client
    int n = send(to, payload, len+sizeof(p_header), 0);
    if(n == -1) {
        DEBUG_PRINT("Error while sending message...\n");
        exit(1);
    }

    free(payload);

    return true;
}

/* Recv messages from Client and distinguish between
 * heartbeat and telnet messages
 * from: socket to recv message from
 * to: socket to send message to
 * return: true if successful, false otherwise
 */
bool RecvClientMessage (int from, int to){
    
    // Receive message
    char buf[MAX_LEN + 1];
    int len = recv(from, buf, MAX_LEN, 0);
    if (len == -1 || len == 0) {
        // There was an error (maybe connection closed)
        return false;
    }

    // process message
    unpack(buf, len);

    // Send payload messages to local telnet daemon
    for(int i=0; i<msg_count; i++) {
        if(messages[i][0] == 'p') {
            p_header *header = (p_header *) messages[i];
            
            // Update last ackn
            if (header->seqn > last_ackn - 1) last_ackn = header->seqn + 1; // one more than the last seqn received
            DEBUG_PRINT("Got packet with seqn: %d\tackn: %d\tsize: %d\n", header->seqn, header->ackn, header->size);
            DEBUG_PRINT("Set last_ackn to %d\n", last_ackn);

            // Testing function to force a timeout
            if (last_ackn % 5 == 4 && AllowForceFail)
                forceFail = true;

            char *telnet_message = messages[i]+sizeof(p_header);
            int n = send(to, telnet_message, header->size, 0);
            if(n == -1) {
                DEBUG_PRINT("Error while sending message...\n");
                exit(1);
            }
        }
    }

    /*
    // process partial message by calling recv
    // with appropriate byte count
    while(partial) {
        char buf1[MAX_LEN+1], buf2[MAX_LEN+1];
        int len1, len2;
        if(p_buffer[0] == 'h') {
            len1 = recv(from, buf1, sizeof(h_header)-p_index, 0);
            if(len1 == -1 || len1 == 0) {
                partial = 0;
                p_index = 0;
                return false;
            }
            unpack(buf1, len1);
        }
        else if(p_buffer[0] == 'p') {
            if(p_index < sizeof(p_header)) {
                // get remainder of the payload header to find size
                // of payload message
                len1 = recv(from, buf1, sizeof(p_header)-p_index, 0);
                if(len1 == -1 || len1 == 0) {
                    partial = 0;
                    p_index = 0;
                    return false;
                }
                memcpy(p_buffer+p_index, buf1, len1);
                p_index += len1;
            }

            // read payload header
            p_header *header = (p_header *) p_buffer;

            // recv length of payload message size
            len2 = recv(from, buf2, header->size, 0);
            if(len2 == -1 || len2 == 0) {
                partial = 0;
                p_index = 0;
                return false;
            }

            unpack(buf2, len2);
        }
    }

    // Send partial payload messages to local telnet daemon
    for(int i=0; i<msg_count; i++) {
        if(messages[i][0] == 'p') {
            p_header *header = (p_header *) messages[i];
            char *telnet_message = messages[i]+sizeof(p_header);
            int n = send(to, telnet_message, header->size, 0);
            if(n == -1) {
                DEBUG_PRINT("Error while sending message...\n");
                exit(1);
            }
        }
    }*/

    return true;
}

// Create global cliAddr and function to accept
// incoming client connection and return socket
struct sockaddr_in cliAddr;
socklen_t len = sizeof(cliAddr);
int AcceptConnection(int _s) {
    int s;

    DEBUG_PRINT("Waiting to recv client connect\n");

    do {
    s = accept(_s, (struct sockaddr *) &cliAddr, &len);
        if(s == -1) {
            DEBUG_PRINT("Error while accepting client connection...\n");
        }
    } while(s < 0);

    DEBUG_PRINT("Accepted incoming client connection\n");

    remoteSocket = s;

    return s;
}

/* Connects to the local telnet daemon
 * return: socket to telnet daemon
 */
int ConnectToTelnet() {
    int s;

    struct sockaddr_in serverAddr;
    memset(&serverAddr, 0, sizeof(serverAddr));


    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = inet_addr(IP);
    serverAddr.sin_port = htons(outPort);

    s = socket(PF_INET, SOCK_STREAM, 0);
    if(s == -1) {
        DEBUG_PRINT("Error while creating a socket...\n");
        exit(1);
    }

    int n = connect(s, (struct sockaddr *) &serverAddr, sizeof(serverAddr));
    if(n == -1) {
        DEBUG_PRINT("Error while connecting to server...\n");
        exit(1);
    }

    return s;
}

/* Create socket to recv client connections
 * return: socket to recv client connections
 */
int CreateSocket() {

    int n, s;

    // Set up local socket with ip and port
    struct sockaddr_in myAddr;
    memset(&myAddr, 0, sizeof(myAddr));

    myAddr.sin_family = AF_INET;
    myAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    myAddr.sin_port = htons(PORT);
    
    s = socket(PF_INET, SOCK_STREAM, 0);
    if(s == -1) {
        DEBUG_PRINT("Error while creating a socket...\n");
        exit(1);
    }

    n = bind(s, (struct sockaddr *) &myAddr, sizeof(myAddr));
    if(n == -1) {
        DEBUG_PRINT("Error while binding...\n");
        exit(1);
    }

    n = listen(s, 5);
    if(n == -1) {
        DEBUG_PRINT("Error while listening...\n");
        exit(1);
    }

    return s;
}

/* Handles forwarding messages between the client and server
 * s1: socket to client
 * s2: socket to telnet daemon
 */
void HandleForwarding(int s1, int s2) {

    int n;
    fd_set readfds;

    // to keep track of time to send heartbeat every second
    time_t start, end;
    double elapsed;
    time(&start);

    // To keep trach of the time since the last message from the client (Set it to)
    time_t last_msg;
    time(&last_msg);

    // Loop while the connection is open (This calls select() every second)
    while (1) {

        // If it has been 3 seconds since the last message from the client, close the connection
        time_t curTime;
        time(&curTime);
        if(difftime(curTime, last_msg) >= 3) {
            DEBUG_PRINT("Closing connection due to inactivity...\n");
            return;
        }

        // clear the set ahead of time
        FD_ZERO(&readfds);
        // add our descriptors to the set
        FD_SET(s1, &readfds);
        FD_SET(s2, &readfds);

        // find the largest descriptor, and plus one.
        if (s1 > s2) 
            n = s1 + 1;
        else 
            n = s2 + 1;
        
        // set timeout to 1 second
        struct timeval timeout;
        timeout.tv_sec = delay;
        timeout.tv_usec = 0;

        time(&end);
        elapsed = difftime(end, start);

        if(elapsed >= delay) {
            // send heartbeat every second
            SendHeartbeat(s1);

            time(&start);
        }

        // select() will return when at least one of the sockets has incoming traffic,
        // or after timeout.
        int rv = select(n, &readfds, NULL, NULL, &timeout);

        // Force failure if forceFail is true
        if (forceFail){
            forceFail = false;
            DEBUG_PRINT("Forcing failure...\n");
            sleep(4);
            return;
        }

        if (rv == -1) {
            perror("select"); // error occurred in select()
        } else if (rv == 0) {
            // timeout occurred
            counter++;
            // DEBUG_PRINT("Timeout occurred! No data after %d seconds. (%d)\n", delay, counter);

            if (counter >= 3){
                DEBUG_PRINT("Closing connection due to 3 timeouts...\n");
                return;
            }
        } else {
            // one or both of the descriptors have data
            if (FD_ISSET(s1, &readfds)) {
                // Recv'd from client, encode and forward
                // to telnet daemon
                //gettimeofday(&start_time, NULL);

                // Update last_msg from client
                time(&last_msg);

                bool result = RecvClientMessage(s1, s2);

                if (!result) {
                    return;
                }

                // Reset counter
                counter = 0;
            }
            if (FD_ISSET(s2, &readfds)) {
                // Recv'd from telnet, pack and forward
                // to client

                bool result = SendTelnetOutput(s2, s1);

                if (!result) {
                    return;
                }
            }
        }
    }
}

/* Called upon reconnection to send all unacked messages back to the server
 * ackn: the ackn of the last message received from the server
*/
void ResendAllFromList(int ackn){

    DEBUG_PRINT("Resending all messages up to ackn %d\n", ackn);

    // Tell the list to remove all nodes up to the ackn
    ClearSeqn(&list, ackn);

    // Tell the list to start iterating from the beginning
    StartIterator(&list);

    // Iterate through the list and resend all messages
    while (HasNext(&list)){

        Node* node = Next(&list);

        DEBUG_PRINT("Resending message with Seqn %d\n", node->seqn);

        // Send message to server
        int n = send(remoteSocket, node->msg, node->len, 0);
        if(n == -1) {
            printf("Error while sending message...\n");
            exit(1);
        }
    }

}

unsigned long int curr_id = 0;

int main(int argc, char* argv[])
{	
    // Initialize linked list
    list = NewLinkedList();

    // Temp change for testing (Taking an out port)

    // Check if port was given
    if(argc >= 2) {
        PORT = atoi(argv[1]);

        if (argc >= 3) {
            outPort = atoi(argv[2]);
        }
        else {
            DEBUG_PRINT("Using [%d] because argument [outPort] was not given\n", outPort);
        }

    }
    else {
        DEBUG_PRINT("Using [%d] because argument [PORT] was not given\n", PORT);
    }


    int _s1, s1, s2 = -1;

    // Create socket for incoming client connections
    _s1 = CreateSocket();


    char buf[MAX_LEN+1];

    while (1){
        // Accept incoming client connections to s1
        s1 = AcceptConnection(_s1);

        // Reset counter
        counter = 0;

        // Receive heartbeat from client to check ID
        int len = recv(s1, buf, MAX_LEN, 0);
        if(len == -1 || len == 0) {
            DEBUG_PRINT("Failed to receive initial heartbeat from client upon connection\n");
            exit(1);
        }
        unpack(buf, len);
        if(messages[0][0] == 'h') {
            h_header *header = (h_header *) messages[0];
            DEBUG_PRINT("ID of new client connection: %lu\n", header->id);

            if(header->id != curr_id) {
                // reset parameters for new connection
                Clear(&list);
                seqn = 0;
                last_ackn = -1;

                // new client connection detected
                curr_id = header->id;

                // close previous telnet session
                if(s2 != -1) close(s2);

                // create new telnet session
                s2 = ConnectToTelnet();
            }
            else {
                if(header->ackn == -1) {
                    // client telnet session was reset
                    Clear(&list);
                    seqn = 0;
                    last_ackn = -1;

                    // close and restart telnet session
                    if(s2 != -1) close(s2);
                    s2 = ConnectToTelnet();
                }
                // Resend all unacked messages
                ResendAllFromList(header->ackn);

                DEBUG_PRINT("Reconnected to previously connected client\n");
            }
        }
        else {
            DEBUG_PRINT("Client did not send a heartbeat message upon connection\n");
            exit(1);
        }

        // Handle forwarding
        HandleForwarding(s1, s2);

        // Close connection from client
        close(s1);
    }

    close(s1);
    close(s2);
}