I want to split received socket data when encountered with a null byte: b'\0' in Python

Question

I want to split received socket data when encountered with a null byte: b'\0'. However, it is not splitting properly. In client side code, the b'\0' character is used as a delimiter. It is used to separate different parts of the message. The server side code then splits the data when encountered with a b'\0'. However, it mixes things up. Part of one message is inside another part. What seems to be happening and how to fix this?

I have another client and server side code similar to one I have attached below. In that code it splits correctly sometimes but occasionally doesn't split. It almost never splits properly if the encrypted message was longer. Please help. I have no clue.

Here's my server side code:

import struct
import socket
import threading
import random

# Define the server's IP address and port
SERVER_IP = '0.0.0.0'
SERVER_PORT = 55555

# Create a socket for the server
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind((SERVER_IP, SERVER_PORT))
server_socket.listen()

# Dictionary to store client sockets based on their public keys
client_sockets = {}

def handle_client(client_socket):
    try:
        while True:
            # Receive a message
            message_type, received_message = receive_message(client_socket)
            
            if message_type == 3:
                print("Received message type", message_type)

            
    except Exception as e:
        print(f"Error handling client: {e}")
    
    
def receive_message(socket):
    # Receive the header
    print("Received message")
    header = socket.recv(8)

    # Unpack the message type and length
    message_type, length = struct.unpack('!II', header)

    # Receive the message data
    message_data = socket.recv(length)

    if message_type == 3:
        #print("message_data: ",message_data)
        encrypted_prompt, serialized_signing_pub_key, nonce, another_serialized_pub_key = message_data.split(b'\0', 3)
        # Process the received data for message type 3
        
        print("encrypted prompt: ", encrypted_prompt)
        print("serialized signing pub key: ", serialized_signing_pub_key)
        print("nonce: ", nonce)
        print("another_serialized_pub_key: ", another_serialized_pub_key
            )
        return message_type, message_data
    
    else:
        # Handle other message types as needed
        print("Unsupported type)")
        return message_type, message_data


while True:
    # Accept a client connection
    client_socket, client_address = server_socket.accept()

    # Create a new thread to handle the client
    client_thread = threading.Thread(target=handle_client, args=(client_socket,))
    client_thread.start()

Here's my client side code:

from ast import While
from inspect import signature
import socket
import struct 
import threading
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.asymmetric import x25519, ed25519
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import serialization
from os import urandom
import time

shared_key = b'\xa1b\x0f\xc3\xf0e"\xa3\xc8\x95v\xa5\xf8nK\xf3\x1c9i\x88\xe9\xa3\\\xa9\xfa\xc4\xfeH\x1b\x91\xd1p'

def generate_key_pair():
    private_key = x25519.X25519PrivateKey.generate()
    public_key = private_key.public_key()
    return private_key, public_key

def sign_and_encrypt_message(signing_private_key, shared_key, nonce, prompt):
    prompt = prompt.encode('utf-8')
    # Sign the plaintext message
    signature = signing_private_key.sign(prompt)

    # Encrypt the message along with the signature
    cipher = Cipher(algorithms.AES(shared_key), modes.CTR(nonce), backend=default_backend())
    encryptor = cipher.encryptor()
    ciphertext = encryptor.update(prompt + signature) + encryptor.finalize()

    return ciphertext

def send_message(socket, message_type, serialized_pub_key=None, serialized_signing_pub_key=None, encrypted_prompt=None, nonce=None, another_serialized_pub_key=None):
    # Pack the message type and length
    header = struct.pack('!II', message_type, len(serialized_pub_key))

    if message_type == 3:
        print("encrypted prompt: ", encrypted_prompt)
        print("serilized signing pub key: ", serialized_signing_pub_key)
        print("nonce: ", nonce)
        print("another serialized pub key", another_serialized_pub_key)        
        
        header = struct.pack('!II', message_type, len(encrypted_prompt) + len(serialized_signing_pub_key) + len(nonce) + len(another_serialized_pub_key) + 3)
        socket.sendall(header)
        combined_data = encrypted_prompt + b'\0' + serialized_signing_pub_key + b'\0' + nonce + b'\0' + another_serialized_pub_key
        socket.sendall(combined_data)
    else:
        # Handle other message types as needed
        print("Unsupported message type")

# Create a client socket
client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client_socket.connect(('44.216.25.153', 55555))

# Generate client key pair
private_key, public_key = generate_key_pair()
# client generates keys
signing_private_key = ed25519.Ed25519PrivateKey.generate()
signing_public_key = signing_private_key.public_key()

another_private_key, another_public_key = generate_key_pair()

# Serialize public key to bytes
serialized_pub_key = public_key.public_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PublicFormat.Raw
)

another_serialized_pub_key = another_public_key.public_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PublicFormat.Raw
)

print(serialized_pub_key)

while True:
    # Get user input or any other data to send
    user_input = input("Enter a message to send: ")        
    serialized_signing_pub_key = signing_public_key.public_bytes(
    encoding=serialization.Encoding.Raw,
    format=serialization.PublicFormat.Raw
)
    prompt = user_input
    nonce = urandom(16)
    encrypted_prompt = sign_and_encrypt_message(signing_private_key, shared_key, nonce, prompt)
    # Sending a regular message
    prompt = user_input  
    send_message(client_socket, 3, serialized_pub_key, serialized_signing_pub_key,  encrypted_prompt, nonce, another_serialized_pub_key)

Here's the server output:

~$ python3 server1.py
Received message
encrypted prompt:  b"\xe9I\xfd!\xdb\xf9\xe3\xd4\xe9\x97\x83{\xbd\xab\xacq\x7f\xdc\xb9T\xda\x8dG-4jZ\x03\xa4\tC\xc1\xa6h\x15$\x1dH\x82\xd1@,\x13WW\xcc'\xdd\x01\x12\x87#Vn#+\x80\xf8\xbd\xfc\xf4Xm\n\x18\x0e"
serialized signing pub key:  b'\x83\x99B\x86I%\xb4D\xfaq\x9ds\xe7\x14\xd8\xeb@\xde\xc7'
nonce:  b'Jg+Y\xf0\x96\xba\xf0\x18}\xb6)'
another_serialized_pub_key:  b'\xae\xb0\x87\x9a\xd0z\x10\x167\xe7\xb4\xbcj\xba%\xf0\x00\xd6\\\x98\xbe5\x8bpO\x86L!Er\xae\xb0\x8c\xacO\x08\x9d\x9a6\n\xcf"TmJ\xbf\t\xfc-'
Received message type 3
Received message

Here's the client output:

~$ python3 client1.py 
b'\xd1\x8b\xa0\xab\xa1\xd2~\xbc\x00"Z\\\xd4=\x0cd\x9c!:\xa6\xa7\xe1\xe2d\xad/M;\xed\x9f\x07\x1e'
Enter a message to send: hi
encrypted prompt:  b"\xe9I\xfd!\xdb\xf9\xe3\xd4\xe9\x97\x83{\xbd\xab\xacq\x7f\xdc\xb9T\xda\x8dG-4jZ\x03\xa4\tC\xc1\xa6h\x15$\x1dH\x82\xd1@,\x13WW\xcc'\xdd\x01\x12\x87#Vn#+\x80\xf8\xbd\xfc\xf4Xm\n\x18\x0e"
serilized signing pub key:  b'\x83\x99B\x86I%\xb4D\xfaq\x9ds\xe7\x14\xd8\xeb@\xde\xc7\x00Jg+Y\xf0\x96\xba\xf0\x18}\xb6)'
nonce:  b'\xae\xb0\x87\x9a\xd0z\x10\x167\xe7\xb4\xbcj\xba%\xf0'
another serialized pub key b'\xd6\\\x98\xbe5\x8bpO\x86L!Er\xae\xb0\x8c\xacO\x08\x9d\x9a6\n\xcf"TmJ\xbf\t\xfc-'