struct message_str
{
    Message*  next; // For kernel use only.
    ...
    // Users should not modify anything above this line.
    int mac_protocol_type       // Newly added
    void *mac_hdr               // Newly added
    int network_protocol_type   // Newly added
    void *network_hdr           // Newly added
    int transport_protocol_type // Newly added
    void *transport_hdr         // Newly added
};

// Retrun the ProtocolType of a specific layer
int MESSAGE_GetMacProtocolType(Node* node, Message* msg)
{
	return msg->mac_protocol_type;
}

int MESSAGE_GetNetworkProtocolType(Node* node, Message* msg)
{
	return msg->network_protocol_type;
}

int MESSAGE_GetTransportProtocolType(Node* node, Message* msg)
{
	return msg->transport_protocol_type;
}

// Retrun the pointer to the ProtocolHeader for a specific layer of this packet
void* MESSAGE_GetMacHeader(Node* node, Message* msg)
{
	return msg->mac_hdr;
}

void* MESSAGE_GetNetworkHeader(Node* node, Message* msg)
{
	return msg->network_hdr;
}

void* MESSAGE_GetTransportHeader(Node* node, Message* msg)
{
	return msg->transport_hdr;
}

void tcp_output(
    Node *node,
    struct tcpcb *tp,
    UInt32 tcp_now,
    struct tcpstat *tcp_stat)
{
    ...
    // Send the packet to the IP model.
    // Remove the ipovly pseudo-header beforehand.
    {
        ...
        anEcnCapablePacket = TransportTcpSetEcnCapablePacket(node, tp, ti);
        TransportTcpSetEcnEchoAndCWR(node, tp, ti);

        TransportTcpTrace(node, tp, ti, "output");
        if (anEcnCapablePacket) {
            TransportTcpTrace(node, 0, 0, "ECT");
        }
        TransportTcpTrace(node, tp, ti, "variables");

        MESSAGE_RemoveHeader(node, msg, sizeof(struct ipovly), TRACE_TCP);

        memcpy(
            MESSAGE_ReturnPacket(msg),
            &ti->ti_t,
            sizeof (struct tcphdr));

        memcpy(
            MESSAGE_ReturnPacket(msg) + sizeof (struct tcphdr),
            (ti + 1), optlen);

        msg->transport_protocol_type = TransportProtocol_TCP; // Newly added
        msg->transport_hdr = MESSAGE_ReturnPacket(msg);       // Newly added

        // Hack solution for TCP trace.
        msg->headerSizes[msg->numberOfHeaders] -= sizeof(struct ipovly);
        msg->numberOfHeaders++;
        ...
    }
    ...
}

    ...
    switch (MESSAGE_GetTransportProtocolType(node, msg))
    {
        case TransportProtocol_UDP:
        {
            TransportUdpHeader* udphdr = (TransportUdpHeader *) MESSAGE_GetTransportHeader(node, msg);
             ...
        }
        case TransportProtocol_TCP:
        {
            struct tcphdr* tcphdr = (struct tcphdr *) MESSAGE_GetTransportHeader(node, msg);
            ...
        }
        case TransportProtocol_RSVP:
        default:
        {
            // NOOP
            break;
        }
    }
    ...