/*
 * 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>

/*
 * VxLAN tunnel gateway
 * This application allows the user have a customized tunnel gateway, that can stitch VxLAN
 * packets across tenant domains. End point traffic destined to local bare metal hosts can 
 * be decapsulated and forwarded, while gateway traffic can be decapsulated, then
 * encapsulated back to the wire. This program can be easily extended to be a gateway across
 * different tunnel types as well.
 */

/*
* Table sizes.
*/
const bit<32> DECAP_TABLE_SIZE = 32768;
const bit<32> ENCAP_TABLE_SIZE = 32768;

/* The directionality is based on network to host
 * The user will configure the P4 port IDs in the OVS configuration
 */
const bit<32> WIRE_PORT = 32w0;

struct metadata_t {
    bit<1> was_decapped;
}

struct headers_t {
    NV_FIXED_HEADERS
}

parser packet_parser(packet_in packet, out headers_t headers) {
    NV_FIXED_PARSER(packet, headers)
 }

/**
* This control performs the overlay policy including L2 encap with VxLAN
*/
control overlay_encap(
    inout headers_t headers,
    in nv_standard_metadata_t std_meta,
    inout metadata_t user_meta,
    inout nv_empty_metadata_t pkt_out_meta
) {
    NvDirectCounter(NvCounterType.PACKETS_AND_BYTES) encap_counter;

    action deny_encap_v4_table() {
        encap_counter.count();
        nv_drop();
    }

    action vxlan_v4_encap_encap_v4_table(nv_mac_addr_t underlay_src_mac, nv_mac_addr_t underlay_dst_mac,
                 nv_ipv4_addr_t underlay_sip, nv_ipv4_addr_t underlay_dip, bit<24> vni, nv_logical_port_t port) {
        nv_set_vxlan_v4_underlay(headers, false, underlay_dst_mac, underlay_src_mac, 0, underlay_sip, underlay_dip, vni);
        encap_counter.count();
        nv_send_to_port(port);
    }
  
    table encap_v4_table {
        key = {
            headers.ipv4.dst_addr : exact;
        }
        actions = {
            vxlan_v4_encap_encap_v4_table;
            deny_encap_v4_table;
        }
        size = ENCAP_TABLE_SIZE;
        default_action = deny_encap_v4_table;
        direct_counter = encap_counter;
    }

    apply {
       if (headers.ipv4.isValid() && (user_meta.was_decapped == 1)) {
            encap_v4_table.apply();
       }
    }
}


/**
* This control is for packets from wire to host (RX)
* and includes policy for L2 decap
*/
control decap_flow(
    inout headers_t headers,
    in nv_standard_metadata_t std_meta,
    inout metadata_t user_meta,
    inout nv_empty_metadata_t pkt_out_meta
) {
    NvDirectCounter(NvCounterType.PACKETS_AND_BYTES) decap_counter;

    action deny_decap_v4_table() {
        decap_counter.count();
        nv_drop();
    }

    action decap_decap_v4_table() {
        decap_counter.count();
        nv_l2_decap(headers);
        user_meta.was_decapped = 1;
    }

    action to_port_decap_v4_table(nv_logical_port_t port) {
        decap_counter.count();
        nv_send_to_port(port);
    }

    action decap_to_port_decap_v4_table(nv_logical_port_t port) {
        decap_counter.count();
        user_meta.was_decapped = 1;
        nv_l2_decap(headers);
        nv_send_to_port(port);
    }

    action no_action_decap_v4_table() {
        decap_counter.count();
    }

    table decap_v4_table {
        key = {
            headers.vxlan.vni : exact;
        }
        actions = {
            decap_decap_v4_table;
            to_port_decap_v4_table;
            decap_to_port_decap_v4_table;
            deny_decap_v4_table;
            no_action_decap_v4_table;
        }
        size = DECAP_TABLE_SIZE;
        direct_counter = decap_counter;
        default_action = deny_decap_v4_table;
    }

    apply {
        if (headers.vxlan.isValid()) {
            decap_v4_table.apply();
        }
    }
}

control gateway(
    inout headers_t headers,
    in nv_standard_metadata_t std_meta,
    inout metadata_t user_meta,
    inout nv_empty_metadata_t pkt_out_meta
) {
    overlay_encap() over;
    decap_flow() decap;

    /* 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 {
        user_meta.was_decapped = 0;
        if (direction_table.apply().hit) {
            decap.apply(headers, std_meta, user_meta, pkt_out_meta);
        }
        else {
            over.apply(headers, std_meta, user_meta, pkt_out_meta);
        }
    }
}

NvDocaPipeline(
    packet_parser(),
    gateway()
) main;