From a0abf10c8e6dd2c05815434b7ef78dbf80c013e0 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 7 Oct 2021 19:15:48 +0300
Subject: [PATCH] [#857] config: allow to parse sizes with suffix

The code is taken from viper as `cast` package that we use
doesn't have needed converter.

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 cmd/neofs-node/config/cast.go          | 62 ++++++++++++++++++++++++++
 cmd/neofs-node/config/cast_test.go     | 19 ++++++++
 cmd/neofs-node/config/test/config.json | 10 +++++
 cmd/neofs-node/config/test/config.yaml |  9 ++++
 4 files changed, 100 insertions(+)

diff --git a/cmd/neofs-node/config/cast.go b/cmd/neofs-node/config/cast.go
index b8e4b4938..c4b7e602b 100644
--- a/cmd/neofs-node/config/cast.go
+++ b/cmd/neofs-node/config/cast.go
@@ -1,7 +1,9 @@
 package config
 
 import (
+	"strings"
 	"time"
+	"unicode"
 
 	"github.com/spf13/cast"
 )
@@ -144,3 +146,63 @@ func Int(c *Config, name string) int64 {
 func IntSafe(c *Config, name string) int64 {
 	return cast.ToInt64(c.Value(name))
 }
+
+// SizeInBytesSafe reads configuration value
+// from c by name and casts it to size in bytes (uint64).
+//
+// The suffix can be single-letter (b, k, m, g, t) or with
+// an additional b at the end. Spaces between the number and the suffix
+// are allowed. All multipliers are power of 2 (i.e. k is for kibi-byte).
+//
+// Returns 0 if a value can't be casted.
+func SizeInBytesSafe(c *Config, name string) uint64 {
+	s := StringSafe(c, name)
+	return parseSizeInBytes(s)
+}
+
+// The following code is taken from https://github.com/spf13/viper/blob/master/util.go
+// with minor corrections (allow to use both `k` and `kb` forms.
+// Seems like viper allows to convert sizes but corresponding parser in `cast` package
+// is missing.
+func safeMul(a, b uint64) uint64 {
+	c := a * b
+	if a > 1 && b > 1 && c/b != a {
+		return 0
+	}
+	return c
+}
+
+// parseSizeInBytes converts strings like 1GB or 12 mb into an unsigned integer number of bytes
+func parseSizeInBytes(sizeStr string) uint64 {
+	sizeStr = strings.TrimSpace(sizeStr)
+	lastChar := len(sizeStr) - 1
+	multiplier := uint64(1)
+
+	if lastChar > 0 {
+		if sizeStr[lastChar] == 'b' || sizeStr[lastChar] == 'B' {
+			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])
+			case 't':
+				multiplier = 1 << 40
+				sizeStr = strings.TrimSpace(sizeStr[:lastChar])
+			default:
+				multiplier = 1
+				sizeStr = strings.TrimSpace(sizeStr[:lastChar+1])
+			}
+		}
+	}
+
+	size := cast.ToUint64(sizeStr)
+	return safeMul(size, multiplier)
+}
diff --git a/cmd/neofs-node/config/cast_test.go b/cmd/neofs-node/config/cast_test.go
index b6ea84d2f..8ff493a63 100644
--- a/cmd/neofs-node/config/cast_test.go
+++ b/cmd/neofs-node/config/cast_test.go
@@ -115,3 +115,22 @@ func TestNumbers(t *testing.T) {
 		require.Zero(t, config.UintSafe(c, incorrect))
 	})
 }
+
+func TestSizeInBytes(t *testing.T) {
+	const (
+		kb = 1024
+		mb = 1024 * kb
+		gb = 1024 * mb
+		tb = 1024 * gb
+	)
+	configtest.ForEachFileType("test/config", func(c *config.Config) {
+		c = c.Sub("sizes")
+		require.EqualValues(t, kb, config.SizeInBytesSafe(c, "size_kb"))
+		require.EqualValues(t, 2*kb, config.SizeInBytesSafe(c, "size_kb_no_space"))
+		require.EqualValues(t, 12*mb, config.SizeInBytesSafe(c, "size_mb"))
+		require.EqualValues(t, 4*gb, config.SizeInBytesSafe(c, "size_gb"))
+		require.EqualValues(t, 5*tb, config.SizeInBytesSafe(c, "size_tb"))
+		require.EqualValues(t, 2048, config.SizeInBytesSafe(c, "size_bytes"))
+		require.EqualValues(t, 123456, config.SizeInBytesSafe(c, "size_bytes_no_suffix"))
+	})
+}
diff --git a/cmd/neofs-node/config/test/config.json b/cmd/neofs-node/config/test/config.json
index 81bae1cdb..4c1c0caf9 100644
--- a/cmd/neofs-node/config/test/config.json
+++ b/cmd/neofs-node/config/test/config.json
@@ -46,5 +46,15 @@
     "correct": true,
     "correct_string": "true",
     "incorrect": "not true"
+  },
+
+  "sizes": {
+    "size_kb": "1 kb",
+    "size_kb_no_space": "2kb",
+    "size_mb": "12m",
+    "size_gb": "4g",
+    "size_tb": "5 TB",
+    "size_bytes": "2048b",
+    "size_bytes_no_suffix": 123456
   }
 }
diff --git a/cmd/neofs-node/config/test/config.yaml b/cmd/neofs-node/config/test/config.yaml
index 2b080977b..062bed4d0 100644
--- a/cmd/neofs-node/config/test/config.yaml
+++ b/cmd/neofs-node/config/test/config.yaml
@@ -40,3 +40,12 @@ bool:
   correct: true
   correct_string: "true"
   incorrect: not true
+
+sizes:
+  size_kb: 1 kb
+  size_kb_no_space: 2kb
+  size_mb: 12m
+  size_gb: 4g
+  size_tb: 5 TB
+  size_bytes: 2048b
+  size_bytes_no_suffix: 123456