diff --git a/block_class.module b/block_class.module
index 130158f..7e4a1fe 100644
--- a/block_class.module
+++ b/block_class.module
@@ -1,57 +1,45 @@
 <?php
 
 /**
- * @file Module for adding classes to blocks.
+ * @file
+ * Module for adding classes to blocks.
  */
 
+use Drupal\block\BlockInterface;
+use Drupal\Core\Block\BlockPluginInterface;
+
 /**
  * Implements hook_form_FORM_ID_alter().
  */
 function block_class_form_block_form_alter(&$form, \Drupal\Core\Form\FormStateInterface $form_state, $form_id) {
   if (\Drupal::currentUser()->hasPermission('administer block classes')) {
-    $config = \Drupal::configFactory()->getEditable('block_class.settings');
-    $block_id = $form['id']['#default_value'];
 
-    $form['block_classes'] = array(
+    /** @var \Drupal\block\BlockInterface $block */
+    $block = $form_state->getFormObject()->getEntity();
+
+    // This will automatically be saved in the third party settings.
+    $form['third_party_settings']['#tree'] = TRUE;
+    $form['third_party_settings']['block_class']['classes'] = array(
       '#type' => 'textfield',
       '#title' => t('CSS class(es)'),
       '#description' => t('Customize the styling of this block by adding CSS classes. Separate multiple classes by spaces.'),
-      '#default_value' => $config->get($block_id.'.block_classes'),
+      '#default_value' => $block->getThirdPartySetting('block_class', 'classes'),
     );
 
-    $form['actions']['submit']['#submit'][] = 'block_class_form_submit';
-
-  }
-}
-
-/**
- * Form submission callback.
- */
-function block_class_form_submit($form, \Drupal\Core\Form\FormStateInterface &$form_state) {
-  $config = \Drupal::configFactory()->getEditable('block_class.settings');
-  $block_id = $form['id']['#default_value'];
-
-  $classes = trim($form_state->getValue('block_classes'));
-
-  if(empty($classes)) {
-    $config->clear($block_id.'.block_classes',$classes);
-    $config->save();
-  }
-  else {
-    $config->set($block_id.'.block_classes',$classes);
-    $config->save();
   }
 }
 
 /**
  * Implements hook_preprocess_HOOK().
  */
-function block_class_preprocess_block(&$variables) {
-  $config = \Drupal::config('block_class.settings');
-  // Blocks coming from page manager widget does not have id.
-  if (!empty($variables['elements']['#id'])) {
-    $block_id = $variables['elements']['#id'];
-    $classes = $config->get($block_id . '.block_classes');
-    $variables['attributes']['class'][] = $classes;
+function block_class_block_view_alter(&$build, BlockPluginInterface $plugin) {
+
+  // Get the classes from the third party setting.
+  if (isset($build['#block'])) {
+    /** @var \Drupal\block\BlockInterface $block */
+    $block = $build['#block'];
+    if ($classes = $block->getThirdPartySetting('block_class', 'classes')) {
+      $build['#attributes']['class'][] = $classes;
+    }
   }
 }
diff --git a/config/schema/block_class.schema.yml b/config/schema/block_class.schema.yml
new file mode 100644
index 0000000..cb6e83c
--- /dev/null
+++ b/config/schema/block_class.schema.yml
@@ -0,0 +1,7 @@
+block.block.*.third_party.block_class:
+  type: mapping
+  label: Block class third party settings
+  mapping:
+    classes:
+      type: string
+      label: Additional classes for the block
diff --git a/src/Tests/BlockClassTest.php b/src/Tests/BlockClassTest.php
index 678a9af..7e5c792 100644
--- a/src/Tests/BlockClassTest.php
+++ b/src/Tests/BlockClassTest.php
@@ -16,11 +16,6 @@ use Drupal\simpletest\WebTestBase;
 class BlockClassTest extends WebTestBase {
 
   /**
-   * Disabled config schema checking temporarily until all errors are resolved.
-   */
-  protected $strictConfigSchema = FALSE;
-
-  /**
    * Modules to enable.
    *
    * @var array
@@ -43,7 +38,7 @@ class BlockClassTest extends WebTestBase {
     $this->clickLink(t('Main page content'));
     $edit = array(
       'region' => 'content',
-      'block_classes' => 'TestClass_content'
+      'third_party_settings[block_class][classes]' => 'TestClass_content'
     );
     $this->drupalPostForm(NULL, $edit, t('Save block'));
 
@@ -51,7 +46,7 @@ class BlockClassTest extends WebTestBase {
     $this->clickLink(t('User account menu'));
     $edit = array(
       'region' => 'secondary_menu',
-      'block_classes' => 'TestClass_menu'
+      'third_party_settings[block_class][classes]' => 'TestClass_menu'
     );
     $this->drupalPostForm(NULL, $edit, t('Save block'));
 
@@ -60,7 +55,7 @@ class BlockClassTest extends WebTestBase {
     // Assert the custom class in the content block.
     $this->assertRaw('<div id="block-mainpagecontent" class="TestClass_content block block-system">');
     // Assert the custom class in user menu.
-    $this->assertRaw('<nav id="block-useraccountmenu" class="TestClass_menu block block-menu navigation menu--account" role="navigation" aria-labelledby="block-useraccountmenu-menu">');
+    $this->assertRaw('class="TestClass_menu block block-menu navigation menu--account"');
   }
 
 }