Merge pull request #792 from schollz/schollz/issue791

feature: gracefully exit when not enough disk space

src/croc/croc.go

@@ -31,6 +31,7 @@ import (
 	"github.com/schollz/croc/v10/src/comm"
 	"github.com/schollz/croc/v10/src/compress"
 	"github.com/schollz/croc/v10/src/crypt"
+	"github.com/schollz/croc/v10/src/diskusage"
 	"github.com/schollz/croc/v10/src/message"
 	"github.com/schollz/croc/v10/src/models"
 	"github.com/schollz/croc/v10/src/tcp"
@@ -1268,6 +1269,12 @@ func (c *Client) processMessageFileInfo(m message.Message) (done bool, err error
 			}
 		}
 	}
+	// check the totalSize does not exceed disk space
+	usage := diskusage.NewDiskUsage(".")
+	if usage.Available() < uint64(totalSize) {
+		return true, fmt.Errorf("not enough disk space")
+	}
+
 	// c.spinner.Stop()
 	action := "Accept"
 	if c.Options.SendingText {

src/diskusage/diskusage.go

@@ -0,0 +1,49 @@
+//go:build !windows
+// +build !windows
+
+package diskusage
+
+import (
+	"golang.org/x/sys/unix"
+)
+
+// DiskUsage contains usage data and provides user-friendly access methods
+type DiskUsage struct {
+	stat *unix.Statfs_t
+}
+
+// NewDiskUsages returns an object holding the disk usage of volumePath
+// or nil in case of error (invalid path, etc)
+func NewDiskUsage(volumePath string) *DiskUsage {
+	stat := unix.Statfs_t{}
+	err := unix.Statfs(volumePath, &stat)
+	if err != nil {
+		return nil
+	}
+	return &DiskUsage{&stat}
+}
+
+// Free returns total free bytes on file system
+func (du *DiskUsage) Free() uint64 {
+	return du.stat.Bfree * uint64(du.stat.Bsize)
+}
+
+// Available return total available bytes on file system to an unprivileged user
+func (du *DiskUsage) Available() uint64 {
+	return uint64(du.stat.Bavail) * uint64(du.stat.Bsize)
+}
+
+// Size returns total size of the file system
+func (du *DiskUsage) Size() uint64 {
+	return uint64(du.stat.Blocks) * uint64(du.stat.Bsize)
+}
+
+// Used returns total bytes used in file system
+func (du *DiskUsage) Used() uint64 {
+	return du.Size() - du.Free()
+}
+
+// Usage returns percentage of use on the file system
+func (du *DiskUsage) Usage() float32 {
+	return float32(du.Used()) / float32(du.Size())
+}

src/diskusage/diskusage_test.go

@@ -0,0 +1,17 @@
+package diskusage
+
+import (
+	"fmt"
+	"testing"
+)
+
+var KB = uint64(1024)
+
+func TestNewDiskUsage(t *testing.T) {
+	usage := NewDiskUsage(".")
+	fmt.Println("Free:", usage.Free()/(KB*KB))
+	fmt.Println("Available:", usage.Available()/(KB*KB))
+	fmt.Println("Size:", usage.Size()/(KB*KB))
+	fmt.Println("Used:", usage.Used()/(KB*KB))
+	fmt.Println("Usage:", usage.Usage()*100, "%")
+}

src/diskusage/diskusage_windows.go

@@ -0,0 +1,55 @@
+package diskusage
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+type DiskUsage struct {
+	freeBytes  int64
+	totalBytes int64
+	availBytes int64
+}
+
+// NewDiskUsages returns an object holding the disk usage of volumePath
+// or nil in case of error (invalid path, etc)
+func NewDiskUsage(volumePath string) *DiskUsage {
+	h := windows.MustLoadDLL("kernel32.dll")
+	c := h.MustFindProc("GetDiskFreeSpaceExW")
+
+	du := &DiskUsage{}
+
+	c.Call(
+		uintptr(unsafe.Pointer(windows.StringToUTF16Ptr(volumePath))),
+		uintptr(unsafe.Pointer(&du.freeBytes)),
+		uintptr(unsafe.Pointer(&du.totalBytes)),
+		uintptr(unsafe.Pointer(&du.availBytes)))
+
+	return du
+}
+
+// Free returns total free bytes on file system
+func (du *DiskUsage) Free() uint64 {
+	return uint64(du.freeBytes)
+}
+
+// Available returns total available bytes on file system to an unprivileged user
+func (du *DiskUsage) Available() uint64 {
+	return uint64(du.availBytes)
+}
+
+// Size returns total size of the file system
+func (du *DiskUsage) Size() uint64 {
+	return uint64(du.totalBytes)
+}
+
+// Used returns total bytes used in file system
+func (du *DiskUsage) Used() uint64 {
+	return du.Size() - du.Free()
+}
+
+// Usage returns percentage of use on the file system
+func (du *DiskUsage) Usage() float32 {
+	return float32(du.Used()) / float32(du.Size())
+}