1package atc
2
3import (
4	"encoding/json"
5	"errors"
6	"math"
7	"regexp"
8	"strconv"
9	"strings"
10)
11
12const MemoryRegex = "^([0-9]+)([G|M|K|g|m|k]?[b|B])?$"
13
14func (c *ContainerLimits) UnmarshalJSON(limit []byte) error {
15	var data interface{}
16
17	err := json.Unmarshal(limit, &data)
18	if err != nil {
19		return err
20	}
21
22	climits, err := ParseContainerLimits(data)
23	if err != nil {
24		return err
25	}
26
27	c.CPU = climits.CPU
28	c.Memory = climits.Memory
29
30	return nil
31}
32
33func ParseContainerLimits(data interface{}) (ContainerLimits, error) {
34	mapData, ok := data.(map[string]interface{})
35	if !ok {
36		mapData = make(map[string]interface{})
37	}
38
39	var c ContainerLimits
40
41	var memoryBytes uint64
42	var uVal int
43	var err error
44
45	// the json unmarshaller returns numbers as float64 while yaml returns int
46	for key, val := range mapData {
47		if key == "memory" {
48			switch val.(type) {
49			case string:
50				memoryBytes, err = parseMemoryLimit(val.(string))
51				if err != nil {
52					return ContainerLimits{}, err
53				}
54			case *string:
55				if val.(*string) == nil {
56					c.Memory = nil
57					continue
58				}
59				memoryBytes, err = parseMemoryLimit(*val.(*string))
60				if err != nil {
61					return ContainerLimits{}, err
62				}
63			case float64:
64				memoryBytes = uint64(int(val.(float64)))
65			case int:
66				memoryBytes = uint64(val.(int))
67			}
68			c.Memory = &memoryBytes
69
70		} else if key == "cpu" {
71			switch val.(type) {
72			case float64:
73				uVal = int(val.(float64))
74			case int:
75				uVal = val.(int)
76			case *int:
77				if val.(*int) == nil {
78					c.CPU = nil
79					continue
80				}
81				uVal = *val.(*int)
82			default:
83				return ContainerLimits{}, errors.New("cpu limit must be an integer")
84			}
85			helper := uint64(uVal)
86			c.CPU = &helper
87
88		}
89	}
90
91	return c, nil
92}
93
94func parseMemoryLimit(limit string) (uint64, error) {
95	limit = strings.ToUpper(limit)
96	var sizeRegex *regexp.Regexp = regexp.MustCompile(MemoryRegex)
97	matches := sizeRegex.FindStringSubmatch(limit)
98
99	if len(matches) > 3 || len(matches) < 1 {
100		return 0, errors.New("could not parse container memory limit")
101	}
102
103	value, err := strconv.ParseUint(matches[1], 10, 64)
104	if err != nil {
105		return 0, err
106	}
107
108	var power float64
109	var base float64 = 2
110	var unit string = matches[2]
111	switch unit {
112	case "KB":
113		power = 10
114	case "MB":
115		power = 20
116	case "GB":
117		power = 30
118	default:
119		power = 0
120	}
121
122	return value * uint64(math.Pow(base, power)), nil
123}
124