feat: add tailnet to agent RPC service

spikecurtis · spikecurtis · commit a66b15e8c129 · 2024-01-02T08:08:27.000+04:00
diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go
@@ -26,6 +26,7 @@ import (
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/tailnet"
+	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
 )
 
 const AgentAPIVersionDRPC = "2.0"
@@ -42,6 +43,7 @@ type API struct {
 	*AppsAPI
 	*MetadataAPI
 	*LogsAPI
+	*tailnet.DRPCService
 
 	mu                sync.Mutex
 	cachedWorkspaceID uuid.UUID
@@ -145,6 +147,13 @@ func New(opts Options) *API {
 		PublishWorkspaceAgentLogsUpdateFn: opts.PublishWorkspaceAgentLogsUpdateFn,
 	}
 
+	api.DRPCService = &tailnet.DRPCService{
+		CoordPtr:               opts.TailnetCoordinator,
+		Logger:                 opts.Log,
+		DerpMapUpdateFrequency: opts.DerpMapUpdateFrequency,
+		DerpMapFn:              opts.DerpMapFn,
+	}
+
 	return api
 }
 
@@ -155,6 +164,11 @@ func (a *API) Server(ctx context.Context) (*drpcserver.Server, error) {
 		return nil, xerrors.Errorf("register agent API protocol in DRPC mux: %w", err)
 	}
 
+	err = tailnetproto.DRPCRegisterTailnet(mux, a)
+	if err != nil {
+		return nil, xerrors.Errorf("register tailnet API protocol in DRPC mux: %w", err)
+	}
+
 	return drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux},
 		drpcserver.Options{
 			Log: func(err error) {
diff --git a/coderd/workspaceagentsrpc.go b/coderd/workspaceagentsrpc.go
@@ -3,6 +3,7 @@ package coderd
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"net/http"
 	"runtime/pprof"
 	"sync/atomic"
@@ -22,6 +23,7 @@ import (
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/coderd/util/ptr"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/tailnet"
 )
 
 // @Summary Workspace agent RPC API
@@ -128,6 +130,13 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		UpdateAgentMetricsFn: api.UpdateAgentMetrics,
 	})
 
+	streamID := tailnet.StreamID{
+		Name: fmt.Sprintf("%s-%s-%s", owner.Username, workspace.Name, workspaceAgent.Name),
+		ID:   workspaceAgent.ID,
+		Auth: tailnet.AgentTunnelAuth{},
+	}
+	ctx = tailnet.WithStreamID(ctx, streamID)
+
 	closeCtx, closeCtxCancel := context.WithCancel(ctx)
 	go func() {
 		defer closeCtxCancel()
