diff --git a/src/croc/croc.go b/src/croc/croc.go index 141afb4..a10b70b 100644 --- a/src/croc/croc.go +++ b/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 { diff --git a/src/diskusage/diskusage.go b/src/diskusage/diskusage.go new file mode 100644 index 0000000..8e8bc1b --- /dev/null +++ b/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()) +} diff --git a/src/diskusage/diskusage_test.go b/src/diskusage/diskusage_test.go new file mode 100644 index 0000000..3fe88a9 --- /dev/null +++ b/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, "%") +} diff --git a/src/diskusage/diskusage_windows.go b/src/diskusage/diskusage_windows.go new file mode 100644 index 0000000..a4b0865 --- /dev/null +++ b/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()) +}