server.c

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <time.h>
#include "sham.h"
#include <openssl/md5.h>
#include <stdarg.h>

#define BUFFER_SIZE 1040 // 1024 (payload) + 16 (header)
#define PAYLOAD_SIZE 1024
#define WINDOW_SIZE 10

typedef struct {
    char data[PAYLOAD_SIZE];
    uint32_t seq_num;
    int received;
    size_t data_len; // To handle variable-size packets
} PacketBuffer;

PacketBuffer buffer_storage[WINDOW_SIZE * 2];
int buffered_packets_count = 0;
uint32_t expected_seq_num = 1;

FILE *log_file = NULL;

void log_event(const char *format, ...) {
    if (!log_file) {
        return;
    }

    char time_buffer[30];
    struct timeval tv;
    time_t curtime;
    gettimeofday(&tv, NULL);
    curtime = tv.tv_sec;

    strftime(time_buffer, 30, "%Y-%m-%d %H:%M:%S", localtime(&curtime));

    fprintf(log_file, "[%s.%06d] [LOG] ", time_buffer, tv.tv_usec);

    va_list args;
    va_start(args, format);
    vfprintf(log_file, format, args);
    va_end(args);

    fprintf(log_file, "\n");
    fflush(log_file);
}

// Function to simulate packet loss
int simulate_packet_loss(double loss_rate) {
    if (loss_rate > 0.0 && ((double)rand() / RAND_MAX) < loss_rate) {
        return 1; // Drop the packet
    }
    return 0; // Send the packet
}

void send_ack(int sockfd, struct sockaddr_in *client_addr, socklen_t addr_len, uint32_t ack_num, double loss_rate) {
    if (simulate_packet_loss(loss_rate)) {
        log_event("DROP ACK=%u", ack_num);
        return;
    }
    char ack_buffer[sizeof(struct sham_header)];
    struct sham_header *header = (struct sham_header *)ack_buffer;
    
    header->seq_num = 0;
    header->ack_num = ack_num;
    header->flags = ACK_FLAG;

    uint32_t remaining_window_slots = (WINDOW_SIZE * 2) - buffered_packets_count;
    header->window_size = remaining_window_slots * PAYLOAD_SIZE;
    
    sham_header_to_network(header);
    sendto(sockfd, ack_buffer, sizeof(struct sham_header), 0, (const struct sockaddr *)client_addr, addr_len);
    log_event("SND ACK=%u WIN=%u", ack_num, header->window_size);
}

void perform_four_way_handshake(int sockfd, struct sockaddr_in *client_addr, socklen_t addr_len, double loss_rate) {
    char buffer[BUFFER_SIZE];
    struct sham_header *header = (struct sham_header *)buffer;
    
    printf("\nPerforming four-way handshake from server side...\n");

    // Wait for client's FIN with retransmission logic
    printf("Waiting for client's FIN...\n");
    fd_set read_fds;
    struct timeval timeout;
    
    int fin_received = 0;
    uint32_t client_fin_seq = 0;
    
    // This loop is the key change. The server will now retry receiving the FIN.
    for (int i = 0; i < 5; i++) { // Retry loop for receiving FIN
        FD_ZERO(&read_fds);
        FD_SET(sockfd, &read_fds);
        timeout.tv_sec = 10;
        timeout.tv_usec = 0;

        if (select(sockfd + 1, &read_fds, NULL, NULL, &timeout) > 0) {
            ssize_t bytes_received = recvfrom(sockfd, buffer, BUFFER_SIZE, 0, (struct sockaddr *)client_addr, &addr_len);
            if (bytes_received > 0) {
                sham_header_to_host(header);
                if (header->flags & FIN_FLAG) {
                    log_event("RCV FIN SEQ=%u", header->seq_num);
                    client_fin_seq = header->seq_num;
                    fin_received = 1;
                    break;
                }
            }
        } else {
            printf("Timeout waiting for client's FIN. Retrying...\n");
        }
    }
    
    if (!fin_received) {
        printf("Failed to receive client's FIN after multiple retries. Aborting handshake.\n");
        return;
    }
    
    // Send ACK for client's FIN
    printf("Sending ACK for client's FIN...\n");
    memset(buffer, 0, sizeof(struct sham_header));
    header->seq_num = 0;
    header->ack_num = client_fin_seq + 1;
    header->flags = ACK_FLAG;
    header->window_size = 0;
    
    if (!simulate_packet_loss(loss_rate)) {
        sham_header_to_network(header);
        sendto(sockfd, buffer, sizeof(struct sham_header), 0, (const struct sockaddr *)client_addr, addr_len);
        sham_header_to_host(header);
        log_event("SND ACK FOR FIN ACK=%u", header->ack_num);
    } else {
        printf("[Packet Dropped] Server ACK for FIN packet dropped.\n");
    }
    
    // Send own FIN
    printf("Sending server's FIN...\n");
    memset(buffer, 0, sizeof(struct sham_header));
    header->seq_num = 1;
    header->ack_num = 0;
    header->flags = FIN_FLAG;
    header->window_size = 0;
    
    int fin_retransmissions = 0;
    int final_ack_received = 0;
    
    while(fin_retransmissions < 5 && !final_ack_received) {
        if (!simulate_packet_loss(loss_rate)) {
            sham_header_to_network(header);
            sendto(sockfd, buffer, sizeof(struct sham_header), 0, (const struct sockaddr *)client_addr, addr_len);
            sham_header_to_host(header);
            log_event("SND FIN SEQ=%u", header->seq_num);
        } else {
            printf("[Packet Dropped] Server FIN packet dropped.\n");
        }

        // Wait for final ACK from client
        FD_ZERO(&read_fds);
        FD_SET(sockfd, &read_fds);
        timeout.tv_sec = 2;
        timeout.tv_usec = 0;

        if (select(sockfd + 1, &read_fds, NULL, NULL, &timeout) > 0) {
            recvfrom(sockfd, buffer, BUFFER_SIZE, 0, NULL, NULL);
            sham_header_to_host(header);
            if (header->flags & ACK_FLAG) {
                log_event("RCV ACK FOR FIN");
                final_ack_received = 1;
                break;
            }
        }
        fin_retransmissions++;
    }
    
    if (final_ack_received) {
        printf("Four-way handshake completed successfully.\n");
    } else {
        printf("Timeout waiting for final ACK. Connection closed.\n");
    }
}

// --- Chat Mode Logic ---
void run_chat_mode(int sockfd, struct sockaddr_in *client_addr, socklen_t addr_len, double loss_rate) {
    char buffer[BUFFER_SIZE];
    struct sham_header *header = (struct sham_header *)buffer;
    fd_set read_fds;
    
    printf("Entering Chat Mode. Type messages and press Enter to send.\n");
    
    while(1) {
        FD_ZERO(&read_fds);
        FD_SET(STDIN_FILENO, &read_fds);
        FD_SET(sockfd, &read_fds);
        
        int max_fd = (STDIN_FILENO > sockfd) ? STDIN_FILENO : sockfd;
        
        if (select(max_fd + 1, &read_fds, NULL, NULL, NULL) == -1) {
            perror("select");
            exit(EXIT_FAILURE);
        }

        // Handle incoming data from the socket
        if (FD_ISSET(sockfd, &read_fds)) {
            ssize_t bytes_received = recvfrom(sockfd, buffer, BUFFER_SIZE, 0, (struct sockaddr *)client_addr, &addr_len);
            if (bytes_received > (ssize_t)sizeof(struct sham_header)) {
                 sham_header_to_host(header);
                 
                 if (header->flags & FIN_FLAG) {
                    printf("Client is closing the connection.\n");
                    return; // Exit chat mode to handle handshake
                 }

                 printf("Client: %s\n", buffer + sizeof(struct sham_header));
            } else if (bytes_received > 0) {
                // Handle control packets (like FIN)
                sham_header_to_host(header);
                if (header->flags & FIN_FLAG) {
                    printf("Received FIN from client in chat mode.\n");
                    return;
                }
            }
        }
        
        // Handle input from standard input
        if (FD_ISSET(STDIN_FILENO, &read_fds)) {
            char chat_msg[PAYLOAD_SIZE];
            if (fgets(chat_msg, PAYLOAD_SIZE, stdin) != NULL) {
                chat_msg[strcspn(chat_msg, "\n")] = 0;
                
                // Prepare response packet
                memset(buffer, 0, BUFFER_SIZE);
                header->seq_num = 0;
                header->ack_num = 0;
                header->flags = 0;
                header->window_size = 1024;
                memcpy(buffer + sizeof(struct sham_header), chat_msg, strlen(chat_msg) + 1);

                // Add this check to simulate packet loss for chat messages
                if (!simulate_packet_loss(loss_rate)) {
                    sham_header_to_network(header);
                    sendto(sockfd, buffer, sizeof(struct sham_header) + strlen(chat_msg) + 1, 0, (const struct sockaddr *)client_addr, addr_len);
                    printf("You: %s\n", chat_msg);
                } else {
                    printf("[Packet Dropped] You tried to send: %s\n", chat_msg);
                }
            }
        }
    }
}

int main(int argc, char *argv[]) {
    srand(time(NULL));

    // Check for logging environment variable
    if (getenv("RUDP_LOG") != NULL && strcmp(getenv("RUDP_LOG"), "1") == 0) {
        log_file = fopen("server_log.txt", "a");
        if (log_file == NULL) {
            perror("Failed to open server_log.txt");
        }
    }

    int port;
    int chat_mode = 0;
    double loss_rate = 0.0;
    
    if (argc < 2) {
        fprintf(stderr, "Usage: ./server <port> [--chat] [loss_rate]\n");
        if (log_file) fclose(log_file);
        exit(EXIT_FAILURE);
    }
    
    port = atoi(argv[1]);
    
    if (argc > 2) {
        for (int i = 2; i < argc; i++) {
            if (strcmp(argv[i], "--chat") == 0) {
                chat_mode = 1;
            } else {
                loss_rate = atof(argv[i]);
                if (loss_rate < 0.0 || loss_rate > 1.0) {
                    fprintf(stderr, "Error: loss_rate must be between 0.0 and 1.0\n");
                    if (log_file) fclose(log_file);
                    exit(EXIT_FAILURE);
                }
            }
        }
    }

    int sockfd;
    struct sockaddr_in server_addr, client_addr;
    socklen_t addr_len = sizeof(client_addr);
    char recv_buffer[BUFFER_SIZE];
    struct sham_header *header = (struct sham_header *)recv_buffer;
    FILE *received_file;
    char received_file_name[256] = "received_file.txt";

    if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("socket creation failed");
        if (log_file) fclose(log_file);
        exit(EXIT_FAILURE);
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);

    if (bind(sockfd, (const struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind failed");
        if (log_file) fclose(log_file);
        exit(EXIT_FAILURE);
    }

    printf("Server listening on port %d with a packet loss rate of %.2f\n", port, loss_rate);
    
    // Initial recvfrom to get client's address for the handshake
    ssize_t bytes_received = recvfrom(sockfd, recv_buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &addr_len);
    
    // --- Three-Way Handshake (Server Side) ---
    printf("Waiting for SYN packet from client...\n");
    if (bytes_received > 0) {
        sham_header_to_host(header);
        if (header->flags == SYN_FLAG) {
            uint32_t client_isn = header->seq_num;
            log_event("RCV SYN SEQ=%u", client_isn);
            
            uint32_t server_isn = 5000;
            header->seq_num = server_isn;
            header->ack_num = client_isn + 1;
            header->flags = SYN_FLAG | ACK_FLAG;
            header->window_size = WINDOW_SIZE * PAYLOAD_SIZE;
            
            sham_header_to_network(header);
            sendto(sockfd, recv_buffer, sizeof(struct sham_header), 0, (const struct sockaddr *)&client_addr, addr_len);
            sham_header_to_host(header);
            log_event("SND SYN-ACK SEQ=%u ACK=%u", header->seq_num, header->ack_num);
            
            recvfrom(sockfd, recv_buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &addr_len);
            sham_header_to_host(header);
            
            if (header->flags == ACK_FLAG && header->ack_num == server_isn + 1) {
                log_event("RCV ACK FOR SYN");
                printf("Connection established.\n\n");
            } else {
                fprintf(stderr, "Handshake failed.\n");
                close(sockfd);
                if (log_file) fclose(log_file);
                exit(EXIT_FAILURE);
            }
        }
    }
    
    if (chat_mode) {
        run_chat_mode(sockfd, &client_addr, addr_len, loss_rate);
        // After chat mode ends, perform handshake
        perform_four_way_handshake(sockfd, &client_addr, addr_len, loss_rate);
    } else {
        // --- Data Transfer Loop ---
        printf("Connection is active. Ready to receive data...\n");

        // Wait for and process the filename packet
        ssize_t filename_bytes_received = recvfrom(sockfd, recv_buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &addr_len);
        sham_header_to_host(header);
        if (filename_bytes_received > (ssize_t)sizeof(struct sham_header)) {
            size_t filename_len = filename_bytes_received - sizeof(struct sham_header);
            memcpy(received_file_name, recv_buffer + sizeof(struct sham_header), filename_len);
            received_file_name[filename_len] = '\0';
            log_event("RCV DATA SEQ=%u LEN=%zu", header->seq_num, filename_len);
            expected_seq_num = header->seq_num + filename_len;
            send_ack(sockfd, &client_addr, addr_len, expected_seq_num, loss_rate);
        }

        received_file = fopen(received_file_name, "wb");
        if (received_file == NULL) {
            perror("Failed to open file for writing");
            close(sockfd);
            if (log_file) fclose(log_file);
            exit(EXIT_FAILURE);
        }

        while (1) {
            bytes_received = recvfrom(sockfd, recv_buffer, BUFFER_SIZE, 0, (struct sockaddr *)&client_addr, &addr_len);
            if (bytes_received <= 0) {
                continue;
            }
            sham_header_to_host(header);

            if (header->flags & FIN_FLAG) {
                printf("Received FIN from client. Terminating transfer...\n");
                break;
            }

            uint32_t payload_len = bytes_received - sizeof(struct sham_header);
            char *payload_data = recv_buffer + sizeof(struct sham_header);

            log_event("RCV DATA SEQ=%u LEN=%u", header->seq_num, payload_len);

            if (header->seq_num == expected_seq_num) {
                printf("Received in-order packet. Seq: %u\n", header->seq_num);
                fwrite(payload_data, 1, payload_len, received_file);
                expected_seq_num += payload_len;
                
                int sorted_packets = 0;
                while(sorted_packets < buffered_packets_count) {
                    int packet_found = 0;
                    for(int i = 0; i < WINDOW_SIZE * 2; i++) {
                        if (buffer_storage[i].received && buffer_storage[i].seq_num == expected_seq_num) {
                            printf("Processing buffered packet. Seq: %u\n", buffer_storage[i].seq_num);
                            fwrite(buffer_storage[i].data, 1, buffer_storage[i].data_len, received_file);
                            expected_seq_num += buffer_storage[i].data_len;
                            buffer_storage[i].received = 0;
                            buffered_packets_count--;
                            packet_found = 1;
                            break;
                        }
                    }
                    if (!packet_found) {
                        break;
                    }
                }
                
                send_ack(sockfd, &client_addr, addr_len, expected_seq_num, loss_rate);
            } else if (header->seq_num > expected_seq_num) {
                printf("Received out-of-order packet. Seq: %u. Buffering...\n", header->seq_num);
                int i;
                for (i = 0; i < WINDOW_SIZE * 2; i++) {
                    if (!buffer_storage[i].received) {
                        memcpy(buffer_storage[i].data, payload_data, payload_len);
                        buffer_storage[i].seq_num = header->seq_num;
                        buffer_storage[i].received = 1;
                        buffer_storage[i].data_len = payload_len;
                        buffered_packets_count++;
                        break;
                    }
                }
                send_ack(sockfd, &client_addr, addr_len, expected_seq_num, loss_rate);
            } else {
                // Duplicate packet - just send ACK
                printf("Received duplicate packet. Seq: %u. Sending ACK.\n", header->seq_num);
                send_ack(sockfd, &client_addr, addr_len, expected_seq_num, loss_rate);
            }
        }
        fclose(received_file);
        
        // Calculate and print MD5 checksum
        MD5_CTX md5_context;
        unsigned char md5_digest[MD5_DIGEST_LENGTH];
        
        // Initialize the MD5 context
        MD5_Init(&md5_context);
        
        // Open the received file for reading to calculate the hash
        received_file = fopen(received_file_name, "rb");
        if (received_file) {
            char file_buffer[1024];
            size_t bytes;
            // Read the file in chunks and update the MD5 context
            while ((bytes = fread(file_buffer, 1, sizeof(file_buffer), received_file)) != 0) {
                MD5_Update(&md5_context, file_buffer, bytes);
            }
            // Finalize the MD5 hash
            MD5_Final(md5_digest, &md5_context);
            fclose(received_file);
            
            // Print the MD5 hash in the required format
            printf("MD5: ");
            for (int i = 0; i < MD5_DIGEST_LENGTH; i++) {
                printf("%02x", md5_digest[i]);
            }
            printf("\n");
        }

        perform_four_way_handshake(sockfd, &client_addr, addr_len, loss_rate);
    }
    
    printf("Connection terminated.\n");
    close(sockfd);
    if (log_file) {
        fclose(log_file);
    }
    return 0;
}