Adjust WG plugin

Zadamsa · Oct 8, 2024 · 107401e · 107401e
1 parent 80cfca3
commit 107401e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/process/wg.cpp b/process/wg.cpp
@@ -70,32 +70,34 @@ ProcessPlugin *WGPlugin::copy()
    return new WGPlugin(*this);
 }
 
-int WGPlugin::post_create(Flow &rec, const Packet &pkt)
+ProcessPlugin::FlowAction WGPlugin::post_create(Flow &rec, const Packet &pkt)
 {
    if (pkt.ip_proto == IPPROTO_UDP) {
       add_ext_wg(reinterpret_cast<const char *>(pkt.payload), pkt.payload_len, pkt.source_pkt, rec);
+      return ProcessPlugin::FlowAction::GET_ALL_DATA;
    }
 
-   return 0;
+   return ProcessPlugin::FlowAction::NO_PROCESS;
 }
 
-int WGPlugin::pre_update(Flow &rec, Packet &pkt)
+ProcessPlugin::FlowAction WGPlugin::pre_update(Flow &rec, Packet &pkt)
 {
    RecordExtWG *vpn_data = (RecordExtWG *) rec.get_extension(RecordExtWG::REGISTERED_ID);
    if (vpn_data != nullptr && vpn_data->possible_wg) {
       bool res = parse_wg(reinterpret_cast<const char *>(pkt.payload), pkt.payload_len, pkt.source_pkt, vpn_data);
       // In case of new flow, flush
       if (flow_flush) {
          flow_flush = false;
-         return FLOW_FLUSH_WITH_REINSERT;
+         return ProcessPlugin::FlowAction::FLUSH_WITH_REINSERT;
       }
       // In other cases, when WG was not detected
       if (!res) {
          vpn_data->possible_wg = 0;
+         return ProcessPlugin::FlowAction::NO_PROCESS;
       }
+      return ProcessPlugin::FlowAction::GET_ALL_DATA;
    }
-
-   return 0;
+   return ProcessPlugin::FlowAction::NO_PROCESS;
 }
 
 void WGPlugin::pre_export(Flow &rec)

diff --git a/process/wg.hpp b/process/wg.hpp
@@ -158,8 +158,8 @@ class WGPlugin : public ProcessPlugin
    RecordExt *get_ext() const { return new RecordExtWG(); }
    ProcessPlugin *copy();
 
-   int post_create(Flow &rec, const Packet &pkt);
-   int pre_update(Flow &rec, Packet &pkt);
+   ProcessPlugin::FlowAction post_create(Flow &rec, const Packet &pkt);
+   ProcessPlugin::FlowAction pre_update(Flow &rec, Packet &pkt);
    void pre_export(Flow &rec);
    void finish(bool print_stats);