1package command
2
3import (
4	"bytes"
5	"encoding/json"
6	"fmt"
7	"io/ioutil"
8	"os"
9	"strings"
10
11	multierror "github.com/hashicorp/go-multierror"
12	"github.com/hashicorp/hcl"
13	"github.com/hashicorp/hcl/hcl/ast"
14	"github.com/hashicorp/nomad/api"
15	"github.com/hashicorp/nomad/helper"
16	"github.com/hashicorp/nomad/jobspec"
17	"github.com/mitchellh/mapstructure"
18	"github.com/posener/complete"
19)
20
21type QuotaApplyCommand struct {
22	Meta
23}
24
25func (c *QuotaApplyCommand) Help() string {
26	helpText := `
27Usage: nomad quota apply [options] <input>
28
29  Apply is used to create or update a quota specification. The specification file
30  will be read from stdin by specifying "-", otherwise a path to the file is
31  expected.
32
33General Options:
34
35  ` + generalOptionsUsage() + `
36
37Apply Options:
38
39  -json
40    Parse the input as a JSON quota specification.
41`
42
43	return strings.TrimSpace(helpText)
44}
45
46func (c *QuotaApplyCommand) AutocompleteFlags() complete.Flags {
47	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
48		complete.Flags{
49			"-json": complete.PredictNothing,
50		})
51}
52
53func (c *QuotaApplyCommand) AutocompleteArgs() complete.Predictor {
54	return complete.PredictFiles("*")
55}
56
57func (c *QuotaApplyCommand) Synopsis() string {
58	return "Create or update a quota specification"
59}
60
61func (c *QuotaApplyCommand) Name() string { return "quota apply" }
62
63func (c *QuotaApplyCommand) Run(args []string) int {
64	var jsonInput bool
65	flags := c.Meta.FlagSet(c.Name(), FlagSetClient)
66	flags.Usage = func() { c.Ui.Output(c.Help()) }
67	flags.BoolVar(&jsonInput, "json", false, "")
68
69	if err := flags.Parse(args); err != nil {
70		return 1
71	}
72
73	// Check that we get exactly one argument
74	args = flags.Args()
75	if l := len(args); l != 1 {
76		c.Ui.Error("This command takes one argument: <input>")
77		c.Ui.Error(commandErrorText(c))
78		return 1
79	}
80
81	// Read the file contents
82	file := args[0]
83	var rawQuota []byte
84	var err error
85	if file == "-" {
86		rawQuota, err = ioutil.ReadAll(os.Stdin)
87		if err != nil {
88			c.Ui.Error(fmt.Sprintf("Failed to read stdin: %v", err))
89			return 1
90		}
91	} else {
92		rawQuota, err = ioutil.ReadFile(file)
93		if err != nil {
94			c.Ui.Error(fmt.Sprintf("Failed to read file: %v", err))
95			return 1
96		}
97	}
98
99	var spec *api.QuotaSpec
100	if jsonInput {
101		var jsonSpec api.QuotaSpec
102		dec := json.NewDecoder(bytes.NewBuffer(rawQuota))
103		if err := dec.Decode(&jsonSpec); err != nil {
104			c.Ui.Error(fmt.Sprintf("Failed to parse quota: %v", err))
105			return 1
106		}
107		spec = &jsonSpec
108	} else {
109		hclSpec, err := parseQuotaSpec(rawQuota)
110		if err != nil {
111			c.Ui.Error(fmt.Sprintf("Error parsing quota specification: %s", err))
112			return 1
113		}
114
115		spec = hclSpec
116	}
117
118	// Get the HTTP client
119	client, err := c.Meta.Client()
120	if err != nil {
121		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
122		return 1
123	}
124
125	_, err = client.Quotas().Register(spec, nil)
126	if err != nil {
127		c.Ui.Error(fmt.Sprintf("Error applying quota specification: %s", err))
128		return 1
129	}
130
131	c.Ui.Output(fmt.Sprintf("Successfully applied quota specification %q!", spec.Name))
132	return 0
133}
134
135// parseQuotaSpec is used to parse the quota specification from HCL
136func parseQuotaSpec(input []byte) (*api.QuotaSpec, error) {
137	root, err := hcl.ParseBytes(input)
138	if err != nil {
139		return nil, err
140	}
141
142	// Top-level item should be a list
143	list, ok := root.Node.(*ast.ObjectList)
144	if !ok {
145		return nil, fmt.Errorf("error parsing: root should be an object")
146	}
147
148	var spec api.QuotaSpec
149	if err := parseQuotaSpecImpl(&spec, list); err != nil {
150		return nil, err
151	}
152
153	return &spec, nil
154}
155
156// parseQuotaSpecImpl parses the quota spec taking as input the AST tree
157func parseQuotaSpecImpl(result *api.QuotaSpec, list *ast.ObjectList) error {
158	// Check for invalid keys
159	valid := []string{
160		"name",
161		"description",
162		"limit",
163	}
164	if err := helper.CheckHCLKeys(list, valid); err != nil {
165		return err
166	}
167
168	// Decode the full thing into a map[string]interface for ease
169	var m map[string]interface{}
170	if err := hcl.DecodeObject(&m, list); err != nil {
171		return err
172	}
173
174	// Manually parse
175	delete(m, "limit")
176
177	// Decode the rest
178	if err := mapstructure.WeakDecode(m, result); err != nil {
179		return err
180	}
181
182	// Parse limits
183	if o := list.Filter("limit"); len(o.Items) > 0 {
184		if err := parseQuotaLimits(&result.Limits, o); err != nil {
185			return multierror.Prefix(err, "limit ->")
186		}
187	}
188
189	return nil
190}
191
192// parseQuotaLimits parses the quota limits
193func parseQuotaLimits(result *[]*api.QuotaLimit, list *ast.ObjectList) error {
194	for _, o := range list.Elem().Items {
195		// Check for invalid keys
196		valid := []string{
197			"region",
198			"region_limit",
199		}
200		if err := helper.CheckHCLKeys(o.Val, valid); err != nil {
201			return err
202		}
203
204		var m map[string]interface{}
205		if err := hcl.DecodeObject(&m, o.Val); err != nil {
206			return err
207		}
208
209		// Manually parse
210		delete(m, "region_limit")
211
212		// Decode the rest
213		var limit api.QuotaLimit
214		if err := mapstructure.WeakDecode(m, &limit); err != nil {
215			return err
216		}
217
218		// We need this later
219		var listVal *ast.ObjectList
220		if ot, ok := o.Val.(*ast.ObjectType); ok {
221			listVal = ot.List
222		} else {
223			return fmt.Errorf("limit should be an object")
224		}
225
226		// Parse limits
227		if o := listVal.Filter("region_limit"); len(o.Items) > 0 {
228			limit.RegionLimit = new(api.Resources)
229			if err := parseQuotaResource(limit.RegionLimit, o); err != nil {
230				return multierror.Prefix(err, "region_limit ->")
231			}
232		}
233
234		*result = append(*result, &limit)
235	}
236
237	return nil
238}
239
240// parseQuotaResource parses the region_limit resources
241func parseQuotaResource(result *api.Resources, list *ast.ObjectList) error {
242	list = list.Elem()
243	if len(list.Items) == 0 {
244		return nil
245	}
246	if len(list.Items) > 1 {
247		return fmt.Errorf("only one 'region_limit' block allowed per limit")
248	}
249
250	// Get our resource object
251	o := list.Items[0]
252
253	// We need this later
254	var listVal *ast.ObjectList
255	if ot, ok := o.Val.(*ast.ObjectType); ok {
256		listVal = ot.List
257	} else {
258		return fmt.Errorf("resource: should be an object")
259	}
260
261	// Check for invalid keys
262	valid := []string{
263		"cpu",
264		"memory",
265		"network",
266	}
267	if err := helper.CheckHCLKeys(listVal, valid); err != nil {
268		return multierror.Prefix(err, "resources ->")
269	}
270
271	var m map[string]interface{}
272	if err := hcl.DecodeObject(&m, o.Val); err != nil {
273		return err
274	}
275
276	if err := mapstructure.WeakDecode(m, result); err != nil {
277		return err
278	}
279
280	// Find the network ObjectList, parse it
281	nw := listVal.Filter("network")
282	if len(nw.Items) > 0 {
283		rl, err := jobspec.ParseNetwork(nw)
284		if err != nil {
285			return multierror.Prefix(err, "resources ->")
286		}
287		if rl != nil {
288			if rl.Mode != "" || rl.HasPorts() {
289				return fmt.Errorf("resources -> network only allows mbits")
290			}
291			result.Networks = []*api.NetworkResource{rl}
292		}
293	}
294
295	return nil
296}
297