/*
 * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#include <doca_model.p4>
#include <doca_headers.p4>
#include <doca_externs.p4>
#include <doca_parser.p4>


/* Basic Connection Tracking example for SDN */

@nv_controller_metadata("packet_in")
struct connection_meta_t {
    bit<2> type;
    bit<2> _reserved;
    bit<28> zone;
}

/* Meta connection type */
enum bit<2> ct_type {
    PASS = 0,
    NEW = 1,
    RST = 2,
    FIN = 3
}

/*
   VF port to SW entity that manages the flow insertion,
   deletion and expiry times
*/
const bit<32> WIRE_PORT = 32w00;
const bit<32> HAIRPIN_PORT = 32w01;


control conn_track(
    inout nv_headers_t headers,
    in nv_standard_metadata_t std_meta,
    inout nv_empty_metadata_t user_meta,
    inout nv_empty_metadata_t pkt_out_meta
) {
    connection_meta_t ctrl_meta;

    /* 5-tuple matching for L4 TCP/UDP flows */
    NvDirectCounter(NvCounterType.PACKETS_AND_BYTES) table_t5_counter;

    action set_connection_id_ct_table_t5(bit<28> zone) {
        table_t5_counter.count();
        ctrl_meta.zone = zone;
        ctrl_meta.type = ct_type.PASS;
    }

    action no_action_ct_table_t5() {
        table_t5_counter.count();
    }

    table ct_table_t5 {
        key = {
            headers.ipv4.src_addr : exact;
            headers.ipv4.dst_addr : exact;
            headers.ipv4.protocol : exact;
            headers.tcp.src_port : exact;
            headers.tcp.dst_port : exact;
        }
        actions = {
            set_connection_id_ct_table_t5;
            no_action_ct_table_t5;
        }
        direct_counter = table_t5_counter;
        // on hit table ct_table_known
        // on miss table ct_table_tcp_miss
        default_action = no_action_ct_table_t5;
        size = 1048576;
    }

    /*
    * Known connections handling.
    * Precondition - must be a TCP packet
    *  RST:    TYPE_RST
    *  FIN:    TYPE_FIN
    *  FINRST: TYPE_RST
    *
    * Match: tcp.flags
    * Actions: meta.type, next PIPE
    * MISS: next PIPE
    */
    NvDirectCounter(NvCounterType.PACKETS_AND_BYTES) table_known_counter;

    action set_connection_type_ct_table_known(bit<2> type) {
        table_known_counter.count();
        ctrl_meta.type = type;
    }

    table ct_table_known {
        key = {
            headers.tcp.flags : exact;
        }
        actions = {
            set_connection_type_ct_table_known;
        }
        direct_counter = table_known_counter;
        default_action = set_connection_type_ct_table_known(ct_type.PASS);
        const entries = {
            ( 0x1) : set_connection_type_ct_table_known(ct_type.FIN);
            ( 0x4) : set_connection_type_ct_table_known(ct_type.RST);
            (/*NV_TCP_PROTOCOL,*/ 0x5) : set_connection_type_ct_table_known(ct_type.RST); /* RST+FIN */
        }
    }

    /*
    * Unknown TCP connections handling
    * Precondition - must be a TCP packet
    *  SYN: TYPE_NEW
    *
    * Match: tcp.flags
    * Actions: PIPE, meta.type=NEW
    * MISS: NEXT PIPE
    */
    NvDirectCounter(NvCounterType.PACKETS_AND_BYTES) table_tcp_miss_counter;

    action set_connection_type_ct_table_tcp(bit<2> type) {
        table_tcp_miss_counter.count();
        ctrl_meta.type = type;
    }

    action no_action_ct_table_tcp() {
        table_tcp_miss_counter.count();
    }

    table ct_table_tcp {
        key = {
            headers.tcp.flags : exact;
        }
        actions = {
            set_connection_type_ct_table_tcp;
            no_action_ct_table_tcp;
        }
        direct_counter = table_tcp_miss_counter;
        default_action = no_action_ct_table_tcp;
        const entries = {
            (0x02) : set_connection_type_ct_table_tcp(ct_type.NEW);  /* SYN=1 ACK=0 */
        }
    }

    /*
    * Unknown UDP connections handling
    *
    * Match: UDP
    * Actions: PIPE, meta.type=NEW
    * MISS: DROP
    */
    NvDirectCounter(NvCounterType.PACKETS_AND_BYTES) table_udp_miss_counter;

    action set_connection_type_ct_table_udp(bit<2> type) {
        table_udp_miss_counter.count();
        ctrl_meta.type = type;
    }

    action drop_ct_table_udp() {
        table_udp_miss_counter.count();
        nv_drop();
    }

    table ct_table_udp {
        key = {
            std_meta.l4_type : exact;
        }
        actions = {
            set_connection_type_ct_table_udp;
            drop_ct_table_udp();
        }
        direct_counter = table_udp_miss_counter;
        default_action = drop_ct_table_udp();
        const entries = {
            (L4_TYPE_UDP) : set_connection_type_ct_table_udp(ct_type.NEW);  /* SYN=1 ACK=0 */
        }
    }

    /* user should add entries that correspond to the wire ports
     * A hit means this is an RX packet, miss means a TX packet
    */
    table direction_table {
        key = {
            std_meta.ingress_port : exact;
        }
        actions = {
            NoAction;
        }
        default_action = NoAction;
        const entries = {
            (WIRE_PORT) : NoAction();
        }
    }    

    apply {
        ctrl_meta.type = ct_type.PASS;
        ctrl_meta._reserved = 0;
        ctrl_meta.zone = 0;

        if (direction_table.apply().hit) {
            if (std_meta.is_l4_ok == 1w1) {
                if (ct_table_t5.apply().hit) {
                    if (ct_table_known.apply().hit) {
                        nv_send_to_controller(ctrl_meta);
                    }
                } else if (headers.tcp.isValid()) {
                    if (ct_table_tcp.apply().hit) {
                        nv_send_to_controller(ctrl_meta);
                    }
                } else {
                    if (ct_table_udp.apply().hit) {
                        nv_send_to_controller(ctrl_meta);
                    }               
                }
            }
        }
        nv_send_to_port(HAIRPIN_PORT);
    }
}

// Instantiate the top-level DOCA Rx package
NvDocaPipeline(
    nv_fixed_parser(),
    conn_track()
) main;