2 files changed, 37 insertions, 28 deletions

diff --git a/firmware/drivers/ata.c b/firmware/drivers/ata.c
index 549a7bf920..db8bdf92a0 100644
--- a/firmware/drivers/ata.c
+++ b/firmware/drivers/ata.c
@@ -89,6 +89,7 @@ static bool initialized = false;
 static long last_user_activity = -1;
 long last_disk_activity = -1;
 
+static unsigned long total_sectors;
 static int multisectors; /* number of supported multisectors */
 static unsigned short identify_info[SECTOR_SIZE/2];
 
@@ -284,6 +285,11 @@ int ata_read_sectors(IF_MV2(int drive,)
     mutex_lock(&ata_mtx);
 #endif
 
+    if (start + incount > total_sectors) {
+        ret = -1;
+        goto error;
+    }
+
     last_disk_activity = current_tick;
     spinup_start = current_tick;
 
@@ -293,16 +299,14 @@ int ata_read_sectors(IF_MV2(int drive,)
         spinup = true;
         if (poweroff) {
             if (ata_power_on()) {
-                mutex_unlock(&ata_mtx);
-                ata_led(false);
-                return -1;
+                ret = -2;
+                goto error;
             }
         }
         else {
             if (perform_soft_reset()) {
-                mutex_unlock(&ata_mtx);
-                ata_led(false);
-                return -1;
+                ret = -2;
+                goto error;
             }
         }
     }
@@ -312,9 +316,8 @@ int ata_read_sectors(IF_MV2(int drive,)
     SET_REG(ATA_SELECT, ata_device);
     if (!wait_for_rdy())
     {
-        mutex_unlock(&ata_mtx);
-        ata_led(false);
-        return -2;
+        ret = -3;
+        goto error;
     }
 
  retry:
@@ -371,7 +374,7 @@ int ata_read_sectors(IF_MV2(int drive,)
                    We choose alternative 2.
                 */
                 perform_soft_reset();
-                ret = -4;
+                ret = -5;
                 goto retry;
             }
 
@@ -403,7 +406,7 @@ int ata_read_sectors(IF_MV2(int drive,)
             */
             if ( status & (STATUS_BSY | STATUS_ERR | STATUS_DF) ) {
                 perform_soft_reset();
-                ret = -5;
+                ret = -6;
                 goto retry;
             }
 
@@ -415,13 +418,14 @@ int ata_read_sectors(IF_MV2(int drive,)
 
         if(!ret && !wait_for_end_of_transfer()) {
             perform_soft_reset();
-            ret = -3;
+            ret = -4;
             goto retry;
         }
         break;
     }
-    ata_led(false);
 
+  error:
+    ata_led(false);
 #ifndef MAX_PHYS_SECTOR_SIZE
     mutex_unlock(&ata_mtx);
 #endif
@@ -489,6 +493,9 @@ int ata_write_sectors(IF_MV2(int drive,)
     mutex_lock(&ata_mtx);
 #endif
     
+    if (start + count > total_sectors)
+        panicf("Writing past end of disk");
+
     last_disk_activity = current_tick;
     spinup_start = current_tick;
 
@@ -498,16 +505,14 @@ int ata_write_sectors(IF_MV2(int drive,)
         spinup = true;
         if (poweroff) {
             if (ata_power_on()) {
-                mutex_unlock(&ata_mtx);
-                ata_led(false);
-                return -1;
+                ret = -1;
+                goto error;
             }
         }
         else {
             if (perform_soft_reset()) {
-                mutex_unlock(&ata_mtx);
-                ata_led(false);
-                return -1;
+                ret = -1;
+                goto error;
             }
         }
     }
@@ -515,9 +520,8 @@ int ata_write_sectors(IF_MV2(int drive,)
     SET_REG(ATA_SELECT, ata_device);
     if (!wait_for_rdy())
     {
-        mutex_unlock(&ata_mtx);
-        ata_led(false);
-        return -2;
+        ret = -2;
+        goto error;
     }
 
 #ifdef HAVE_LBA48
@@ -575,8 +579,8 @@ int ata_write_sectors(IF_MV2(int drive,)
         ret = -4;
     }
 
+  error:
     ata_led(false);
-
 #ifndef MAX_PHYS_SECTOR_SIZE
     mutex_unlock(&ata_mtx);
 #endif
@@ -1240,11 +1244,16 @@ int ata_init(void)
                    phys_sector_mult * SECTOR_SIZE);
 #endif
 
+        total_sectors = identify_info[60] | (identify_info[60] << 16);
+
 #ifdef HAVE_LBA48
-        if (identify_info[83] & 0x0400 /* 48 bit address support */
-            && identify_info[60] == 0xFFFF  /* and disk size >= 128 GiB */
-            && identify_info[61] == 0x0FFF) /* (needs BigLBA addressing) */
-        {
+        if (identify_info[83] & 0x0400       /* 48 bit address support */
+            && total_sectors == 0x0FFFFFFF)  /* and disk size >= 128 GiB */
+        {                                    /* (needs BigLBA addressing) */
+            if (identify_info[102] || identify_info[103])
+                panicf("Unsupported disk size, >= 2^32 sectors");
+                
+            total_sectors = identify_info[100] | (identify_info[101] << 16);
             lba48 = true; /* use BigLBA */
         }
 #endif
diff --git a/firmware/drivers/ata_mmc.c b/firmware/drivers/ata_mmc.c
index fdd091248f..e0a2b17500 100644
--- a/firmware/drivers/ata_mmc.c
+++ b/firmware/drivers/ata_mmc.c
@@ -814,7 +814,7 @@ int ata_write_sectors(IF_MV2(int drive,)
     }
 
     if (c_end_addr  > card->size)
-        panicf("Writing past end of card\n");
+        panicf("Writing past end of card");
 
     blocksize = card->blocksize;
     offset = c_addr & (blocksize - 1);