Make WorkspaceWatcher

f0ssel · f0ssel · commit da9c744fafed · 2022-05-18T16:44:12.000Z
diff --git a/codersdk/client.go b/codersdk/client.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"golang.org/x/xerrors"
+	"nhooyr.io/websocket"
 
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
@@ -80,6 +81,41 @@ 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.
+func (c *Client) websocket(ctx context.Context, path string) (*websocket.Conn, error) {
+	serverURL, err := c.URL.Parse(path)
+	if err != nil {
+		return nil, xerrors.Errorf("parse url: %w", err)
+	}
+
+	apiURL, err := url.Parse(serverURL.String())
+	apiURL.Scheme = "ws"
+	if serverURL.Scheme == "https" {
+		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)
+
+	//nolint:bodyclose
+	conn, _, err := websocket.Dial(context.Background(), apiURL.String(), &websocket.DialOptions{
+		HTTPClient: c.HTTPClient,
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("dial websocket: %w", err)
+	}
+
+	return conn, nil
+}
+
 // readBodyAsError reads the response as an httpapi.Message, and
 // wraps it in a codersdk.Error type for easy marshaling.
 func readBodyAsError(res *http.Response) error {
diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go
@@ -9,6 +9,8 @@ import (
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
+	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/wsjson"
 
 	"github.com/coder/coder/coderd/database"
 )
@@ -98,6 +100,40 @@ func (c *Client) WorkspaceBuildByName(ctx context.Context, workspace uuid.UUID,
 	return workspaceBuild, json.NewDecoder(res.Body).Decode(&workspaceBuild)
 }
 
+type WorkspaceWatcher struct {
+	conn *websocket.Conn
+}
+
+func (w *WorkspaceWatcher) Read(ctx context.Context) (Workspace, error) {
+	var ws Workspace
+	err := wsjson.Read(ctx, w.conn, &ws)
+	if err != nil {
+		return ws, xerrors.Errorf("read workspace: %w")
+	}
+
+	return ws, nil
+}
+
+func (w *WorkspaceWatcher) Close() error {
+	err := w.conn.Close(websocket.StatusNormalClosure, "")
+	if err != nil {
+		return xerrors.Errorf("closing workspace watcher: %w", err)
+	}
+
+	return nil
+}
+
+func (c *Client) WorkspaceWatcher(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
+	}
+
+	return &WorkspaceWatcher{
+		conn: conn,
+	}, nil
+}
+
 // UpdateWorkspaceAutostartRequest is a request to update a workspace's autostart schedule.
 type UpdateWorkspaceAutostartRequest struct {
 	Schedule string `json:"schedule"`
