2 files changed, 27 insertions, 5 deletions

diff --git a/arch/x86/kernel/apic.c b/arch/x86/kernel/apic.c
index 784116933c70..1f48bd1c9f2d 100644
--- a/arch/x86/kernel/apic.c
+++ b/arch/x86/kernel/apic.c
@@ -1360,7 +1360,12 @@ void enable_IR_x2apic(void)
 
         local_irq_save(flags);
         mask_8259A();
-        save_mask_IO_APIC_setup();
+
+        ret = save_mask_IO_APIC_setup();
+        if (ret) {
+                printk(KERN_INFO "Saving IO-APIC state failed: %d\n", ret);
+                goto end;
+        }
 
         ret = enable_intr_remapping(1);
 
@@ -1370,14 +1375,15 @@ void enable_IR_x2apic(void)
         }
 
         if (ret)
-                goto end;
+                goto end_restore;
 
         if (!x2apic) {
                 x2apic = 1;
                 apic_ops = &x2apic_ops;
                 enable_x2apic();
         }
-end:
+
+end_restore:
         if (ret)
                 /*
                  * IR enabling failed
@@ -1386,6 +1392,7 @@ end:
         else
                 reinit_intr_remapped_IO_APIC(x2apic_preenabled);
 
+end:
         unmask_8259A();
         local_irq_restore(flags);
 
diff --git a/arch/x86/kernel/io_apic.c b/arch/x86/kernel/io_apic.c
index 01419fdc1d5d..4cc9cb64c811 100644
--- a/arch/x86/kernel/io_apic.c
+++ b/arch/x86/kernel/io_apic.c
@@ -812,7 +812,7 @@ int save_mask_IO_APIC_setup(void)
                         kzalloc(sizeof(struct IO_APIC_route_entry) *
                                 nr_ioapic_registers[apic], GFP_KERNEL);
                 if (!early_ioapic_entries[apic])
-                        return -ENOMEM;
+                        goto nomem;
         }
 
         for (apic = 0; apic < nr_ioapics; apic++)
@@ -826,17 +826,32 @@ int save_mask_IO_APIC_setup(void)
                                 ioapic_write_entry(apic, pin, entry);
                         }
                 }
+
         return 0;
+
+nomem:
+        for (; apic > 0; apic--)
+                kfree(early_ioapic_entries[apic]);
+        kfree(early_ioapic_entries[apic]);
+        memset(early_ioapic_entries, 0,
+                ARRAY_SIZE(early_ioapic_entries));
+
+        return -ENOMEM;
 }
 
 void restore_IO_APIC_setup(void)
 {
         int apic, pin;
 
-        for (apic = 0; apic < nr_ioapics; apic++)
+        for (apic = 0; apic < nr_ioapics; apic++) {
+                if (!early_ioapic_entries[apic])
+                        break;
                 for (pin = 0; pin < nr_ioapic_registers[apic]; pin++)
                         ioapic_write_entry(apic, pin,
                                            early_ioapic_entries[apic][pin]);
+                kfree(early_ioapic_entries[apic]);
+                early_ioapic_entries[apic] = NULL;
+        }
 }
 
 void reinit_intr_remapped_IO_APIC(int intr_remapping)