add tests

f0ssel · f0ssel · commit dbaeb2cf596f · 2022-05-18T16:44:12.000Z
diff --git a/coderd/httpmw/apikey.go b/coderd/httpmw/apikey.go
@@ -43,14 +43,20 @@ type OAuth2Configs struct {
 func ExtractAPIKey(db database.Store, oauth *OAuth2Configs) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			var cookieValue string
 			cookie, err := r.Cookie(AuthCookie)
 			if err != nil {
+				cookieValue = r.URL.Query().Get(AuthCookie)
+			} else {
+				cookieValue = cookie.Value
+			}
+			if cookieValue == "" {
 				httpapi.Write(rw, http.StatusUnauthorized, httpapi.Response{
-					Message: fmt.Sprintf("%q cookie must be provided", AuthCookie),
+					Message: fmt.Sprintf("%q cookie or query parameter must be provided", AuthCookie),
 				})
 				return
 			}
-			parts := strings.Split(cookie.Value, "-")
+			parts := strings.Split(cookieValue, "-")
 			// APIKeys are formatted: ID-SECRET
 			if len(parts) != 2 {
 				httpapi.Write(rw, http.StatusUnauthorized, httpapi.Response{
diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go
@@ -621,3 +621,28 @@ func mustLocation(t *testing.T, location string) *time.Location {
 
 	return loc
 }
+
+func TestWorkspaceWatcher(t *testing.T) {
+	t.Parallel()
+	client := coderdtest.New(t, nil)
+	user := coderdtest.CreateFirstUser(t, client)
+	coderdtest.NewProvisionerDaemon(t, client)
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, nil)
+	coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+	w, err := client.Workspace(context.Background(), workspace.ID)
+	require.NoError(t, err)
+
+	ww, err := client.WatchWorkspace(context.Background(), w.ID)
+	require.NoError(t, err)
+	defer ww.Close()
+	for i := 0; i < 5; i++ {
+		_, err := ww.Read(context.Background())
+		require.NoError(t, err)
+	}
+	err = ww.Close()
+	require.NoError(t, err)
+	_, err = ww.Read(context.Background())
+	require.Error(t, err)
+}
diff --git a/codersdk/client.go b/codersdk/client.go
@@ -81,8 +81,8 @@ func (c *Client) Request(ctx context.Context, method, path string, body interfac
 	return resp, err
 }
 
-// request performs an HTTP request with the body provided.
-// The caller is responsible for closing the response body.
+// websocket opens a websocket connection on that path provided.
+// The caller is responsible for closing the websocket.Conn.
 func (c *Client) websocket(ctx context.Context, path string) (*websocket.Conn, error) {
 	serverURL, err := c.URL.Parse(path)
 	if err != nil {
@@ -95,15 +95,9 @@ func (c *Client) websocket(ctx context.Context, path string) (*websocket.Conn, e
 		apiURL.Scheme = "wss"
 	}
 	apiURL.Path = path
-
-	client := &http.Client{
-		Jar: c.HTTPClient.Jar,
-	}
-	cookies := append(client.Jar.Cookies(c.URL), &http.Cookie{
-		Name:  httpmw.AuthCookie,
-		Value: c.SessionToken,
-	})
-	client.Jar.SetCookies(c.URL, cookies)
+	q := apiURL.Query()
+	q.Add(httpmw.AuthCookie, c.SessionToken)
+	apiURL.RawQuery = q.Encode()
 
 	//nolint:bodyclose
 	conn, _, err := websocket.Dial(context.Background(), apiURL.String(), &websocket.DialOptions{
diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go
@@ -123,7 +123,7 @@ func (w *WorkspaceWatcher) Close() error {
 	return nil
 }
 
-func (c *Client) WorkspaceWatcher(ctx context.Context, id uuid.UUID) (*WorkspaceWatcher, error) {
+func (c *Client) WatchWorkspace(ctx context.Context, id uuid.UUID) (*WorkspaceWatcher, error) {
 	conn, err := c.websocket(ctx, fmt.Sprintf("/api/v2/workspaces/%s/watch", id))
 	if err != nil {
 		return nil, err

Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,7 @@ func (w *WorkspaceWatcher) Close() error {`
`123`	`123`	`return nil`
`124`	`124`	`}`
`125`	`125`
`126`		`-func (c Client) WorkspaceWatcher(ctx context.Context, id uuid.UUID) (WorkspaceWatcher, error) {`
	`126`	`+func (c Client) WatchWorkspace(ctx context.Context, id uuid.UUID) (WorkspaceWatcher, error) {`
`127`	`127`	`conn, err := c.websocket(ctx, fmt.Sprintf("/api/v2/workspaces/%s/watch", id))`
`128`	`128`	`if err != nil {`
`129`	`129`	`return nil, err`