github.com/clly/consul@v1.4.5/agent/remote_exec_test.go (about) 1 package agent 2 3 import ( 4 "bytes" 5 "encoding/json" 6 "fmt" 7 "reflect" 8 "testing" 9 10 "time" 11 12 "github.com/hashicorp/consul/agent/structs" 13 "github.com/hashicorp/consul/api" 14 "github.com/hashicorp/consul/testrpc" 15 "github.com/hashicorp/consul/testutil/retry" 16 "github.com/hashicorp/go-uuid" 17 ) 18 19 func generateUUID() (ret string) { 20 var err error 21 if ret, err = uuid.GenerateUUID(); err != nil { 22 panic(fmt.Sprintf("Unable to generate a UUID, %v", err)) 23 } 24 return ret 25 } 26 27 func TestRexecWriter(t *testing.T) { 28 // t.Parallel() // timing test. no parallel 29 writer := &rexecWriter{ 30 BufCh: make(chan []byte, 16), 31 BufSize: 16, 32 BufIdle: 100 * time.Millisecond, 33 CancelCh: make(chan struct{}), 34 } 35 36 // Write short, wait for idle 37 start := time.Now() 38 n, err := writer.Write([]byte("test")) 39 if err != nil { 40 t.Fatalf("err: %v", err) 41 } 42 if n != 4 { 43 t.Fatalf("bad: %v", n) 44 } 45 46 select { 47 case b := <-writer.BufCh: 48 if len(b) != 4 { 49 t.Fatalf("Bad: %v", b) 50 } 51 if time.Since(start) < writer.BufIdle { 52 t.Fatalf("too early") 53 } 54 case <-time.After(2 * writer.BufIdle): 55 t.Fatalf("timeout") 56 } 57 58 // Write in succession to prevent the timeout 59 writer.Write([]byte("test")) 60 time.Sleep(writer.BufIdle / 2) 61 writer.Write([]byte("test")) 62 time.Sleep(writer.BufIdle / 2) 63 start = time.Now() 64 writer.Write([]byte("test")) 65 66 select { 67 case b := <-writer.BufCh: 68 if len(b) != 12 { 69 t.Fatalf("Bad: %v", b) 70 } 71 if time.Since(start) < writer.BufIdle { 72 t.Fatalf("too early") 73 } 74 case <-time.After(2 * writer.BufIdle): 75 t.Fatalf("timeout") 76 } 77 78 // Write large values, multiple flushes required 79 writer.Write([]byte("01234567890123456789012345678901")) 80 81 select { 82 case b := <-writer.BufCh: 83 if string(b) != "0123456789012345" { 84 t.Fatalf("bad: %s", b) 85 } 86 default: 87 t.Fatalf("should have buf") 88 } 89 select { 90 case b := <-writer.BufCh: 91 if string(b) != "6789012345678901" { 92 t.Fatalf("bad: %s", b) 93 } 94 default: 95 t.Fatalf("should have buf") 96 } 97 } 98 99 func TestRemoteExecGetSpec(t *testing.T) { 100 t.Parallel() 101 testRemoteExecGetSpec(t, "", "", true, "") 102 } 103 104 func TestRemoteExecGetSpec_ACLToken(t *testing.T) { 105 t.Parallel() 106 dc := "dc1" 107 testRemoteExecGetSpec(t, ` 108 acl_datacenter = "`+dc+`" 109 acl_master_token = "root" 110 acl_token = "root" 111 acl_default_policy = "deny" 112 `, "root", true, dc) 113 } 114 115 func TestRemoteExecGetSpec_ACLAgentToken(t *testing.T) { 116 t.Parallel() 117 dc := "dc1" 118 testRemoteExecGetSpec(t, ` 119 acl_datacenter = "`+dc+`" 120 acl_master_token = "root" 121 acl_agent_token = "root" 122 acl_default_policy = "deny" 123 `, "root", true, dc) 124 } 125 126 func TestRemoteExecGetSpec_ACLDeny(t *testing.T) { 127 t.Parallel() 128 dc := "dc1" 129 testRemoteExecGetSpec(t, ` 130 acl_datacenter = "`+dc+`" 131 acl_master_token = "root" 132 acl_default_policy = "deny" 133 `, "root", false, dc) 134 } 135 136 func testRemoteExecGetSpec(t *testing.T, hcl string, token string, shouldSucceed bool, dc string) { 137 a := NewTestAgent(t, t.Name(), hcl) 138 defer a.Shutdown() 139 if dc != "" { 140 testrpc.WaitForLeader(t, a.RPC, dc) 141 } else { 142 testrpc.WaitForTestAgent(t, a.RPC, "dc1") 143 } 144 event := &remoteExecEvent{ 145 Prefix: "_rexec", 146 Session: makeRexecSession(t, a.Agent, token), 147 } 148 defer destroySession(t, a.Agent, event.Session, token) 149 150 spec := &remoteExecSpec{ 151 Command: "uptime", 152 Script: []byte("#!/bin/bash"), 153 Wait: time.Second, 154 } 155 buf, err := json.Marshal(spec) 156 if err != nil { 157 t.Fatalf("err: %v", err) 158 } 159 key := "_rexec/" + event.Session + "/job" 160 setKV(t, a.Agent, key, buf, token) 161 162 var out remoteExecSpec 163 if shouldSucceed != a.remoteExecGetSpec(event, &out) { 164 t.Fatalf("bad") 165 } 166 if shouldSucceed && !reflect.DeepEqual(spec, &out) { 167 t.Fatalf("bad spec") 168 } 169 } 170 171 func TestRemoteExecWrites(t *testing.T) { 172 t.Parallel() 173 testRemoteExecWrites(t, "", "", true, "") 174 } 175 176 func TestRemoteExecWrites_ACLToken(t *testing.T) { 177 t.Parallel() 178 dc := "dc1" 179 testRemoteExecWrites(t, ` 180 acl_datacenter = "`+dc+`" 181 acl_master_token = "root" 182 acl_token = "root" 183 acl_default_policy = "deny" 184 `, "root", true, dc) 185 } 186 187 func TestRemoteExecWrites_ACLAgentToken(t *testing.T) { 188 t.Parallel() 189 dc := "dc1" 190 testRemoteExecWrites(t, ` 191 acl_datacenter = "`+dc+`" 192 acl_master_token = "root" 193 acl_agent_token = "root" 194 acl_default_policy = "deny" 195 `, "root", true, dc) 196 } 197 198 func TestRemoteExecWrites_ACLDeny(t *testing.T) { 199 t.Parallel() 200 dc := "dc1" 201 testRemoteExecWrites(t, ` 202 acl_datacenter = "`+dc+`" 203 acl_master_token = "root" 204 acl_default_policy = "deny" 205 `, "root", false, dc) 206 } 207 208 func testRemoteExecWrites(t *testing.T, hcl string, token string, shouldSucceed bool, dc string) { 209 a := NewTestAgent(t, t.Name(), hcl) 210 defer a.Shutdown() 211 if dc != "" { 212 testrpc.WaitForLeader(t, a.RPC, dc) 213 } else { 214 // For slow machines, ensure we wait a bit 215 testrpc.WaitForLeader(t, a.RPC, "dc1") 216 } 217 event := &remoteExecEvent{ 218 Prefix: "_rexec", 219 Session: makeRexecSession(t, a.Agent, token), 220 } 221 defer destroySession(t, a.Agent, event.Session, token) 222 223 if shouldSucceed != a.remoteExecWriteAck(event) { 224 t.Fatalf("bad") 225 } 226 227 output := []byte("testing") 228 if shouldSucceed != a.remoteExecWriteOutput(event, 0, output) { 229 t.Fatalf("bad") 230 } 231 if shouldSucceed != a.remoteExecWriteOutput(event, 10, output) { 232 t.Fatalf("bad") 233 } 234 235 // Bypass the remaining checks if the write was expected to fail. 236 if !shouldSucceed { 237 return 238 } 239 240 exitCode := 1 241 if !a.remoteExecWriteExitCode(event, &exitCode) { 242 t.Fatalf("bad") 243 } 244 245 key := "_rexec/" + event.Session + "/" + a.Config.NodeName + "/ack" 246 d := getKV(t, a.Agent, key, token) 247 if d == nil || d.Session != event.Session { 248 t.Fatalf("bad ack: %#v", d) 249 } 250 251 key = "_rexec/" + event.Session + "/" + a.Config.NodeName + "/out/00000" 252 d = getKV(t, a.Agent, key, token) 253 if d == nil || d.Session != event.Session || !bytes.Equal(d.Value, output) { 254 t.Fatalf("bad output: %#v", d) 255 } 256 257 key = "_rexec/" + event.Session + "/" + a.Config.NodeName + "/out/0000a" 258 d = getKV(t, a.Agent, key, token) 259 if d == nil || d.Session != event.Session || !bytes.Equal(d.Value, output) { 260 t.Fatalf("bad output: %#v", d) 261 } 262 263 key = "_rexec/" + event.Session + "/" + a.Config.NodeName + "/exit" 264 d = getKV(t, a.Agent, key, token) 265 if d == nil || d.Session != event.Session || string(d.Value) != "1" { 266 t.Fatalf("bad output: %#v", d) 267 } 268 } 269 270 func testHandleRemoteExec(t *testing.T, command string, expectedSubstring string, expectedReturnCode string) { 271 a := NewTestAgent(t, t.Name(), "") 272 defer a.Shutdown() 273 testrpc.WaitForTestAgent(t, a.RPC, "dc1") 274 retry.Run(t, func(r *retry.R) { 275 event := &remoteExecEvent{ 276 Prefix: "_rexec", 277 Session: makeRexecSession(t, a.Agent, ""), 278 } 279 defer destroySession(t, a.Agent, event.Session, "") 280 281 spec := &remoteExecSpec{ 282 Command: command, 283 Wait: time.Second, 284 } 285 buf, err := json.Marshal(spec) 286 if err != nil { 287 t.Fatalf("err: %v", err) 288 } 289 key := "_rexec/" + event.Session + "/job" 290 setKV(t, a.Agent, key, buf, "") 291 292 buf, err = json.Marshal(event) 293 if err != nil { 294 t.Fatalf("err: %v", err) 295 } 296 msg := &UserEvent{ 297 ID: generateUUID(), 298 Payload: buf, 299 } 300 301 // Handle the event... 302 a.handleRemoteExec(msg) 303 304 // Verify we have an ack 305 key = "_rexec/" + event.Session + "/" + a.Config.NodeName + "/ack" 306 d := getKV(t, a.Agent, key, "") 307 if d == nil || d.Session != event.Session { 308 t.Fatalf("bad ack: %#v", d) 309 } 310 311 // Verify we have output 312 key = "_rexec/" + event.Session + "/" + a.Config.NodeName + "/out/00000" 313 d = getKV(t, a.Agent, key, "") 314 if d == nil || d.Session != event.Session || 315 !bytes.Contains(d.Value, []byte(expectedSubstring)) { 316 t.Fatalf("bad output: %#v", d) 317 } 318 319 // Verify we have an exit code 320 key = "_rexec/" + event.Session + "/" + a.Config.NodeName + "/exit" 321 d = getKV(t, a.Agent, key, "") 322 if d == nil || d.Session != event.Session || string(d.Value) != expectedReturnCode { 323 t.Fatalf("bad output: %#v", d) 324 } 325 }) 326 } 327 328 func TestHandleRemoteExec(t *testing.T) { 329 t.Parallel() 330 testHandleRemoteExec(t, "uptime", "load", "0") 331 } 332 333 func TestHandleRemoteExecFailed(t *testing.T) { 334 t.Parallel() 335 testHandleRemoteExec(t, "echo failing;exit 2", "failing", "2") 336 } 337 338 func makeRexecSession(t *testing.T, a *Agent, token string) string { 339 args := structs.SessionRequest{ 340 Datacenter: a.config.Datacenter, 341 Op: structs.SessionCreate, 342 Session: structs.Session{ 343 Node: a.config.NodeName, 344 LockDelay: 15 * time.Second, 345 }, 346 WriteRequest: structs.WriteRequest{ 347 Token: token, 348 }, 349 } 350 var out string 351 if err := a.RPC("Session.Apply", &args, &out); err != nil { 352 t.Fatalf("err: %v", err) 353 } 354 return out 355 } 356 357 func destroySession(t *testing.T, a *Agent, session string, token string) { 358 args := structs.SessionRequest{ 359 Datacenter: a.config.Datacenter, 360 Op: structs.SessionDestroy, 361 Session: structs.Session{ 362 ID: session, 363 }, 364 WriteRequest: structs.WriteRequest{ 365 Token: token, 366 }, 367 } 368 var out string 369 if err := a.RPC("Session.Apply", &args, &out); err != nil { 370 t.Fatalf("err: %v", err) 371 } 372 } 373 374 func setKV(t *testing.T, a *Agent, key string, val []byte, token string) { 375 write := structs.KVSRequest{ 376 Datacenter: a.config.Datacenter, 377 Op: api.KVSet, 378 DirEnt: structs.DirEntry{ 379 Key: key, 380 Value: val, 381 }, 382 WriteRequest: structs.WriteRequest{ 383 Token: token, 384 }, 385 } 386 var success bool 387 if err := a.RPC("KVS.Apply", &write, &success); err != nil { 388 t.Fatalf("err: %v", err) 389 } 390 } 391 392 func getKV(t *testing.T, a *Agent, key string, token string) *structs.DirEntry { 393 req := structs.KeyRequest{ 394 Datacenter: a.config.Datacenter, 395 Key: key, 396 QueryOptions: structs.QueryOptions{ 397 Token: token, 398 }, 399 } 400 var out structs.IndexedDirEntries 401 if err := a.RPC("KVS.Get", &req, &out); err != nil { 402 t.Fatalf("err: %v", err) 403 } 404 if len(out.Entries) > 0 { 405 return out.Entries[0] 406 } 407 return nil 408 }