Add database functions for authenticating the agent

coder · kylecarbs · Feb 24, 2022 · Feb 25, 2022 · Feb 25, 2022 · Feb 26, 2022
commit e4b51dc03c4982c47ae59e91b1a9190bd4cdb37c
diff --git a/agent/agent.go b/agent/agent.go
@@ -25,6 +25,7 @@ import (
 	"golang.org/x/xerrors"
 )
 
+// DialSSH creates an SSH DataChannel on the connection provided.
 func DialSSH(conn *peer.Conn) (net.Conn, error) {
 	channel, err := conn.Dial(context.Background(), "ssh", &peer.ChannelOptions{
 		Protocol: "ssh",
@@ -35,6 +36,7 @@ func DialSSH(conn *peer.Conn) (net.Conn, error) {
 	return channel.NetConn(), nil
 }
 
+// DialSSHClient wraps the DialSSH function with a Go SSH client.
 func DialSSHClient(conn *peer.Conn) (*gossh.Client, error) {
 	netConn, err := DialSSH(conn)
 	if err != nil {
@@ -59,6 +61,13 @@ type Options struct {
 	Logger slog.Logger
 }
 
+// Dialer to return the peerbroker.Listener, but that hinges on
+// a proper authentication token. If it fails to dial for that
+// reason, we should check the API error and renegotiate a new
+// authentication method.
+//
+// This also needs to update it's own metadata and such.
+
 type Dialer func(ctx context.Context) (*peerbroker.Listener, error)
 
 func New(dialer Dialer, options *Options) io.Closer {

diff --git a/coderd/workspaceagent.go b/coderd/workspaceagent.go
@@ -0,0 +1,153 @@
+package coderd
+
+import (
+	"database/sql"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+
+	"github.com/hashicorp/yamux"
+	"nhooyr.io/websocket"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/database"
+	"github.com/coder/coder/httpapi"
+	"github.com/coder/coder/httpmw"
+	"github.com/coder/coder/peerbroker"
+	"github.com/coder/coder/peerbroker/proto"
+	"github.com/coder/coder/provisionersdk"
+)
+
+type AgentResourceMetadata struct {
+	MemoryTotal uint64  `json:"memory_total"`
+	DiskTotal   uint64  `json:"disk_total"`
+	CPUCores    uint64  `json:"cpu_cores"`
+	CPUModel    string  `json:"cpu_model"`
+	CPUMhz      float64 `json:"cpu_mhz"`
+}
+
+type AgentInstanceMetadata struct {
+	JailOrchestrator   string `json:"jail_orchestrator"`
+	OperatingSystem    string `json:"operating_system"`
+	Platform           string `json:"platform"`
+	PlatformFamily     string `json:"platform_family"`
+	KernelVersion      string `json:"kernel_version"`
+	KernelArchitecture string `json:"kernel_architecture"`
+	Cloud              string `json:"cloud"`
+	Jail               string `json:"jail"`
+	VNC                bool   `json:"vnc"`
+}
+
+func (api *api) workspaceAgentUpdate() {
+
+}
+
+func (api *api) workspaceAgentConnectByResource(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitGroup.Add(1)
+	defer api.websocketWaitGroup.Done()
+
+	agent := httpmw.WorkspaceAgent(r)
+	if !agent.UpdatedAt.Valid {
+		httpapi.Write(rw, http.StatusPreconditionRequired, httpapi.Response{
+			Message: "Agent hasn't connected yet!",
+		})
+		return
+	}
+
+	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
+			Message: fmt.Sprintf("accept websocket: %s", err),
+		})
+		return
+	}
+	defer func() {
+		_ = conn.Close(websocket.StatusNormalClosure, "")
+	}()
+	config := yamux.DefaultConfig()
+	config.LogOutput = io.Discard
+	session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), config)
+	if err != nil {
+		_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
+		return
+	}
+	err = peerbroker.ProxyListen(r.Context(), session, peerbroker.ProxyOptions{
+		ChannelID: resource.WorkspaceAgentID.UUID.String(),
+		Logger:    api.Logger.Named("peerbroker-proxy-dial"),
+		Pubsub:    api.Pubsub,
+	})
+	if err != nil {
+		_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("serve: %s", err))
+		return
+	}
+}
+
+func (api *api) workspaceAgentServe(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitGroup.Add(1)
+	defer api.websocketWaitGroup.Done()
+
+	workspaceAgent := httpmw.WorkspaceAgent(r)
+	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
+			Message: fmt.Sprintf("accept websocket: %s", err),
+		})
+		return
+	}
+	defer func() {
+		_ = conn.Close(websocket.StatusNormalClosure, "")
+	}()
+	config := yamux.DefaultConfig()
+	config.LogOutput = io.Discard
+	session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), config)
+	if err != nil {
+		_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
+		return
+	}
+	closer, err := peerbroker.ProxyDial(proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(session)), peerbroker.ProxyOptions{
+		ChannelID: workspaceAgent.ID.String(),
+		Pubsub:    api.Pubsub,
+		Logger:    api.Logger.Named("peerbroker-proxy-listen"),
+	})
+	if err != nil {
+		_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
+		return
+	}
+
+	err = api.Database.UpdateWorkspaceAgentByID(r.Context(), database.UpdateWorkspaceAgentByIDParams{
+		ID: workspaceAgent.ID,
+		UpdatedAt: sql.NullTime{
+			Time:  database.Now(),
+			Valid: true,
+		},
+	})
+	if err != nil {
+		_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
+		return
+	}
+	defer closer.Close()
+	ticker := time.NewTicker(5 * time.Second)
+	for {
+		select {
+		case <-ticker.C:
+			err = api.Database.UpdateWorkspaceAgentByID(r.Context(), database.UpdateWorkspaceAgentByIDParams{
+				ID: workspaceAgent.ID,
+				UpdatedAt: sql.NullTime{
+					Time:  database.Now(),
+					Valid: true,
+				},
+			})
+			if err != nil {
+				api.Logger.Error(r.Context(), "update workspace agent by id", slog.Error(err), slog.F("id", workspaceAgent.ID.String()))
+				return
+			}
+		case <-r.Context().Done():
+			return
+		}
+	}
+}
diff --git a/database/databasefake/databasefake.go b/database/databasefake/databasefake.go
@@ -522,6 +522,18 @@ func (q *fakeQuerier) GetProvisionerJobAgentByInstanceID(_ context.Context, inst
 	return database.ProvisionerJobAgent{}, sql.ErrNoRows
 }
 
+func (q *fakeQuerier) GetProvisionerJobAgentByAuthToken(ctx context.Context, authToken uuid.UUID) (database.ProvisionerJobAgent, error) {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	for _, agent := range q.provisionerJobAgent {
+		if agent.AuthToken.String() == authToken.String() {
+			return agent, nil
+		}
+	}
+	return database.ProvisionerJobAgent{}, sql.ErrNoRows
+}
+
 func (q *fakeQuerier) GetProvisionerJobAgentsByResourceIDs(_ context.Context, ids []uuid.UUID) ([]database.ProvisionerJobAgent, error) {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
@@ -956,6 +968,23 @@ func (q *fakeQuerier) UpdateProvisionerDaemonByID(_ context.Context, arg databas
 	return sql.ErrNoRows
 }
 
+func (q *fakeQuerier) UpdateProvisionerJobAgentByID(ctx context.Context, arg database.UpdateProvisionerJobAgentByIDParams) error {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	for index, agent := range q.provisionerJobAgent {
+		if arg.ID.String() != agent.ID.String() {
+			continue
+		}
+		agent.UpdatedAt = arg.UpdatedAt
+		agent.InstanceMetadata = arg.InstanceMetadata
+		agent.ResourceMetadata = arg.ResourceMetadata
+		q.provisionerJobAgent[index] = agent
+		return nil
+	}
+	return sql.ErrNoRows
+}
+
 func (q *fakeQuerier) UpdateProvisionerJobByID(_ context.Context, arg database.UpdateProvisionerJobByIDParams) error {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()

diff --git a/database/querier.go b/database/querier.go
diff --git a/database/query.sql b/database/query.sql
@@ -226,6 +226,14 @@ SELECT
 FROM
   provisioner_daemon;
 
+-- name: GetProvisionerJobAgentByAuthToken :one
+SELECT
+  *
+FROM
+  provisioner_job_agent
+WHERE
+  auth_token = $1;
+
 -- name: GetProvisionerJobAgentByInstanceID :one
 SELECT
   *
@@ -624,6 +632,16 @@ SET
 WHERE
   id = $1;
 
+-- name: UpdateProvisionerJobAgentByID :exec
+UPDATE
+  provisioner_job_agent
+SET
+  updated_at = $2,
+  instance_metadata = $3,
+  resource_metadata = $4
+WHERE
+  id = $1;
+
 -- name: UpdateProvisionerDaemonByID :exec
 UPDATE
   provisioner_daemon

diff --git a/database/query.sql.go b/database/query.sql.go
diff --git a/go.mod b/go.mod
@@ -37,6 +37,7 @@ require (
 	github.com/hashicorp/go-version v1.4.0
 	github.com/hashicorp/terraform-config-inspect v0.0.0-20211115214459-90acf1ca460f
 	github.com/hashicorp/terraform-exec v0.15.0
+	github.com/hashicorp/terraform-json v0.13.0
 	github.com/hashicorp/terraform-plugin-sdk/v2 v2.10.1
 	github.com/hashicorp/yamux v0.0.0-20211028200310-0bc27b27de87
 	github.com/justinas/nosurf v1.1.1
@@ -109,7 +110,6 @@ require (
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/hashicorp/hcl/v2 v2.11.1 // indirect
 	github.com/hashicorp/logutils v1.0.0 // indirect
-	github.com/hashicorp/terraform-json v0.13.0 // indirect
 	github.com/hashicorp/terraform-plugin-go v0.5.0 // indirect
 	github.com/hashicorp/terraform-plugin-log v0.2.0 // indirect
 	github.com/hashicorp/terraform-registry-address v0.0.0-20210412075316-9b2996cce896 // indirect