Check disk space before downloading file

Signed-off-by: Jianhui Zhao <zhaojh329@gmail.com>

src/downfile.c

@@ -64,6 +64,10 @@ static void on_socket_read(struct ev_loop *loop, struct ev_io *w, int revents)
         puts("");
         ev_break(loop, EVBREAK_ALL);
         break;
+    case RTTY_FILE_MSG_NO_SPACE:
+        fprintf(stderr, "No enough space\n");
+        ev_break(loop, EVBREAK_ALL);
+        break;
     default:
         break;
     }

src/file.c

@@ -29,6 +29,8 @@
 #include <unistd.h>
 #include <sys/socket.h>
 #include <netinet/in.h>
+#include <sys/statvfs.h>
+#include <mntent.h>
 
 #include "log.h"
 #include "file.h"
@@ -118,15 +120,39 @@ static void start_upload_file(struct file_context *ctx, struct buffer *info)
     log_info("upload file: %s\n", path);
 }
 
+static void send_canceled_msg(struct rtty *rtty)
+{
+    buffer_put_u8(&rtty->wb, MSG_TYPE_FILE);
+    buffer_put_u16be(&rtty->wb, 2);
+    buffer_put_u8(&rtty->wb, rtty->file_context.sid);
+    buffer_put_u8(&rtty->wb, RTTY_FILE_MSG_CANCELED);
+    ev_io_start(rtty->loop, &rtty->iow);
+}
+
 static void start_download_file(struct file_context *ctx, struct buffer *info, int len)
 {
+    struct rtty *rtty = container_of(ctx, struct rtty, file_context);
+    struct mntent *ment;
+    struct statvfs sfs;
     char *name;
     int fd;
 
     ctx->total_size = ctx->remain_size = buffer_pull_u32be(info);
     name = strndup(buffer_data(info), len - 4);
     buffer_pull(info, NULL, len - 4);
 
+    ment = find_mount_point(abspath);
+    if (ment) {
+        if (statvfs(ment->mnt_dir, &sfs) == 0 && ctx->total_size > sfs.f_bavail * sfs.f_frsize) {
+            int type = RTTY_FILE_MSG_NO_SPACE;
+
+            send_canceled_msg(rtty);
+            sendto(ctx->sock, &type, 1, 0, (struct sockaddr *) &ctx->peer_sun, sizeof(struct sockaddr_un));
+            log_err("download file fail: no enough space\n");
+            return;
+        }
+    }
+
     strcat(abspath, name);
 
     fd = open(abspath, O_WRONLY | O_TRUNC | O_CREAT, 0644);
@@ -175,12 +201,7 @@ static void on_socket_read(struct ev_loop *loop, struct ev_io *w, int revents)
         }
 
         ctx->busy = false;
-
-        buffer_put_u8(&rtty->wb, MSG_TYPE_FILE);
-        buffer_put_u16be(&rtty->wb, 2);
-        buffer_put_u8(&rtty->wb, ctx->sid);
-        buffer_put_u8(&rtty->wb, RTTY_FILE_MSG_CANCELED);
-        ev_io_start(rtty->loop, &rtty->iow);
+        send_canceled_msg(rtty);
         break;
     case RTTY_FILE_MSG_SAVE_PATH:
         strcpy(abspath, buffer_data(&b));

src/file.h

@@ -38,7 +38,8 @@ enum {
     RTTY_FILE_MSG_BUSY,
     RTTY_FILE_MSG_PROGRESS,
     RTTY_FILE_MSG_REQUEST_ACCEPT,
-    RTTY_FILE_MSG_SAVE_PATH
+    RTTY_FILE_MSG_SAVE_PATH,
+    RTTY_FILE_MSG_NO_SPACE
 };
 
 struct file_context {

src/utils.c

@@ -22,14 +22,14 @@
  * SOFTWARE.
  */
 
-#include "utils.h"
+#include <sys/stat.h>
+#include <mntent.h>
 #include <stdio.h>
-#include <unistd.h>
 #include <stdint.h>
 #include <string.h>
-#include <errno.h>
 #include <ctype.h>
-#include <sys/time.h>
+
+#include "utils.h"
 
 int find_login(char *buf, int len)
 {
@@ -115,3 +115,42 @@ const char *format_size(size_t size)
     return str;
 }
 
+/*
+ * Given any file (or directory), find the mount table entry for its
+ * filesystem.
+ */
+struct mntent *find_mount_point(const char *name)
+{
+    struct mntent *ment;
+    dev_t devno_of_name;
+    struct stat s;
+    FILE *mtab_fp;
+
+    if (stat(name, &s) < 0)
+        return NULL;
+
+    devno_of_name = s.st_dev;
+
+    if (S_ISBLK(s.st_mode) || S_ISCHR(s.st_mode))
+        return NULL;
+
+    mtab_fp = setmntent("/etc/mtab", "r");
+    if (!mtab_fp)
+        return NULL;
+
+    while ((ment = getmntent(mtab_fp))) {
+        if (!strcmp(ment->mnt_fsname, "rootfs"))
+            continue;
+
+        /* string match */
+        if (!strcmp(name, ment->mnt_dir))
+            break;
+
+        /* match the directory's mount point. */
+        if (stat(ment->mnt_dir, &s) == 0 && s.st_dev == devno_of_name)
+            break;
+    }
+    endmntent(mtab_fp);
+
+    return ment;
+}

src/utils.h

@@ -36,4 +36,6 @@ int b64_encode(const void *src, size_t srclen, void *dest, size_t destsize);
 
 const char *format_size(size_t size);
 
+struct mntent *find_mount_point(const char *name);
+
 #endif