From 3e70cb3062f2854107569e5276d02eb4762891f2 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <stratonikov@runbox.com>
Date: Thu, 7 Oct 2021 19:35:43 +0300
Subject: [PATCH] Implement ToSizeInBytes(E)

Allow to convert string with size units (such as "10kb") to bytes.
This converter is privately used in viper. I think it is great to
allow using it from cast. Compared to viper implementation this one
allows to have single-letter suffixes.

Signed-off-by: Evgenii Stratonikov <stratonikov@runbox.com>
---
 cast.go      |  6 ++++++
 cast_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 caste.go     | 50 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+)

diff --git a/cast.go b/cast.go
index 0cfe941..38fe341 100644
--- a/cast.go
+++ b/cast.go
@@ -25,6 +25,12 @@ func ToTimeInDefaultLocation(i interface{}, location *time.Location) time.Time {
 	return v
 }
 
+// ToSizeInBytes casts an interface to the uint size in bytes.
+func ToSizeInBytes(i interface{}) uint {
+	v, _ := ToSizeInBytesE(i)
+	return v
+}
+
 // ToDuration casts an interface to a time.Duration type.
 func ToDuration(i interface{}) time.Duration {
 	v, _ := ToDurationE(i)
diff --git a/cast_test.go b/cast_test.go
index c254c57..2af46c0 100644
--- a/cast_test.go
+++ b/cast_test.go
@@ -9,7 +9,9 @@ import (
 	"errors"
 	"fmt"
 	"html/template"
+	"math"
 	"path"
+	"strconv"
 	"testing"
 	"time"
 
@@ -1298,6 +1300,50 @@ func TestToDurationE(t *testing.T) {
 	}
 }
 
+func TestToSizeInBytes(t *testing.T) {
+	const (
+		kb = 1024
+		mb = 1024 * kb
+		gb = 1024 * mb
+	)
+
+	tests := []struct {
+		input  interface{}
+		expect uint
+		iserr  bool
+	}{
+		{123, 123, false},
+		{ "  42  ", 42, false},
+		{"1", 1, false},
+		{"2b", 2, false},
+		{"3B", 3, false},
+		{"4 b", 4, false},
+		{"5k", 5 * kb, false},
+		{"6 KB", 6 * kb, false},
+		{"7m", 7 * mb, false},
+		{"8MB", 8 * mb, false},
+		{"9 g", 9 * gb, false},
+		{"10 GB", 10 * gb, false},
+		// errors
+		{"test", 0, true},
+		{strconv.FormatUint(uint64(math.MaxUint64 / kb), 10) + "9999k", 0, true}, // overflow with suffix
+		{strconv.FormatUint(uint64(math.MaxUint), 10) + "0", 0, true}, // overflow without suffix
+	}
+
+	for i, test := range tests {
+		errmsg := fmt.Sprintf("i = %d", i) // assert helper message
+
+		v, err := ToSizeInBytesE(test.input)
+		if test.iserr {
+			assert.Error(t, err, errmsg)
+			continue
+		}
+
+		assert.NoError(t, err, errmsg)
+		assert.Equal(t, test.expect, v)
+	}
+}
+
 func TestToTimeWithTimezones(t *testing.T) {
 
 	est, err := time.LoadLocation("EST")
diff --git a/caste.go b/caste.go
index c04af6a..d68b1c8 100644
--- a/caste.go
+++ b/caste.go
@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 	"time"
+	"unicode"
 )
 
 var errNegativeNotAllowed = errors.New("unable to cast negative value")
@@ -51,6 +52,55 @@ func ToTimeInDefaultLocationE(i interface{}, location *time.Location) (tim time.
 	}
 }
 
+func safeMul(a, b uint) (uint, error) {
+	c := a * b
+	if a > 1 && b > 1 && c/b != a {
+		return 0, errors.New("overflow")
+	}
+	return c, nil
+}
+
+
+// ToSizeInBytesE casts an empty interface to a size in bytes
+// interpreting the usual (k, m, g, kb, mB etc.) suffixes.
+func ToSizeInBytesE(i interface{}) (uint, error) {
+	sizeStr := strings.TrimSpace(ToString(i))
+	lastChar := len(sizeStr) - 1
+	multiplier := uint(1)
+
+	if lastChar > 0 {
+		if sizeStr[lastChar] == 'b' || sizeStr[lastChar] == 'B' {
+			sizeStr = sizeStr[:lastChar]
+			lastChar--
+		}
+		if lastChar > 0 {
+			switch unicode.ToLower(rune(sizeStr[lastChar])) {
+			case 'k':
+				multiplier = 1 << 10
+				sizeStr = strings.TrimSpace(sizeStr[:lastChar])
+			case 'm':
+				multiplier = 1 << 20
+				sizeStr = strings.TrimSpace(sizeStr[:lastChar])
+			case 'g':
+				multiplier = 1 << 30
+				sizeStr = strings.TrimSpace(sizeStr[:lastChar])
+			default:
+				multiplier = 1
+				sizeStr = strings.TrimSpace(sizeStr[:lastChar+1])
+			}
+		}
+	}
+
+	size, err := ToUintE(sizeStr)
+	if err != nil {
+		return 0, err
+	}
+
+	return safeMul(size, multiplier)
+}
+
+
+
 // ToDurationE casts an interface to a time.Duration type.
 func ToDurationE(i interface{}) (d time.Duration, err error) {
 	i = indirect(i)