2022-12-08 13:05:19 -03:00

122 lines
2.4 KiB
Go

// Copyright (c) 2020 Gitpod GmbH. All rights reserved.
// Licensed under the GNU Affero General Public License (AGPL).
// See License.AGPL.txt in the project root for license information.
package quota
import (
"encoding/json"
"errors"
"fmt"
"regexp"
"strconv"
"golang.org/x/xerrors"
)
// Size denotes a filesize / amount of bytes
type Size int64
const (
// Byte is a single byte
Byte Size = 1
// Kilobyte is 1024 bytes
Kilobyte Size = 1024 * Byte
// Megabyte is 1024 kilobytes
Megabyte Size = 1024 * Kilobyte
// Gigabyte is 1024 megabytes
Gigabyte Size = 1024 * Megabyte
// Terabyte is 1024 gigabyte
Terabyte Size = 1024 * Gigabyte
)
var (
sizeRegexp = regexp.MustCompile(`^(\d+)(k|m|g|t)?$`)
// ErrInvalidSize is returned by ParseSize if input was not a valid size
ErrInvalidSize = errors.New("invalid size")
)
// ParseSize parses a number of bytes (e.g. 10) into a size.
// ParseSize supports units: k(ilobyte), m(egabyte), (g)igatebyte, (t)erabyte. E.g. 50m or 200g
func ParseSize(q string) (Size, error) {
matches := sizeRegexp.FindSubmatch([]byte(q))
if len(matches) == 0 {
return 0, ErrInvalidSize
}
val, err := strconv.ParseUint(string(matches[1]), 10, 64)
if err != nil {
return 0, ErrInvalidSize
}
size := Size(val)
unit := string(matches[2])
switch unit {
case "":
size = size * Byte
case "k":
size = size * Kilobyte
case "m":
size = size * Megabyte
case "g":
size = size * Gigabyte
case "t":
size = size * Terabyte
}
return size, nil
}
// String converts the size to a string
func (s Size) String() string {
if s == 0 {
return "0"
}
steps := []struct {
Prefix string
Q Size
}{
{"t", Terabyte},
{"g", Gigabyte},
{"m", Megabyte},
{"k", Kilobyte},
}
for _, stp := range steps {
if s%stp.Q == 0 {
return fmt.Sprintf("%d%s", s/stp.Q, stp.Prefix)
}
}
return fmt.Sprintf("%d", s)
}
// MarshalJSON marshals a duration to JSON
func (s Size) MarshalJSON() ([]byte, error) {
return json.Marshal(s.String())
}
// UnmarshalJSON unmarshals a duration from JSON
func (s *Size) UnmarshalJSON(b []byte) error {
var v interface{}
if err := json.Unmarshal(b, &v); err != nil {
return err
}
switch value := v.(type) {
case string:
if value == "" {
*s = 0
return nil
}
tmp, err := ParseSize(value)
if err != nil {
return err
}
*s = tmp
return nil
default:
return xerrors.Errorf("invalid size: %s", string(b))
}
}