Fix(dft+u): validate occupation tags and MPI counts

Author: liu-687

source/source_estate/module_charge/charge_mixing.cpp

@@ -1,6 +1,7 @@
 #include "charge_mixing.h"
 
 #include <functional>
+#include <limits>
 
 #include "source_io/module_parameter/parameter.h"
 #include "source_base/module_mixing/broyden_mixing.h"
@@ -361,7 +362,11 @@ void Charge_Mixing::mix_uom(std::vector<double>& uom_in, std::vector<double>& uo
 #ifdef __MPI
     // Synchronize mixed uom across all ranks to prevent divergence
     // after multiple mixing steps.
-    Parallel_Common::bcast_double(uom_in.data(), uom_in.size());
+    if (uom_in.size() > static_cast<std::size_t>(std::numeric_limits<int>::max()))
+    {
+        ModuleBase::WARNING_QUIT("Charge_Mixing::mix_uom", "UOM buffer is too large for MPI broadcast count");
+    }
+    Parallel_Common::bcast_double(uom_in.data(), static_cast<int>(uom_in.size()));
 #endif
     return;
 }

source/source_pw/module_pwdft/dftu_base.cpp

@@ -6,10 +6,11 @@
 #include "source_base/timer.h"
 #include "source_io/module_parameter/parameter.h"
 
-#include <cstdlib>
 #include <cstring>
 #include <fstream>
+#include <limits>
 #include <sstream>
+#include <stdexcept>
 #include <vector>
 
 // local inline helpers for eigenvalue calculation (JacobiRotate, CalculateEigenvalues)
@@ -553,7 +554,32 @@ void Plus_U_Base::read_occup_m(const UnitCell& ucell,
         }
         else
         {
-            value = std::atoi(token.c_str() + tag.size());
+            const std::string number = token.substr(tag.size());
+            if (number.empty())
+            {
+                return false;
+            }
+
+            std::size_t consumed = 0;
+            try
+            {
+                const long parsed = std::stol(number, &consumed, 10);
+                if (consumed != number.size()
+                    || parsed < static_cast<long>(std::numeric_limits<int>::min())
+                    || parsed > static_cast<long>(std::numeric_limits<int>::max()))
+                {
+                    return false;
+                }
+                value = static_cast<int>(parsed);
+            }
+            catch (const std::invalid_argument&)
+            {
+                return false;
+            }
+            catch (const std::out_of_range&)
+            {
+                return false;
+            }
         }
         return static_cast<bool>(ifdftu);
     };
@@ -574,8 +600,12 @@ void Plus_U_Base::read_occup_m(const UnitCell& ucell,
             break;
         }
 
-        if (read_tagged_int(word, "Atom=", iat))
+        if (word.compare(0, 5, "Atom=") == 0)
         {
+            if (!read_tagged_int(word, "Atom=", iat) || iat < 1 || iat > ucell.nat)
+            {
+                ModuleBase::WARNING_QUIT("Plus_U_Base::read_occup_m", "WRONG ATOM INDEX IN LOCAL OCCUPATION NUMBER MATRIX FROM Plus_U FILE");
+            }
             iat -= 1;
             ifdftu >> word;