Refactoring: extract common functions

- Extract database connection healthcheck to function - Extract check for user existence in logins database to function Signed-off-by: Ondrej Vasko <[email protected]>
Vnet-as · Oct 5, 2021 · 83106f4 · 83106f4
1 parent 7243ffe
commit 83106f4
Showing 1 changed file with 50 additions and 54 deletions.
diff --git a/postfwd-anti-spam.plugin b/postfwd-anti-spam.plugin
@@ -272,6 +272,46 @@ mylog_info(
 # Setup initial time for flushing database records older than interval set in config file
 my $last_cache_flush = time;
 
+
+# Function: Test if database connection is still alive
+sub is_db_connection_alive {
+    my $rc = $dbh->ping;
+    if ( !$rc ) {
+        mylog_info(
+           "Database connection dropped (rc=$rc). Reconnecting to database."
+        );
+        $dbh = DBI->connect_cached(
+            $dsn,
+            $config{database}{userid},
+            $config{database}{password}, \%attr
+        ) or mylog_fatal( DBI->errstr );
+    }
+    return 1
+}
+
+# Function: Check if user exists in logins database
+sub user_exists {
+    my ($user) = @_;
+
+    my $check_user_existence_sth =
+      $dbh->prepare( $config_sql{check_user_existence_st} )
+      or do {
+          mylog_err( $dbh->errstr );
+          return 0;
+      };
+
+    my $row_count = $check_user_existence_sth->execute($user);
+    if ( $row_count == 0 ) {
+        if ( $check_user_existence_sth->err ) {
+            mylog_err( $check_user_existence_sth->errstr );
+        }
+        return 0;
+    }
+
+    return 1;
+}
+
+
 %postfwd_items_plugin = (
 
     'incr_client_country_login_count' => sub {
@@ -280,18 +320,8 @@ my $last_cache_flush = time;
         my ($result) = undef;
         $result->{incr_client_country_login_count} = 0;
 
-        # Test if database connection is still alive
-        my $rc = $dbh->ping;
-        if ( !$rc ) {
-            mylog_info(
-"Database connection dropped (rc=$rc). Reconnecting to database."
-            );
-            $dbh = DBI->connect_cached(
-                $dsn,
-                $config{database}{userid},
-                $config{database}{password}, \%attr
-            ) or mylog_fatal( DBI->errstr );
-        }
+        # Check if we still have DB connection
+        is_db_connection_alive();
 
         # Clear old records after flush interval expired
         if ( ( $last_cache_flush + $config{app}{db_flush_interval} ) < time ) {
@@ -417,34 +447,17 @@ my $last_cache_flush = time;
         my ($result) = undef;
         $result->{client_uniq_country_login_count} = 0;
 
-        # Test if database connection is still alive
-        my $rc = $dbh->ping;
-        if ( !$rc ) {
-            mylog_info(
-"Database connection dropped (rc=$rc). Reconnecting to database."
-            );
-            $dbh = DBI->connect_cached(
-                $dsn,
-                $config{database}{userid},
-                $config{database}{password}, \%attr
-            ) or mylog_fatal( DBI->errstr );
-        }
+        # Check if we still have DB connection
+        is_db_connection_alive();
 
         # Get sasl_username
         my $user = $request->{sasl_username};
         if ( !length $user || !($user) ) {
             return $result;
         }
 
-        # Check if user with given IP already has record
-        my $check_user_existence_sth =
-          $dbh->prepare( $config_sql{check_user_existence_st} )
-          or do { mylog_err( $dbh->errstr ); return $result; };
-        my $row_count = $check_user_existence_sth->execute($user);
-        if ( $row_count == 0 ) {
-            if ( $check_user_existence_sth->err ) {
-                mylog_err( $check_user_existence_sth->errstr );
-            }
+        # Check if user already exists, if not return from function
+        if (!user_exists($user)) {
             return $result;
         }
 
@@ -484,34 +497,17 @@ my $last_cache_flush = time;
         my ($result) = undef;
         $result->{client_ip_login_count} = 0;
 
-        # Test if database connection is still alive
-        my $rc = $dbh->ping;
-        if ( !$rc ) {
-            mylog_info(
-"Database connection dropped (rc=$rc). Reconnecting to database."
-            );
-            $dbh = DBI->connect_cached(
-                $dsn,
-                $config{database}{userid},
-                $config{database}{password}, \%attr
-            ) or mylog_fatal( DBI->errstr );
-        }
+        # Check if we still have DB connection
+        is_db_connection_alive();
 
         # Get sasl_username
         my $user = $request->{sasl_username};
         if ( !length $user || !($user) ) {
             return $result;
         }
 
-        # Check if user with given IP already has record
-        my $check_user_existence_sth =
-          $dbh->prepare( $config_sql{check_user_existence_st} )
-          or do { mylog_err( $dbh->errstr ); return $result; };
-        my $row_count = $check_user_existence_sth->execute($user);
-        if ( $row_count == 0 ) {
-            if ( $check_user_existence_sth->err ) {
-                mylog_err( $check_user_existence_sth->errstr );
-            }
+        # Check if user already exists, if not return from function
+        if (!user_exists($user)) {
             return $result;
         }