Added pre-processor macros to prng.c to prevent referencing to non-ex…

…isting functions in case of lacking SSE4 or AVX2 support
martijnvanbrummelen · Sep 12, 2024 · 7ccd95a · 7ccd95a
1 parent 32542d2
commit 7ccd95a
Showing 1 changed file with 24 additions and 4 deletions.
diff --git a/src/prng.c b/src/prng.c
@@ -408,10 +408,10 @@ int nwipe_rc4_prng_init( NWIPE_PRNG_INIT_SIGNATURE )
 // The main RC4 PRNG read function with AVX2 and SSE4.2 detection
 int nwipe_rc4_prng_read( NWIPE_PRNG_READ_SIGNATURE )
 {
-    u8* restrict bufpos = buffer;
-    size_t words = count / SIZE_OF_RC4_PRNG;
+    u8* restrict bufpos = buffer;  // Buffer position pointer
+    size_t words = count / SIZE_OF_RC4_PRNG;  // Number of 4096-byte blocks
 
-    // Check if the CPU supports AVX2 and SSE4.2
+    // Check if the CPU supports AVX2 or SSE4.2
     int use_avx2 = check_avx2_support();
     int use_sse4 = check_sse42_support();
 
@@ -420,13 +420,23 @@ int nwipe_rc4_prng_read( NWIPE_PRNG_READ_SIGNATURE )
     {
         if( use_avx2 )
         {
+#if defined( __AVX2__ )
             // Use AVX2-optimized version
             rc4_genrand_4096_to_buf_avx2( (rc4_state_t*) *state, bufpos );
+#else
+            // Fallback to generic version if AVX2 is not compiled
+            rc4_genrand_4096_to_buf( (rc4_state_t*) *state, bufpos );
+#endif
         }
         else if( use_sse4 )
         {
+#if defined( __SSE4_2__ )
             // Use SSE4.2-optimized version
             rc4_genrand_4096_to_buf_sse42( (rc4_state_t*) *state, bufpos );
+#else
+            // Fallback to generic version if SSE4.2 is not compiled
+            rc4_genrand_4096_to_buf( (rc4_state_t*) *state, bufpos );
+#endif
         }
         else
         {
@@ -444,21 +454,31 @@ int nwipe_rc4_prng_read( NWIPE_PRNG_READ_SIGNATURE )
 
         if( use_avx2 )
         {
+#if defined( __AVX2__ )
             // Use AVX2-optimized version
             rc4_genrand_4096_to_buf_avx2( (rc4_state_t*) *state, temp_output );
+#else
+            // Fallback to generic version if AVX2 is not compiled
+            rc4_genrand_4096_to_buf( (rc4_state_t*) *state, temp_output );
+#endif
         }
         else if( use_sse4 )
         {
+#if defined( __SSE4_2__ )
             // Use SSE4.2-optimized version
             rc4_genrand_4096_to_buf_sse42( (rc4_state_t*) *state, temp_output );
+#else
+            // Fallback to generic version if SSE4.2 is not compiled
+            rc4_genrand_4096_to_buf( (rc4_state_t*) *state, temp_output );
+#endif
         }
         else
         {
             // Fallback to generic version
             rc4_genrand_4096_to_buf( (rc4_state_t*) *state, temp_output );
         }
 
-        // Copy the remaining bytes
+        // Copy the remaining bytes to the buffer
         memcpy( bufpos, temp_output, remain );
     }