|
| 1 | +package coderd |
| 2 | + |
| 3 | +import ( |
| 4 | +"bytes" |
| 5 | +"context" |
| 6 | +"database/sql" |
| 7 | +"fmt" |
| 8 | +"io" |
| 9 | +"net/http" |
| 10 | +"net/http/httptest" |
| 11 | +"net/http/httputil" |
| 12 | +"net/url" |
| 13 | +"strings" |
| 14 | +"testing" |
| 15 | + |
| 16 | +"github.com/go-chi/chi/v5" |
| 17 | +"github.com/google/uuid" |
| 18 | +"github.com/stretchr/testify/require" |
| 19 | +"go.uber.org/mock/gomock" |
| 20 | + |
| 21 | +"cdr.dev/slog" |
| 22 | +"cdr.dev/slog/sloggers/slogtest" |
| 23 | +"github.com/coder/coder/v2/coderd/database" |
| 24 | +"github.com/coder/coder/v2/coderd/database/dbmock" |
| 25 | +"github.com/coder/coder/v2/coderd/database/dbtime" |
| 26 | +"github.com/coder/coder/v2/coderd/httpmw" |
| 27 | +"github.com/coder/coder/v2/coderd/workspaceapps/appurl" |
| 28 | +"github.com/coder/coder/v2/codersdk" |
| 29 | +"github.com/coder/coder/v2/codersdk/workspacesdk" |
| 30 | +"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" |
| 31 | +"github.com/coder/coder/v2/codersdk/wsjson" |
| 32 | +"github.com/coder/coder/v2/tailnet" |
| 33 | +"github.com/coder/coder/v2/tailnet/tailnettest" |
| 34 | +"github.com/coder/coder/v2/testutil" |
| 35 | +"github.com/coder/websocket" |
| 36 | +) |
| 37 | + |
| 38 | +typefakeAgentProviderstruct { |
| 39 | +agentConnfunc(ctx context.Context,agentID uuid.UUID) (_ workspacesdk.AgentConn,releasefunc(),_error) |
| 40 | +} |
| 41 | + |
| 42 | +func (fakeAgentProvider)ReverseProxy(targetURL,dashboardURL*url.URL,agentID uuid.UUID,app appurl.ApplicationURL,wildcardHoststring)*httputil.ReverseProxy { |
| 43 | +panic("unimplemented") |
| 44 | +} |
| 45 | + |
| 46 | +func (ffakeAgentProvider)AgentConn(ctx context.Context,agentID uuid.UUID) (_ workspacesdk.AgentConn,releasefunc(),_error) { |
| 47 | +iff.agentConn!=nil { |
| 48 | +returnf.agentConn(ctx,agentID) |
| 49 | +} |
| 50 | + |
| 51 | +panic("unimplemented") |
| 52 | +} |
| 53 | + |
| 54 | +func (fakeAgentProvider)ServeHTTPDebug(w http.ResponseWriter,r*http.Request) { |
| 55 | +panic("unimplemented") |
| 56 | +} |
| 57 | + |
| 58 | +func (fakeAgentProvider)Close()error { |
| 59 | +returnnil |
| 60 | +} |
| 61 | + |
| 62 | +funcTestWatchAgentContainers(t*testing.T) { |
| 63 | +t.Parallel() |
| 64 | + |
| 65 | +t.Run("WebSocketClosesProperly",func(t*testing.T) { |
| 66 | +t.Parallel() |
| 67 | + |
| 68 | +// This test ensures that the agent containers `/watch` websocket can gracefully |
| 69 | +// handle the underlying websocket unexpectedly closing. This test was created in |
| 70 | +// response to this issue: https://github.com/coder/coder/issues/19372 |
| 71 | + |
| 72 | +var ( |
| 73 | +ctx=testutil.Context(t,testutil.WaitShort) |
| 74 | +logger=slogtest.Make(t,&slogtest.Options{IgnoreErrors:true}).Leveled(slog.LevelDebug).Named("coderd") |
| 75 | + |
| 76 | +mCtrl=gomock.NewController(t) |
| 77 | +mDB=dbmock.NewMockStore(mCtrl) |
| 78 | +mCoordinator=tailnettest.NewMockCoordinator(mCtrl) |
| 79 | +mAgentConn=agentconnmock.NewMockAgentConn(mCtrl) |
| 80 | + |
| 81 | +fAgentProvider=fakeAgentProvider{ |
| 82 | +agentConn:func(ctx context.Context,agentID uuid.UUID) (_ workspacesdk.AgentConn,releasefunc(),_error) { |
| 83 | +returnmAgentConn,func() {},nil |
| 84 | +}, |
| 85 | +} |
| 86 | + |
| 87 | +workspaceID=uuid.New() |
| 88 | +agentID=uuid.New() |
| 89 | +resourceID=uuid.New() |
| 90 | +jobID=uuid.New() |
| 91 | +buildID=uuid.New() |
| 92 | + |
| 93 | +containersCh=make(chan codersdk.WorkspaceAgentListContainersResponse) |
| 94 | + |
| 95 | +r=chi.NewMux() |
| 96 | + |
| 97 | +api=API{ |
| 98 | +ctx:ctx, |
| 99 | +Options:&Options{ |
| 100 | +AgentInactiveDisconnectTimeout:testutil.WaitShort, |
| 101 | +Database:mDB, |
| 102 | +Logger:logger, |
| 103 | +DeploymentValues:&codersdk.DeploymentValues{}, |
| 104 | +TailnetCoordinator:tailnettest.NewFakeCoordinator(), |
| 105 | +}, |
| 106 | +} |
| 107 | +) |
| 108 | + |
| 109 | +vartailnetCoordinator tailnet.Coordinator=mCoordinator |
| 110 | +api.TailnetCoordinator.Store(&tailnetCoordinator) |
| 111 | +api.agentProvider=fAgentProvider |
| 112 | + |
| 113 | +// Setup: Allow `ExtractWorkspaceAgentParams` to complete. |
| 114 | +mDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(),agentID).Return(database.WorkspaceAgent{ |
| 115 | +ID:agentID, |
| 116 | +ResourceID:resourceID, |
| 117 | +LifecycleState:database.WorkspaceAgentLifecycleStateReady, |
| 118 | +FirstConnectedAt: sql.NullTime{Valid:true,Time:dbtime.Now()}, |
| 119 | +LastConnectedAt: sql.NullTime{Valid:true,Time:dbtime.Now()}, |
| 120 | +},nil) |
| 121 | +mDB.EXPECT().GetWorkspaceResourceByID(gomock.Any(),resourceID).Return(database.WorkspaceResource{ |
| 122 | +ID:resourceID, |
| 123 | +JobID:jobID, |
| 124 | +},nil) |
| 125 | +mDB.EXPECT().GetProvisionerJobByID(gomock.Any(),jobID).Return(database.ProvisionerJob{ |
| 126 | +ID:jobID, |
| 127 | +Type:database.ProvisionerJobTypeWorkspaceBuild, |
| 128 | +},nil) |
| 129 | +mDB.EXPECT().GetWorkspaceBuildByJobID(gomock.Any(),jobID).Return(database.WorkspaceBuild{ |
| 130 | +WorkspaceID:workspaceID, |
| 131 | +ID:buildID, |
| 132 | +},nil) |
| 133 | + |
| 134 | +// And: Allow `db2dsk.WorkspaceAgent` to complete. |
| 135 | +mCoordinator.EXPECT().Node(gomock.Any()).Return(nil) |
| 136 | + |
| 137 | +// And: Allow `WatchContainers` to be called, returing our `containersCh` channel. |
| 138 | +mAgentConn.EXPECT().WatchContainers(gomock.Any(),gomock.Any()). |
| 139 | +Return(containersCh,io.NopCloser(&bytes.Buffer{}),nil) |
| 140 | + |
| 141 | +// And: We mount the HTTP Handler |
| 142 | +r.With(httpmw.ExtractWorkspaceAgentParam(mDB)). |
| 143 | +Get("/workspaceagents/{workspaceagent}/containers/watch",api.watchWorkspaceAgentContainers) |
| 144 | + |
| 145 | +// Given: We create the HTTP server |
| 146 | +srv:=httptest.NewServer(r) |
| 147 | +defersrv.Close() |
| 148 | + |
| 149 | +// And: Dial the WebSocket |
| 150 | +wsURL:=strings.Replace(srv.URL,"http://","ws://",1) |
| 151 | +conn,resp,err:=websocket.Dial(ctx,fmt.Sprintf("%s/workspaceagents/%s/containers/watch",wsURL,agentID),nil) |
| 152 | +require.NoError(t,err) |
| 153 | +ifresp.Body!=nil { |
| 154 | +deferresp.Body.Close() |
| 155 | +} |
| 156 | + |
| 157 | +// And: Create a streaming decoder |
| 158 | +decoder:=wsjson.NewDecoder[codersdk.WorkspaceAgentListContainersResponse](conn,websocket.MessageText,logger) |
| 159 | +deferdecoder.Close() |
| 160 | +decodeCh:=decoder.Chan() |
| 161 | + |
| 162 | +// And: We can successfully send through the channel. |
| 163 | +testutil.RequireSend(ctx,t,containersCh, codersdk.WorkspaceAgentListContainersResponse{ |
| 164 | +Containers: []codersdk.WorkspaceAgentContainer{{ |
| 165 | +ID:"test-container-id", |
| 166 | +}}, |
| 167 | +}) |
| 168 | + |
| 169 | +// And: Receive the data. |
| 170 | +containerResp:=testutil.RequireReceive(ctx,t,decodeCh) |
| 171 | +require.Len(t,containerResp.Containers,1) |
| 172 | +require.Equal(t,"test-container-id",containerResp.Containers[0].ID) |
| 173 | + |
| 174 | +// When: We close the `containersCh` |
| 175 | +close(containersCh) |
| 176 | + |
| 177 | +// Then: We expect `decodeCh` to be closed. |
| 178 | +select { |
| 179 | +case<-ctx.Done(): |
| 180 | +t.Fail() |
| 181 | + |
| 182 | +case_,ok:=<-decodeCh: |
| 183 | +require.False(t,ok,"channel is expected to be closed") |
| 184 | +} |
| 185 | +}) |
| 186 | +} |