done

server.cpp

@@ -34,15 +34,36 @@ void load_env(const std::string& path) {
     }
 }
 
+// hex to byte helper function
+int hex_to_bytes_upper(const char* hex, unsigned char* out, size_t out_size)
+{
+    size_t i = 0;
+    while (hex[0] && hex[1]) {
+        if (i >= out_size) return -1;
+        unsigned char h = hex[0];
+        unsigned char l = hex[1];
+        int hi = (h <= '9') ? (h - '0') : (h - 'A' + 10);
+        int lo = (l <= '9') ? (l - '0') : (l - 'A' + 10);
+        // minimal sanity check
+        if (hi < 0 || hi > 15 || lo < 0 || lo > 15)
+            return -1;
+        out[i++] = (unsigned char)((hi << 4) | lo);
+        hex += 2;
+    }
+    return (int)i;
+}
+
 // Define the decryption function
 int stream_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char *key, unsigned char *iv, unsigned char *plaintext) {
     /* Declare cipher context */
     EVP_CIPHER_CTX *ctx;
-
     int len, plaintext_len;
 
     /* Create and initialise the context */
     ctx = EVP_CIPHER_CTX_new();
+    if (!ctx) {
+        ERR_print_errors_fp(stderr);
+    }
 
     /* Initialise the decryption operation. */
     if (EVP_DecryptInit_ex(ctx, EVP_aes_256_ctr(), NULL, key, iv) != 1){
@@ -62,7 +83,7 @@ int stream_decrypt(unsigned char *ciphertext, int ciphertext_len, unsigned char
     /* Clean up */
     EVP_CIPHER_CTX_free(ctx);
 
-    return plaintext_len;
+    return len + plaintext_len;
 }
 
 int main(void){
@@ -74,11 +95,20 @@ int main(void){
     const int server_port = std::atoi(std::getenv("SERVER_PORT"));
     char client_message[2048];
     char server_message[2048];
-    const char * custom_message="Server: Hello from server, ";
+    const char * custom_message="Server: Hello from server, message recieved!\n";
 
     // decryption parameters
     const char *inital_vector = std::getenv("INITIAL_VECTOR");
     const char *secret_key = std::getenv("SECRET_KEY");
+    unsigned char key_bytes[32], iv_bytes[16];
+    
+    int key_len = hex_to_bytes_upper(secret_key, key_bytes, sizeof(key_bytes));
+    int iv_len  = hex_to_bytes_upper(inital_vector, iv_bytes, sizeof(iv_bytes));
+
+    if (key_len != 32 || iv_len != 16) {
+        fprintf(stderr, "Invalid key/IV size for AES-256\n");
+        return 1;
+    }
 
     // debug
     printf("Server starting at IP: %s, Port: %d\n", server_ip, server_port);
@@ -127,38 +157,40 @@ int main(void){
     memset(client_message, 0, sizeof(client_message));
 
     // Receive client's message
-    while (strcmp(client_message, "\\exit\n") != 0) {
-        
+    while (1) {
         // clean exising buffer
         memset(client_message, 0, sizeof(client_message));
 
-        recv(client_socket, client_message, sizeof(client_message), 0);
+        ssize_t n = recv(client_socket, client_message, sizeof(client_message), 0);
+        if (n == 0) break;
+        if (n < 0) { perror("recv"); break; }
 
         // decrypt client message
-        char plaintext[2048];
-        int plaintext_len = stream_decrypt((unsigned char*)client_message, strlen(client_message), (unsigned char*)secret_key, (unsigned char*)inital_vector, (unsigned char*)plaintext);
+        unsigned char plaintext[2048];
+        int plaintext_len = stream_decrypt((unsigned char*)client_message, (int)n, key_bytes, iv_bytes, plaintext);
+        if (plaintext_len < 0) {
+            fprintf(stderr, "decrypt failed\n");
+            break;
+        }
 
-        printf("Msg from client: %.*s", plaintext_len, plaintext);
+        if (plaintext_len >= (int)sizeof(plaintext)) plaintext_len = (int)sizeof(plaintext) - 1;
+        plaintext[plaintext_len] = '\0';
 
-        // assign client message to plaintext
-        memcpy(client_message, plaintext, plaintext_len);
+        if (strcmp((char*)plaintext, "\\exit\n") == 0) break;
+
+        printf("Msg from client: %s", plaintext);
 
         // Respond to client
         // prepare server message
-        memset(server_message, 0, sizeof(server_message));
-        ;
+        memcpy(server_message, custom_message, strlen(custom_message));
 
-        // add my name in the back as response.
-        std::snprintf(server_message, sizeof(server_message), "%s%s", custom_message, strtok(client_message, ":"));
-        
-        if (send(client_socket, server_message, sizeof(server_message), 0) == -1) {
-            perror("Failed to send message");
-            close(client_socket);
-            close(server_socket);
-            return 1;
+        size_t reply_len = strlen(server_message);
+        if (send(client_socket, server_message, reply_len, 0) == -1) {
+            perror("Send failed");
+            break;
         }
-        printf("Response sent to client: %s\n", server_message);
 
+        printf("Response sent to client: %s\n", server_message);
     }
     // Close the socket
     close(client_socket);