I am creating a socks 5 proxy server for TCP and UDP in python 3

Question

here is my code in python 3. It is working properly for TCP But Not For UDP as per brooktest

I want this code to be working for UDP also as it is working for TCP

This implementation handles both TCP and UDP connections by checking the SOCKS5 command in the client request message. If the command is TCP CONNECT, it connects to the remote server and starts relaying data between the client and the remote server.

import socket


def handle_tcp(sock, remote_addr):
    # Connect to the remote server
    remote_sock = socket.create_connection(remote_addr)

    # Send a success response to the client
    reply = b'\x05\x00\x00\x01' + socket.inet_aton('0.0.0.0') + b'\x00\x00'
    sock.sendall(reply)

    # Start relaying data between the client and the remote server
    try:
        while True:
            data = sock.recv(4096)
            if not data:
                break
            remote_sock.sendall(data)
            data = remote_sock.recv(4096)
            if not data:
                break
            sock.sendall(data)
    finally:
        remote_sock.close()
        sock.close()


def handle_udp(sock, data):
    # Modify data as needed here
    return data


def handle_request(sock):
    # Read initial handshake message from client
    version, nmethods = sock.recv(2)
    assert version == 5, "Invalid SOCKS version"
    assert sock.recv(nmethods) == b'\x00' * nmethods, "Unsupported authentication method"

    # Send server method selection message
    sock.sendall(b'\x05\x00')

    # Read client request message
    version, cmd, _, addrtype = sock.recv(4)
    assert version == 5, "Invalid SOCKS version"
    assert cmd in (1, 3), "Unsupported command"
    if cmd == 1:
        assert addrtype == 1, "Unsupported address type"
        target_addr = socket.inet_ntoa(sock.recv(4))
        target_port = int.from_bytes(sock.recv(2), 'big')
        handle_tcp(sock, (target_addr, target_port))
    elif cmd == 3:
        assert addrtype == 3, "Unsupported address type"
        addr_len = ord(sock.recv(1))
        target_addr = sock.recv(addr_len)
        target_port = sock.recv(2)

        # Send server response message
        reply = b'\x05\x00\x00\x03' + bytes([addr_len]) + target_addr + target_port
        sock.sendall(reply)

        # Create new UDP socket and send/receive data
        udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_sock.connect((target_addr, int.from_bytes(target_port, 'big')))
        while True:
            data, addr = udp_sock.recvfrom(4096)
            data = handle_udp(udp_sock, data)
            sock.sendto(data, addr)


if __name__ == '__main__':
    # Listen on all interfaces on port 1080
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind(('0.0.0.0', 1080))
    server.listen(10)
    while True:
        client, addr = server.accept()
        handle_request(client)