1package execagent
2
3import (
4	"fmt"
5	"io"
6	"io/ioutil"
7	"net"
8	"os"
9	"os/exec"
10	"path/filepath"
11	"text/template"
12
13	"github.com/hashicorp/nomad/api"
14)
15
16type AgentMode int
17
18const (
19	// Conf enum is for configuring either a client, server, or mixed agent.
20	ModeClient AgentMode = 1
21	ModeServer AgentMode = 2
22	ModeBoth             = ModeClient | ModeServer
23)
24
25func init() {
26	if d := os.Getenv("NOMAD_TEST_DIR"); d != "" {
27		BaseDir = d
28	}
29}
30
31var (
32	// BaseDir is where tests will store state and can be overridden by
33	// setting NOMAD_TEST_DIR. Defaults to "/opt/nomadtest"
34	BaseDir = "/opt/nomadtest"
35
36	agentTemplate = template.Must(template.New("agent").Parse(`
37enable_debug = true
38log_level = "{{ or .LogLevel "DEBUG" }}"
39
40ports {
41  http = {{.HTTP}}
42  rpc  = {{.RPC}}
43  serf = {{.Serf}}
44}
45
46{{ if .EnableServer }}
47server {
48  enabled = true
49  bootstrap_expect = 1
50}
51{{ end }}
52
53{{ if .EnableClient }}
54client {
55  enabled = true
56  options = {
57    "driver.raw_exec.enable" = "1"
58  }
59}
60{{ end }}
61`))
62)
63
64type AgentTemplateVars struct {
65	HTTP         int
66	RPC          int
67	Serf         int
68	EnableClient bool
69	EnableServer bool
70	LogLevel     string
71}
72
73func newAgentTemplateVars() (*AgentTemplateVars, error) {
74	httpPort, err := getFreePort()
75	if err != nil {
76		return nil, err
77	}
78	rpcPort, err := getFreePort()
79	if err != nil {
80		return nil, err
81	}
82	serfPort, err := getFreePort()
83	if err != nil {
84		return nil, err
85	}
86
87	vars := AgentTemplateVars{
88		HTTP: httpPort,
89		RPC:  rpcPort,
90		Serf: serfPort,
91	}
92
93	return &vars, nil
94}
95
96func writeConfig(path string, vars *AgentTemplateVars) error {
97	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
98	if err != nil {
99		return err
100	}
101	defer f.Close()
102	return agentTemplate.Execute(f, vars)
103}
104
105// NomadAgent manages an external Nomad agent process.
106type NomadAgent struct {
107	// BinPath is the path to the Nomad binary
108	BinPath string
109
110	// DataDir is the path state will be saved in
111	DataDir string
112
113	// ConfFile is the path to the agent's conf file
114	ConfFile string
115
116	// Cmd is the agent process
117	Cmd *exec.Cmd
118
119	// Vars are the config parameters used to template
120	Vars *AgentTemplateVars
121}
122
123// NewMixedAgent creates a new Nomad agent in mixed server+client mode but does
124// not start the agent process until the Start() method is called.
125func NewMixedAgent(bin string) (*NomadAgent, error) {
126	if err := os.MkdirAll(BaseDir, 0755); err != nil {
127		return nil, err
128	}
129	dir, err := ioutil.TempDir(BaseDir, "agent")
130	if err != nil {
131		return nil, err
132	}
133
134	vars, err := newAgentTemplateVars()
135	if err != nil {
136		return nil, err
137	}
138	vars.EnableClient = true
139	vars.EnableServer = true
140
141	conf := filepath.Join(dir, "config.hcl")
142	if err := writeConfig(conf, vars); err != nil {
143		return nil, err
144	}
145
146	na := &NomadAgent{
147		BinPath:  bin,
148		DataDir:  dir,
149		ConfFile: conf,
150		Vars:     vars,
151		Cmd:      exec.Command(bin, "agent", "-config", conf, "-data-dir", dir),
152	}
153	return na, nil
154}
155
156// NewClientServerPair creates a pair of Nomad agents: 1 server, 1 client.
157func NewClientServerPair(bin string, serverOut, clientOut io.Writer) (
158	server *NomadAgent, client *NomadAgent, err error) {
159
160	if err := os.MkdirAll(BaseDir, 0755); err != nil {
161		return nil, nil, err
162	}
163
164	sdir, err := ioutil.TempDir(BaseDir, "server")
165	if err != nil {
166		return nil, nil, err
167	}
168
169	svars, err := newAgentTemplateVars()
170	if err != nil {
171		return nil, nil, err
172	}
173	svars.LogLevel = "WARN"
174	svars.EnableServer = true
175
176	sconf := filepath.Join(sdir, "config.hcl")
177	if err := writeConfig(sconf, svars); err != nil {
178		return nil, nil, err
179	}
180
181	server = &NomadAgent{
182		BinPath:  bin,
183		DataDir:  sdir,
184		ConfFile: sconf,
185		Vars:     svars,
186		Cmd:      exec.Command(bin, "agent", "-config", sconf, "-data-dir", sdir),
187	}
188	server.Cmd.Stdout = serverOut
189	server.Cmd.Stderr = serverOut
190
191	cdir, err := ioutil.TempDir(BaseDir, "client")
192	if err != nil {
193		return nil, nil, err
194	}
195
196	cvars, err := newAgentTemplateVars()
197	if err != nil {
198		return nil, nil, err
199	}
200	cvars.EnableClient = true
201
202	cconf := filepath.Join(cdir, "config.hcl")
203	if err := writeConfig(cconf, cvars); err != nil {
204		return nil, nil, err
205	}
206
207	client = &NomadAgent{
208		BinPath:  bin,
209		DataDir:  cdir,
210		ConfFile: cconf,
211		Vars:     cvars,
212		Cmd: exec.Command(bin, "agent",
213			"-config", cconf,
214			"-data-dir", cdir,
215			"-servers", fmt.Sprintf("127.0.0.1:%d", svars.RPC),
216		),
217	}
218	client.Cmd.Stdout = clientOut
219	client.Cmd.Stderr = clientOut
220	return
221}
222
223// Start the agent command.
224func (n *NomadAgent) Start() error {
225	return n.Cmd.Start()
226}
227
228// Stop sends an interrupt signal and returns the command's Wait error.
229func (n *NomadAgent) Stop() error {
230	if err := n.Cmd.Process.Signal(os.Interrupt); err != nil {
231		return err
232	}
233
234	return n.Cmd.Wait()
235}
236
237// Destroy stops the agent and removes the data dir.
238func (n *NomadAgent) Destroy() error {
239	if err := n.Stop(); err != nil {
240		return err
241	}
242	return os.RemoveAll(n.DataDir)
243}
244
245// Client returns an api.Client for the agent.
246func (n *NomadAgent) Client() (*api.Client, error) {
247	conf := api.DefaultConfig()
248	conf.Address = fmt.Sprintf("http://127.0.0.1:%d", n.Vars.HTTP)
249	return api.NewClient(conf)
250}
251
252func getFreePort() (int, error) {
253	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
254	if err != nil {
255		return 0, err
256	}
257
258	l, err := net.ListenTCP("tcp", addr)
259	if err != nil {
260		return 0, err
261	}
262	defer l.Close()
263	return l.Addr().(*net.TCPAddr).Port, nil
264}
265