Skip to content

Commit e4b51dc

Browse files
kylecarbskylecarbs
committed
Add database functions for authenticating the agent
1 parent 74d8ae4 commit e4b51dc

File tree

9 files changed

+404
-1
lines changed

9 files changed

+404
-1
lines changed

‎agent/agent.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"golang.org/x/xerrors"
2626
)
2727

28+
// DialSSH creates an SSH DataChannel on the connection provided.
2829
func DialSSH(conn *peer.Conn) (net.Conn, error) {
2930
channel, err := conn.Dial(context.Background(), "ssh", &peer.ChannelOptions{
3031
Protocol: "ssh",
@@ -35,6 +36,7 @@ func DialSSH(conn *peer.Conn) (net.Conn, error) {
3536
return channel.NetConn(), nil
3637
}
3738

39+
// DialSSHClient wraps the DialSSH function with a Go SSH client.
3840
func DialSSHClient(conn *peer.Conn) (*gossh.Client, error) {
3941
netConn, err := DialSSH(conn)
4042
if err != nil {
@@ -59,6 +61,13 @@ type Options struct {
5961
Logger slog.Logger
6062
}
6163

64+
// Dialer to return the peerbroker.Listener, but that hinges on
65+
// a proper authentication token. If it fails to dial for that
66+
// reason, we should check the API error and renegotiate a new
67+
// authentication method.
68+
//
69+
// This also needs to update it's own metadata and such.
70+
6271
type Dialer func(ctx context.Context) (*peerbroker.Listener, error)
6372

6473
func New(dialer Dialer, options *Options) io.Closer {

‎coderd/workspaceagent.go

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
package coderd
2+
3+
import (
4+
"database/sql"
5+
"fmt"
6+
"io"
7+
"net/http"
8+
"time"
9+
10+
"github.com/hashicorp/yamux"
11+
"nhooyr.io/websocket"
12+
13+
"cdr.dev/slog"
14+
"github.com/coder/coder/database"
15+
"github.com/coder/coder/httpapi"
16+
"github.com/coder/coder/httpmw"
17+
"github.com/coder/coder/peerbroker"
18+
"github.com/coder/coder/peerbroker/proto"
19+
"github.com/coder/coder/provisionersdk"
20+
)
21+
22+
type AgentResourceMetadata struct {
23+
MemoryTotal uint64 `json:"memory_total"`
24+
DiskTotal uint64 `json:"disk_total"`
25+
CPUCores uint64 `json:"cpu_cores"`
26+
CPUModel string `json:"cpu_model"`
27+
CPUMhz float64 `json:"cpu_mhz"`
28+
}
29+
30+
type AgentInstanceMetadata struct {
31+
JailOrchestrator string `json:"jail_orchestrator"`
32+
OperatingSystem string `json:"operating_system"`
33+
Platform string `json:"platform"`
34+
PlatformFamily string `json:"platform_family"`
35+
KernelVersion string `json:"kernel_version"`
36+
KernelArchitecture string `json:"kernel_architecture"`
37+
Cloud string `json:"cloud"`
38+
Jail string `json:"jail"`
39+
VNC bool `json:"vnc"`
40+
}
41+
42+
func (api *api) workspaceAgentUpdate() {
43+
44+
}
45+
46+
func (api *api) workspaceAgentConnectByResource(rw http.ResponseWriter, r *http.Request) {
47+
api.websocketWaitGroup.Add(1)
48+
defer api.websocketWaitGroup.Done()
49+
50+
agent := httpmw.WorkspaceAgent(r)
51+
if !agent.UpdatedAt.Valid {
52+
httpapi.Write(rw, http.StatusPreconditionRequired, httpapi.Response{
53+
Message: "Agent hasn't connected yet!",
54+
})
55+
return
56+
}
57+
58+
conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
59+
CompressionMode: websocket.CompressionDisabled,
60+
})
61+
if err != nil {
62+
httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
63+
Message: fmt.Sprintf("accept websocket: %s", err),
64+
})
65+
return
66+
}
67+
defer func() {
68+
_ = conn.Close(websocket.StatusNormalClosure, "")
69+
}()
70+
config := yamux.DefaultConfig()
71+
config.LogOutput = io.Discard
72+
session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), config)
73+
if err != nil {
74+
_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
75+
return
76+
}
77+
err = peerbroker.ProxyListen(r.Context(), session, peerbroker.ProxyOptions{
78+
ChannelID: resource.WorkspaceAgentID.UUID.String(),
79+
Logger: api.Logger.Named("peerbroker-proxy-dial"),
80+
Pubsub: api.Pubsub,
81+
})
82+
if err != nil {
83+
_ = conn.Close(websocket.StatusInternalError, fmt.Sprintf("serve: %s", err))
84+
return
85+
}
86+
}
87+
88+
func (api *api) workspaceAgentServe(rw http.ResponseWriter, r *http.Request) {
89+
api.websocketWaitGroup.Add(1)
90+
defer api.websocketWaitGroup.Done()
91+
92+
workspaceAgent := httpmw.WorkspaceAgent(r)
93+
conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
94+
CompressionMode: websocket.CompressionDisabled,
95+
})
96+
if err != nil {
97+
httpapi.Write(rw, http.StatusBadRequest, httpapi.Response{
98+
Message: fmt.Sprintf("accept websocket: %s", err),
99+
})
100+
return
101+
}
102+
defer func() {
103+
_ = conn.Close(websocket.StatusNormalClosure, "")
104+
}()
105+
config := yamux.DefaultConfig()
106+
config.LogOutput = io.Discard
107+
session, err := yamux.Server(websocket.NetConn(r.Context(), conn, websocket.MessageBinary), config)
108+
if err != nil {
109+
_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
110+
return
111+
}
112+
closer, err := peerbroker.ProxyDial(proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(session)), peerbroker.ProxyOptions{
113+
ChannelID: workspaceAgent.ID.String(),
114+
Pubsub: api.Pubsub,
115+
Logger: api.Logger.Named("peerbroker-proxy-listen"),
116+
})
117+
if err != nil {
118+
_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
119+
return
120+
}
121+
122+
err = api.Database.UpdateWorkspaceAgentByID(r.Context(), database.UpdateWorkspaceAgentByIDParams{
123+
ID: workspaceAgent.ID,
124+
UpdatedAt: sql.NullTime{
125+
Time: database.Now(),
126+
Valid: true,
127+
},
128+
})
129+
if err != nil {
130+
_ = conn.Close(websocket.StatusAbnormalClosure, err.Error())
131+
return
132+
}
133+
defer closer.Close()
134+
ticker := time.NewTicker(5 * time.Second)
135+
for {
136+
select {
137+
case <-ticker.C:
138+
err = api.Database.UpdateWorkspaceAgentByID(r.Context(), database.UpdateWorkspaceAgentByIDParams{
139+
ID: workspaceAgent.ID,
140+
UpdatedAt: sql.NullTime{
141+
Time: database.Now(),
142+
Valid: true,
143+
},
144+
})
145+
if err != nil {
146+
api.Logger.Error(r.Context(), "update workspace agent by id", slog.Error(err), slog.F("id", workspaceAgent.ID.String()))
147+
return
148+
}
149+
case <-r.Context().Done():
150+
return
151+
}
152+
}
153+
}

‎database/databasefake/databasefake.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,6 +522,18 @@ func (q *fakeQuerier) GetProvisionerJobAgentByInstanceID(_ context.Context, inst
522522
return database.ProvisionerJobAgent{}, sql.ErrNoRows
523523
}
524524

525+
func (q *fakeQuerier) GetProvisionerJobAgentByAuthToken(ctx context.Context, authToken uuid.UUID) (database.ProvisionerJobAgent, error) {
526+
q.mutex.Lock()
527+
defer q.mutex.Unlock()
528+
529+
for _, agent := range q.provisionerJobAgent {
530+
if agent.AuthToken.String() == authToken.String() {
531+
return agent, nil
532+
}
533+
}
534+
return database.ProvisionerJobAgent{}, sql.ErrNoRows
535+
}
536+
525537
func (q *fakeQuerier) GetProvisionerJobAgentsByResourceIDs(_ context.Context, ids []uuid.UUID) ([]database.ProvisionerJobAgent, error) {
526538
q.mutex.Lock()
527539
defer q.mutex.Unlock()
@@ -956,6 +968,23 @@ func (q *fakeQuerier) UpdateProvisionerDaemonByID(_ context.Context, arg databas
956968
return sql.ErrNoRows
957969
}
958970

971+
func (q *fakeQuerier) UpdateProvisionerJobAgentByID(ctx context.Context, arg database.UpdateProvisionerJobAgentByIDParams) error {
972+
q.mutex.Lock()
973+
defer q.mutex.Unlock()
974+
975+
for index, agent := range q.provisionerJobAgent {
976+
if arg.ID.String() != agent.ID.String() {
977+
continue
978+
}
979+
agent.UpdatedAt = arg.UpdatedAt
980+
agent.InstanceMetadata = arg.InstanceMetadata
981+
agent.ResourceMetadata = arg.ResourceMetadata
982+
q.provisionerJobAgent[index] = agent
983+
return nil
984+
}
985+
return sql.ErrNoRows
986+
}
987+
959988
func (q *fakeQuerier) UpdateProvisionerJobByID(_ context.Context, arg database.UpdateProvisionerJobByIDParams) error {
960989
q.mutex.Lock()
961990
defer q.mutex.Unlock()

‎database/querier.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎database/query.sql

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,14 @@ SELECT
226226
FROM
227227
provisioner_daemon;
228228

229+
-- name: GetProvisionerJobAgentByAuthToken :one
230+
SELECT
231+
*
232+
FROM
233+
provisioner_job_agent
234+
WHERE
235+
auth_token = $1;
236+
229237
-- name: GetProvisionerJobAgentByInstanceID :one
230238
SELECT
231239
*
@@ -624,6 +632,16 @@ SET
624632
WHERE
625633
id = $1;
626634

635+
-- name: UpdateProvisionerJobAgentByID :exec
636+
UPDATE
637+
provisioner_job_agent
638+
SET
639+
updated_at = $2,
640+
instance_metadata = $3,
641+
resource_metadata = $4
642+
WHERE
643+
id = $1;
644+
627645
-- name: UpdateProvisionerDaemonByID :exec
628646
UPDATE
629647
provisioner_daemon

‎database/query.sql.go

Lines changed: 55 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ require (
3737
github.com/hashicorp/go-version v1.4.0
3838
github.com/hashicorp/terraform-config-inspect v0.0.0-20211115214459-90acf1ca460f
3939
github.com/hashicorp/terraform-exec v0.15.0
40+
github.com/hashicorp/terraform-json v0.13.0
4041
github.com/hashicorp/terraform-plugin-sdk/v2 v2.10.1
4142
github.com/hashicorp/yamux v0.0.0-20211028200310-0bc27b27de87
4243
github.com/justinas/nosurf v1.1.1
@@ -109,7 +110,6 @@ require (
109110
github.com/hashicorp/hcl v1.0.0 // indirect
110111
github.com/hashicorp/hcl/v2 v2.11.1 // indirect
111112
github.com/hashicorp/logutils v1.0.0 // indirect
112-
github.com/hashicorp/terraform-json v0.13.0 // indirect
113113
github.com/hashicorp/terraform-plugin-go v0.5.0 // indirect
114114
github.com/hashicorp/terraform-plugin-log v0.2.0 // indirect
115115
github.com/hashicorp/terraform-registry-address v0.0.0-20210412075316-9b2996cce896 // indirect

0 commit comments

Comments
 (0)