diff --git a/internal/common/helpers.go b/internal/common/helpers.go index 04e8a291a..339e5b73f 100644 --- a/internal/common/helpers.go +++ b/internal/common/helpers.go @@ -1,8 +1,12 @@ package common import ( + "fmt" + "regexp" "runtime" "sort" + "strconv" + "strings" ) var RuntimeGOARCH = runtime.GOARCH @@ -36,3 +40,51 @@ func IsStringInSortedSlice(slice []string, s string) bool { } return false } + +// DataSizeToUint64 converts a size specified as a string in KB/KiB/MB/etc. to +// a number of bytes represented by uint64. +func DataSizeToUint64(size string) (uint64, error) { + // Pre-process the input + size = strings.TrimSpace(size) + + // Get the number from the string + plain_number := regexp.MustCompile(`[[:digit:]]+`) + number_as_str := plain_number.FindString(size) + if number_as_str == "" { + return 0, fmt.Errorf("the size string doesn't contain any number: %s", size) + } + + // Parse the number into integer + return_size, err := strconv.ParseUint(number_as_str, 10, 64) + if err != nil { + return 0, fmt.Errorf("failed to parse size as integer: %s", number_as_str) + } + + // List of all supported units (from kB to TB and KiB to TiB) + supported_units := []struct { + re *regexp.Regexp + multiple uint64 + }{ + {regexp.MustCompile(`^\s*[[:digit:]]+\s*kB$`), 1000}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*KiB$`), 1024}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*MB$`), 1000 * 1000}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*MiB$`), 1024 * 1024}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*GB$`), 1000 * 1000 * 1000}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*GiB$`), 1024 * 1024 * 1024}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*TB$`), 1000 * 1000 * 1000 * 1000}, + {regexp.MustCompile(`^\s*[[:digit:]]+\s*TiB$`), 1024 * 1024 * 1024 * 1024}, + {regexp.MustCompile(`^\s*[[:digit:]]+$`), 1}, + } + + for _, unit := range supported_units { + if unit.re.MatchString(size) { + return_size *= unit.multiple + return return_size, nil + } + } + + // In case the strign didn't match any of the above regexes, return nil + // even if a number was found. This is to prevent users from submitting + // unknown units. + return 0, fmt.Errorf("unknown data size units in string: %s", size) +} diff --git a/internal/common/helpers_test.go b/internal/common/helpers_test.go index 467093221..e8ecaf67e 100644 --- a/internal/common/helpers_test.go +++ b/internal/common/helpers_test.go @@ -2,8 +2,10 @@ package common import ( "errors" - "github.com/stretchr/testify/assert" "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestCurrentArchAMD64(t *testing.T) { @@ -52,3 +54,40 @@ func TestIsStringInSortedSlice(t *testing.T) { assert.False(t, IsStringInSortedSlice([]string{"bart", "lisa", "marge"}, "")) assert.False(t, IsStringInSortedSlice([]string{}, "homer")) } + +func TestDataSizeToUint64(t *testing.T) { + cases := []struct { + input string + success bool + output uint64 + }{ + {"123", true, 123}, + {"123 kB", true, 123000}, + {"123 KiB", true, 123 * 1024}, + {"123 MB", true, 123 * 1000 * 1000}, + {"123 MiB", true, 123 * 1024 * 1024}, + {"123 GB", true, 123 * 1000 * 1000 * 1000}, + {"123 GiB", true, 123 * 1024 * 1024 * 1024}, + {"123 TB", true, 123 * 1000 * 1000 * 1000 * 1000}, + {"123 TiB", true, 123 * 1024 * 1024 * 1024 * 1024}, + {"123kB", true, 123000}, + {"123KiB", true, 123 * 1024}, + {" 123 ", true, 123}, + {" 123kB ", true, 123000}, + {" 123KiB ", true, 123 * 1024}, + {"123 KB", false, 0}, + {"123 mb", false, 0}, + {"123 PB", false, 0}, + {"123 PiB", false, 0}, + } + + for _, c := range cases { + result, err := DataSizeToUint64(c.input) + if c.success { + require.Nil(t, err) + assert.EqualValues(t, c.output, result) + } else { + assert.NotNil(t, err) + } + } +}