fix: use typed wireguard public keys in database structs (#2639)

coadler · web-flow · commit 26e85b0bbc06 · 2022-06-24T15:45:28.000-05:00
diff --git a/coderd/database/dbtypes/dbtypes.go b/coderd/database/dbtypes/dbtypes.go
@@ -0,0 +1,74 @@
+package dbtypes
+
+import (
+	"database/sql/driver"
+
+	"golang.org/x/xerrors"
+	"tailscale.com/types/key"
+)
+
+// NodePublic is a wrapper around a key.NodePublic which represents the
+// Wireguard public key for an agent..
+type NodePublic key.NodePublic
+
+func (n NodePublic) String() string {
+	return key.NodePublic(n).String()
+}
+
+// This is necessary so NodePublic can be serialized in JSON loggers.
+func (n NodePublic) MarshalJSON() ([]byte, error) {
+	j, err := key.NodePublic(n).MarshalText()
+	// surround in quotes to make it a JSON string
+	j = append([]byte{'"'}, append(j, '"')...)
+	return j, err
+}
+
+// Value is so NodePublic can be inserted into the database.
+func (n NodePublic) Value() (driver.Value, error) {
+	return key.NodePublic(n).MarshalText()
+}
+
+// Scan is so NodePublic can be read from the database.
+func (n *NodePublic) Scan(value interface{}) error {
+	switch v := value.(type) {
+	case []byte:
+		return (*key.NodePublic)(n).UnmarshalText(v)
+	case string:
+		return (*key.NodePublic)(n).UnmarshalText([]byte(v))
+	default:
+		return xerrors.Errorf("unexpected type: %T", v)
+	}
+}
+
+// NodePublic is a wrapper around a key.NodePublic which represents the
+// Tailscale disco key for an agent.
+type DiscoPublic key.DiscoPublic
+
+func (n DiscoPublic) String() string {
+	return key.DiscoPublic(n).String()
+}
+
+// This is necessary so DiscoPublic can be serialized in JSON loggers.
+func (n DiscoPublic) MarshalJSON() ([]byte, error) {
+	j, err := key.DiscoPublic(n).MarshalText()
+	// surround in quotes to make it a JSON string
+	j = append([]byte{'"'}, append(j, '"')...)
+	return j, err
+}
+
+// Value is so DiscoPublic can be inserted into the database.
+func (n DiscoPublic) Value() (driver.Value, error) {
+	return key.DiscoPublic(n).MarshalText()
+}
+
+// Scan is so DiscoPublic can be read from the database.
+func (n *DiscoPublic) Scan(value interface{}) error {
+	switch v := value.(type) {
+	case []byte:
+		return (*key.DiscoPublic)(n).UnmarshalText(v)
+	case string:
+		return (*key.DiscoPublic)(n).UnmarshalText([]byte(v))
+	default:
+		return xerrors.Errorf("unexpected type: %T", v)
+	}
+}
diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/sqlc.yaml b/coderd/database/sqlc.yaml
@@ -17,10 +17,10 @@ packages:
     output_db_file_name: db_tmp.go
 
 overrides:
-  - column: workspaces.wireguard_public_key
-    go_type: tailscale.com/types/key.MachinePublic
-  - column: workspaces.disco_public_key
-    go_type: tailscale.com/types/key.DiscoPublic
+  - column: workspace_agents.wireguard_node_public_key
+    go_type: github.com/coder/coder/coderd/database/dbtypes.NodePublic
+  - column: workspace_agents.wireguard_disco_public_key
+    go_type: github.com/coder/coder/coderd/database/dbtypes.DiscoPublic
 
 rename:
   api_key: APIKey
diff --git a/coderd/provisionerdaemons.go b/coderd/provisionerdaemons.go
@@ -19,11 +19,11 @@ import (
 	protobuf "google.golang.org/protobuf/proto"
 	"storj.io/drpc/drpcmux"
 	"storj.io/drpc/drpcserver"
-	"tailscale.com/types/key"
 
 	"cdr.dev/slog"
 
 	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/database/dbtypes"
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/parameter"
 	"github.com/coder/coder/coderd/rbac"
@@ -761,8 +761,8 @@ func insertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.
 				Valid:  prAgent.StartupScript != "",
 			},
 			WireguardNodeIPv6:       peerwg.UUIDToInet(agentID),
-			WireguardNodePublicKey:  key.NodePublic{}.String(),
-			WireguardDiscoPublicKey: key.DiscoPublic{}.String(),
+			WireguardNodePublicKey:  dbtypes.NodePublic{},
+			WireguardDiscoPublicKey: dbtypes.DiscoPublic{},
 		})
 		if err != nil {
 			return xerrors.Errorf("insert agent: %w", err)
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -22,6 +22,7 @@ import (
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/coderd/database"
+	"github.com/coder/coder/coderd/database/dbtypes"
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/rbac"
@@ -488,8 +489,8 @@ func (api *API) postWorkspaceAgentKeys(rw http.ResponseWriter, r *http.Request)
 
 	err := api.Database.UpdateWorkspaceAgentKeysByID(ctx, database.UpdateWorkspaceAgentKeysByIDParams{
 		ID:                      workspaceAgent.ID,
-		WireguardNodePublicKey:  keys.Public.String(),
-		WireguardDiscoPublicKey: keys.Disco.String(),
+		WireguardNodePublicKey:  dbtypes.NodePublic(keys.Public),
+		WireguardDiscoPublicKey: dbtypes.DiscoPublic(keys.Disco),
 	})
 	if err != nil {
 		httpapi.Write(rw, http.StatusInternalServerError, httpapi.Response{
@@ -711,15 +712,8 @@ func convertWorkspaceAgent(dbAgent database.WorkspaceAgent, apps []codersdk.Work
 		Directory:            dbAgent.Directory,
 		Apps:                 apps,
 		IPv6:                 inetToNetaddr(dbAgent.WireguardNodeIPv6),
-	}
-
-	err := workspaceAgent.WireguardPublicKey.UnmarshalText([]byte(dbAgent.WireguardNodePublicKey))
-	if err != nil {
-		return codersdk.WorkspaceAgent{}, xerrors.Errorf("unmarshal wireguard node public key %q: %w", dbAgent.WireguardNodePublicKey, err)
-	}
-	err = workspaceAgent.DiscoPublicKey.UnmarshalText([]byte(dbAgent.WireguardDiscoPublicKey))
-	if err != nil {
-		return codersdk.WorkspaceAgent{}, xerrors.Errorf("unmarshal disco public key %q: %w", dbAgent.WireguardDiscoPublicKey, err)
+		WireguardPublicKey:   key.NodePublic(dbAgent.WireguardNodePublicKey),
+		DiscoPublicKey:       key.DiscoPublic(dbAgent.WireguardDiscoPublicKey),
 	}
 
 	if dbAgent.FirstConnectedAt.Valid {
