WAF: Minor Error Handling Improvements (#29108)

Automattic · Mar 7, 2023 · fc9b6b5 · fc9b6b5
1 parent 128facd
commit fc9b6b5
Show file tree

Hide file tree

Showing 11 changed files with 526 additions and 128 deletions.
diff --git a/projects/packages/waf/changelog/enhance-waf-error-handling b/projects/packages/waf/changelog/enhance-waf-error-handling
@@ -0,0 +1,5 @@
+Significance: patch
+Type: changed
+Comment: Enhance error handling behind the scenes.
+
+
diff --git a/projects/packages/waf/src/class-rest-controller.php b/projects/packages/waf/src/class-rest-controller.php
@@ -17,6 +17,8 @@
 class REST_Controller {
 	/**
 	 * Register REST API endpoints.
+	 *
+	 * @return void
 	 */
 	public static function register_rest_routes() {
 		register_rest_route(
@@ -52,29 +54,29 @@ public static function register_rest_routes() {
 
 	/**
 	 * Update rules endpoint
+	 *
+	 * @return WP_REST_Response|WP_Error
 	 */
 	public static function update_rules() {
-		$success = true;
-		$message = 'Rules updated succesfully';
-
 		try {
 			Waf_Rules_Manager::generate_automatic_rules();
 			Waf_Rules_Manager::generate_rules();
-		} catch ( \Exception $e ) {
-			$success = false;
-			$message = $e->getMessage();
+		} catch ( Waf_Exception $e ) {
+			return $e->get_wp_error();
 		}
 
 		return rest_ensure_response(
 			array(
-				'success' => $success,
-				'message' => $message,
+				'success' => true,
+				'message' => __( 'Rules updated succesfully', 'jetpack-waf' ),
 			)
 		);
 	}
 
 	/**
 	 * WAF Endpoint
+	 *
+	 * @return WP_REST_Response
 	 */
 	public static function waf() {
 		return rest_ensure_response( Waf_Runner::get_config() );
@@ -84,7 +86,8 @@ public static function waf() {
 	 * Update WAF Endpoint
 	 *
 	 * @param WP_REST_Request $request The API request.
-	 * @return WP_REST_Response
+	 *
+	 * @return WP_REST_Response|WP_Error
 	 */
 	public static function update_waf( $request ) {
 		// Automatic Rules Enabled
@@ -112,7 +115,11 @@ public static function update_waf( $request ) {
 			update_option( Waf_Runner::SHARE_DATA_OPTION_NAME, (bool) $request[ Waf_Runner::SHARE_DATA_OPTION_NAME ] );
 		}
 
-		Waf_Runner::update_waf();
+		try {
+			Waf_Runner::update_waf();
+		} catch ( Waf_Exception $e ) {
+			return $e->get_wp_error();
+		}
 
 		return self::waf();
 	}

diff --git a/projects/packages/waf/src/class-waf-initializer.php b/projects/packages/waf/src/class-waf-initializer.php
@@ -49,33 +49,37 @@ public static function init() {
 	}
 
 	/**
-	 * On module activation set up waf mode
+	 * Activate the WAF on module activation.
 	 *
-	 * @return bool|WP_Error True of the WAF activation was successful, WP_Error otherwise.
+	 * @return bool|WP_Error True if the WAF activation is successful, WP_Error otherwise.
 	 */
 	public static function on_activation() {
 		update_option( Waf_Runner::MODE_OPTION_NAME, 'normal' );
 		add_option( Waf_Rules_Manager::AUTOMATIC_RULES_ENABLED_OPTION_NAME, false );
 
-		$waf_activated = Waf_Runner::activate();
-		if ( is_wp_error( $waf_activated ) ) {
-			return $waf_activated;
-		}
-
 		try {
+			Waf_Runner::activate();
 			( new Waf_Standalone_Bootstrap() )->generate();
-		} catch ( \Exception $e ) {
-			return new WP_Error( 'waf_activation_failed', $e->getMessage() );
+		} catch ( Waf_Exception $e ) {
+			return $e->get_wp_error();
 		}
 
 		return true;
 	}
 
 	/**
-	 * On module deactivation, unset waf mode
+	 * Deactivate the WAF on module deactivation.
+	 *
+	 * @return bool|WP_Error True if the WAF deactivation is successful, WP_Error otherwise.
 	 */
 	public static function on_deactivation() {
-		Waf_Runner::deactivate();
+		try {
+			Waf_Runner::deactivate();
+		} catch ( Waf_Exception $e ) {
+			return $e->get_wp_error();
+		}
+
+		return true;
 	}
 
 	/**
@@ -131,23 +135,22 @@ public static function check_for_waf_update() {
 			if ( ! method_exists( Waf_Constants::class, 'define_mode' ) ) {
 				try {
 					( new Waf_Standalone_Bootstrap() )->generate();
-				} catch ( \Exception $e ) {
-					return new WP_Error( 'waf_update_failed', $e->getMessage() );
+				} catch ( Waf_Exception $e ) {
+					return $e->get_wp_error();
 				}
-				return true;
 			}
 
 			Waf_Constants::define_mode();
 			if ( ! Waf_Runner::is_allowed_mode( JETPACK_WAF_MODE ) ) {
-				return new WP_Error( 'waf_update_failed', 'Invalid firewall mode.' );
+				return new WP_Error( 'waf_mode_invalid', 'Invalid firewall mode.' );
 			}
 
 			try {
 				Waf_Rules_Manager::generate_ip_rules();
 				Waf_Rules_Manager::generate_rules();
 				( new Waf_Standalone_Bootstrap() )->generate();
-			} catch ( \Exception $e ) {
-				return new WP_Error( 'waf_update_failed', $e->getMessage() );
+			} catch ( Waf_Exception $e ) {
+				return $e->get_wp_error();
 			}
 		}