|
1 | 1 | package cli_test |
2 | 2 |
|
3 | 3 | import ( |
| 4 | +"bufio" |
4 | 5 | "bytes" |
5 | 6 | "context" |
6 | 7 | "crypto/ecdsa" |
@@ -338,6 +339,157 @@ func TestSSH(t *testing.T) { |
338 | 339 | <-cmdDone |
339 | 340 | }) |
340 | 341 |
|
| 342 | +t.Run("Stdio_StartStoppedWorkspace_CleanStdout",func(t*testing.T) { |
| 343 | +t.Parallel() |
| 344 | + |
| 345 | +authToken:=uuid.NewString() |
| 346 | +ownerClient:=coderdtest.New(t,&coderdtest.Options{IncludeProvisionerDaemon:true}) |
| 347 | +owner:=coderdtest.CreateFirstUser(t,ownerClient) |
| 348 | +client,_:=coderdtest.CreateAnotherUser(t,ownerClient,owner.OrganizationID,rbac.RoleTemplateAdmin()) |
| 349 | +version:=coderdtest.CreateTemplateVersion(t,client,owner.OrganizationID,&echo.Responses{ |
| 350 | +Parse:echo.ParseComplete, |
| 351 | +ProvisionPlan:echo.PlanComplete, |
| 352 | +ProvisionApply:echo.ProvisionApplyWithAgent(authToken), |
| 353 | +}) |
| 354 | +coderdtest.AwaitTemplateVersionJobCompleted(t,client,version.ID) |
| 355 | +template:=coderdtest.CreateTemplate(t,client,owner.OrganizationID,version.ID) |
| 356 | +workspace:=coderdtest.CreateWorkspace(t,client,owner.OrganizationID,template.ID) |
| 357 | +coderdtest.AwaitWorkspaceBuildJobCompleted(t,client,workspace.LatestBuild.ID) |
| 358 | +// Stop the workspace |
| 359 | +workspaceBuild:=coderdtest.CreateWorkspaceBuild(t,client,workspace,database.WorkspaceTransitionStop) |
| 360 | +coderdtest.AwaitWorkspaceBuildJobCompleted(t,client,workspaceBuild.ID) |
| 361 | + |
| 362 | +ctx,cancel:=context.WithTimeout(context.Background(),testutil.WaitLong) |
| 363 | +defercancel() |
| 364 | + |
| 365 | +clientStdinR,clientStdinW:=io.Pipe() |
| 366 | +// Here's a simple flowchart for how these pipes are used: |
| 367 | +// |
| 368 | +// flowchart LR |
| 369 | +// A[ProxyCommand] --> B[captureProxyCommandStdoutW] |
| 370 | +// B --> C[captureProxyCommandStdoutR] |
| 371 | +// C --> VA[Validate output] |
| 372 | +// C --> D[proxyCommandStdoutW] |
| 373 | +// D --> E[proxyCommandStdoutR] |
| 374 | +// E --> F[SSH Client] |
| 375 | +proxyCommandStdoutR,proxyCommandStdoutW:=io.Pipe() |
| 376 | +captureProxyCommandStdoutR,captureProxyCommandStdoutW:=io.Pipe() |
| 377 | +closePipes:=func() { |
| 378 | +for_,c:=range []io.Closer{clientStdinR,clientStdinW,proxyCommandStdoutR,proxyCommandStdoutW,captureProxyCommandStdoutR,captureProxyCommandStdoutW} { |
| 379 | +_=c.Close() |
| 380 | +} |
| 381 | +} |
| 382 | +deferclosePipes() |
| 383 | +tGo(t,func() { |
| 384 | +<-ctx.Done() |
| 385 | +closePipes() |
| 386 | +}) |
| 387 | + |
| 388 | +// Here we start a monitor for the output produced by the proxy command, |
| 389 | +// which is read by the SSH client. This is done to validate that the |
| 390 | +// output is clean. |
| 391 | +proxyCommandOutputBuf:=make(chanbyte,4096) |
| 392 | +tGo(t,func() { |
| 393 | +deferclose(proxyCommandOutputBuf) |
| 394 | + |
| 395 | +gotHeader:=false |
| 396 | +buf:= bytes.Buffer{} |
| 397 | +r:=bufio.NewReader(captureProxyCommandStdoutR) |
| 398 | +for { |
| 399 | +b,err:=r.ReadByte() |
| 400 | +iferr!=nil { |
| 401 | +iferrors.Is(err,io.ErrClosedPipe) { |
| 402 | +return |
| 403 | +} |
| 404 | +assert.NoError(t,err,"read byte failed") |
| 405 | +return |
| 406 | +} |
| 407 | +ifb=='\n'||b=='\r' { |
| 408 | +out:=buf.Bytes() |
| 409 | +t.Logf("monitorServerOutput: %q (%#x)",out,out) |
| 410 | +buf.Reset() |
| 411 | + |
| 412 | +// Ideally we would do further verification, but that would |
| 413 | +// involve parsing the SSH protocol to look for output that |
| 414 | +// doesn't belong. This at least ensures that no garbage is |
| 415 | +// being sent to the SSH client before trying to connect. |
| 416 | +if!gotHeader { |
| 417 | +gotHeader=true |
| 418 | +assert.Equal(t,"SSH-2.0-Go",string(out),"invalid header") |
| 419 | +} |
| 420 | +}else { |
| 421 | +_=buf.WriteByte(b) |
| 422 | +} |
| 423 | +select { |
| 424 | +caseproxyCommandOutputBuf<-b: |
| 425 | +case<-ctx.Done(): |
| 426 | +return |
| 427 | +} |
| 428 | +} |
| 429 | +}) |
| 430 | +tGo(t,func() { |
| 431 | +deferproxyCommandStdoutW.Close() |
| 432 | + |
| 433 | +// Range closed by above goroutine. |
| 434 | +forb:=rangeproxyCommandOutputBuf { |
| 435 | +_,err:=proxyCommandStdoutW.Write([]byte{b}) |
| 436 | +iferr!=nil { |
| 437 | +iferrors.Is(err,io.ErrClosedPipe) { |
| 438 | +return |
| 439 | +} |
| 440 | +assert.NoError(t,err,"write byte failed") |
| 441 | +return |
| 442 | +} |
| 443 | +} |
| 444 | +}) |
| 445 | + |
| 446 | +// Start the SSH stdio command. |
| 447 | +inv,root:=clitest.New(t,"ssh","--stdio",workspace.Name) |
| 448 | +clitest.SetupConfig(t,client,root) |
| 449 | +inv.Stdin=clientStdinR |
| 450 | +inv.Stdout=captureProxyCommandStdoutW |
| 451 | +inv.Stderr=io.Discard |
| 452 | + |
| 453 | +cmdDone:=tGo(t,func() { |
| 454 | +err:=inv.WithContext(ctx).Run() |
| 455 | +assert.NoError(t,err) |
| 456 | +}) |
| 457 | + |
| 458 | +tGo(t,func() { |
| 459 | +// When the agent connects, the workspace was started, and we should |
| 460 | +// have access to the shell. |
| 461 | +_=agenttest.New(t,client.URL,authToken) |
| 462 | +coderdtest.NewWorkspaceAgentWaiter(t,client,workspace.ID).Wait() |
| 463 | +}) |
| 464 | + |
| 465 | +conn,channels,requests,err:=ssh.NewClientConn(&stdioConn{ |
| 466 | +Reader:proxyCommandStdoutR, |
| 467 | +Writer:clientStdinW, |
| 468 | +},"",&ssh.ClientConfig{ |
| 469 | +// #nosec |
| 470 | +HostKeyCallback:ssh.InsecureIgnoreHostKey(), |
| 471 | +}) |
| 472 | +require.NoError(t,err) |
| 473 | +deferconn.Close() |
| 474 | + |
| 475 | +sshClient:=ssh.NewClient(conn,channels,requests) |
| 476 | +session,err:=sshClient.NewSession() |
| 477 | +require.NoError(t,err) |
| 478 | +defersession.Close() |
| 479 | + |
| 480 | +command:="sh -c exit" |
| 481 | +ifruntime.GOOS=="windows" { |
| 482 | +command="cmd.exe /c exit" |
| 483 | +} |
| 484 | +err=session.Run(command) |
| 485 | +require.NoError(t,err) |
| 486 | +err=sshClient.Close() |
| 487 | +require.NoError(t,err) |
| 488 | +_=clientStdinR.Close() |
| 489 | + |
| 490 | +<-cmdDone |
| 491 | +}) |
| 492 | + |
341 | 493 | t.Run("Stdio_RemoteForward_Signal",func(t*testing.T) { |
342 | 494 | t.Parallel() |
343 | 495 | client,workspace,agentToken:=setupWorkspaceForAgent(t) |
|