coder · coadler · Aug 15, 2023 · Aug 14, 2023 · Aug 14, 2023
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -407,6 +407,7 @@ func New(options *Options) *API {
 				return (*api.TailnetCoordinator.Load()).ServeMultiAgent(uuid.New()), nil
 			},
 			wsconncache.New(api._dialWorkspaceAgentTailnet, 0),
+			api.TracerProvider,
 		)
 		if err != nil {
 			panic("failed to setup server tailnet: " + err.Error())

diff --git a/coderd/tailnet.go b/coderd/tailnet.go
@@ -14,11 +14,13 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/xerrors"
 	"tailscale.com/derp"
 	"tailscale.com/tailcfg"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/coderd/tracing"
 	"github.com/coder/coder/coderd/wsconncache"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/site"
@@ -45,6 +47,7 @@ func NewServerTailnet(
 	derpMap *tailcfg.DERPMap,
 	getMultiAgent func(context.Context) (tailnet.MultiAgentConn, error),
 	cache *wsconncache.Cache,
+	traceProvider trace.TracerProvider,
 ) (*ServerTailnet, error) {
 	logger = logger.Named("servertailnet")
 	conn, err := tailnet.NewConn(&tailnet.Options{
@@ -58,15 +61,16 @@ func NewServerTailnet(
 
 	serverCtx, cancel := context.WithCancel(ctx)
 	tn := &ServerTailnet{
-		ctx:           serverCtx,
-		cancel:        cancel,
-		logger:        logger,
-		conn:          conn,
-		getMultiAgent: getMultiAgent,
-		cache:         cache,
-		agentNodes:    map[uuid.UUID]time.Time{},
-		agentTickets:  map[uuid.UUID]map[uuid.UUID]struct{}{},
-		transport:     tailnetTransport.Clone(),
+		ctx:                  serverCtx,
+		cancel:               cancel,
+		logger:               logger,
+		tracer:               traceProvider.Tracer(tracing.TracerName),
+		conn:                 conn,
+		getMultiAgent:        getMultiAgent,
+		cache:                cache,
+		agentConnectionTimes: map[uuid.UUID]time.Time{},
+		agentTickets:         map[uuid.UUID]map[uuid.UUID]struct{}{},
+		transport:            tailnetTransport.Clone(),
 	}
 	tn.transport.DialContext = tn.dialContext
 	tn.transport.MaxIdleConnsPerHost = 10
@@ -139,25 +143,50 @@ func (s *ServerTailnet) expireOldAgents() {
 		case <-ticker.C:
 		}
 
-		s.nodesMu.Lock()
-		agentConn := s.getAgentConn()
-		for agentID, lastConnection := range s.agentNodes {
-			// If no one has connected since the cutoff and there are no active
-			// connections, remove the agent.
-			if time.Since(lastConnection) > cutoff && len(s.agentTickets[agentID]) == 0 {
-				_ = agentConn
-				// err := agentConn.UnsubscribeAgent(agentID)
-				// if err != nil {
-				// 	s.logger.Error(s.ctx, "unsubscribe expired agent", slog.Error(err), slog.F("agent_id", agentID))
-				// }
-				// delete(s.agentNodes, agentID)
-
-				// TODO(coadler): actually remove from the netmap, then reenable
-				// the above
+		s.doExpireOldAgents(cutoff)
+	}
+}
+
+func (s *ServerTailnet) doExpireOldAgents(cutoff time.Duration) {
+	// TODO: add some attrs to this.
+	ctx, span := s.tracer.Start(s.ctx, tracing.FuncName())
+	defer span.End()
+
+	start := time.Now()
+	deletedCount := 0
+
+	s.nodesMu.Lock()
+	s.logger.Debug(ctx, "pruning inactive agents", slog.F("agent_count", len(s.agentConnectionTimes)))
+	agentConn := s.getAgentConn()
+	for agentID, lastConnection := range s.agentConnectionTimes {
+		// If no one has connected since the cutoff and there are no active
+		// connections, remove the agent.
+		if time.Since(lastConnection) > cutoff && len(s.agentTickets[agentID]) == 0 {
+			deleted, err := s.conn.RemovePeer(tailnet.PeerSelector{
+				ID: tailnet.NodeID(agentID),
+				IP: netip.PrefixFrom(tailnet.IPFromUUID(agentID), 128),
+			})
+			if err != nil {
+				s.logger.Warn(ctx, "failed to remove peer from server tailnet", slog.Error(err))
+				continue
+			}
+			if !deleted {
+				s.logger.Warn(ctx, "peer didn't exist in tailnet", slog.Error(err))
+			}
+
+			deletedCount++
+			delete(s.agentConnectionTimes, agentID)
+			err = agentConn.UnsubscribeAgent(agentID)
+			if err != nil {
+				s.logger.Error(ctx, "unsubscribe expired agent", slog.Error(err), slog.F("agent_id", agentID))
 			}
 		}
-		s.nodesMu.Unlock()
 	}
+	s.nodesMu.Unlock()
+	s.logger.Debug(s.ctx, "successfully pruned inactive agents",
+		slog.F("deleted", deletedCount),
+		slog.F("took", time.Since(start)),
+	)
 }
 
 func (s *ServerTailnet) watchAgentUpdates() {
@@ -196,7 +225,7 @@ func (s *ServerTailnet) reinitCoordinator() {
 		s.agentConn.Store(&agentConn)
 
 		// Resubscribe to all of the agents we're tracking.
-		for agentID := range s.agentNodes {
+		for agentID := range s.agentConnectionTimes {
 			err := agentConn.SubscribeAgent(agentID)
 			if err != nil {
 				s.logger.Warn(s.ctx, "resubscribe to agent", slog.Error(err), slog.F("agent_id", agentID))
@@ -212,14 +241,16 @@ type ServerTailnet struct {
 	cancel func()
 
 	logger        slog.Logger
+	tracer        trace.Tracer
 	conn          *tailnet.Conn
 	getMultiAgent func(context.Context) (tailnet.MultiAgentConn, error)
 	agentConn     atomic.Pointer[tailnet.MultiAgentConn]
 	cache         *wsconncache.Cache
 	nodesMu       sync.Mutex
-	// agentNodes is a map of agent tailnetNodes the server wants to keep a
-	// connection to. It contains the last time the agent was connected to.
-	agentNodes map[uuid.UUID]time.Time
+	// agentConnectionTimes is a map of agent tailnetNodes the server wants to
+	// keep a connection to. It contains the last time the agent was connected
+	// to.
+	agentConnectionTimes map[uuid.UUID]time.Time
 	// agentTockets holds a map of all open connections to an agent.
 	agentTickets map[uuid.UUID]map[uuid.UUID]struct{}
 
@@ -268,7 +299,7 @@ func (s *ServerTailnet) ensureAgent(agentID uuid.UUID) error {
 	s.nodesMu.Lock()
 	defer s.nodesMu.Unlock()
 
-	_, ok := s.agentNodes[agentID]
+	_, ok := s.agentConnectionTimes[agentID]
 	// If we don't have the node, subscribe.
 	if !ok {
 		s.logger.Debug(s.ctx, "subscribing to agent", slog.F("agent_id", agentID))
@@ -279,14 +310,27 @@ func (s *ServerTailnet) ensureAgent(agentID uuid.UUID) error {
 		s.agentTickets[agentID] = map[uuid.UUID]struct{}{}
 	}
 
-	s.agentNodes[agentID] = time.Now()
+	s.agentConnectionTimes[agentID] = time.Now()
 	return nil
 }
 
+func (s *ServerTailnet) acquireTicket(agentID uuid.UUID) (release func()) {
+	id := uuid.New()
+	s.nodesMu.Lock()
+	s.agentTickets[agentID][id] = struct{}{}
+	s.nodesMu.Unlock()
+
+	return func() {
+		s.nodesMu.Lock()
+		delete(s.agentTickets[agentID], id)
+		s.nodesMu.Unlock()
+	}
+}
+
 func (s *ServerTailnet) AgentConn(ctx context.Context, agentID uuid.UUID) (*codersdk.WorkspaceAgentConn, func(), error) {
 	var (
 		conn *codersdk.WorkspaceAgentConn
-		ret  = func() {}
+		ret  func()
 	)
 
 	if s.getAgentConn().AgentIsLegacy(agentID) {
@@ -299,12 +343,13 @@ func (s *ServerTailnet) AgentConn(ctx context.Context, agentID uuid.UUID) (*code
 		conn = cconn.WorkspaceAgentConn
 		ret = release
 	} else {
+		s.logger.Debug(s.ctx, "acquiring agent", slog.F("agent_id", agentID))
 		err := s.ensureAgent(agentID)
 		if err != nil {
 			return nil, nil, xerrors.Errorf("ensure agent: %w", err)
 		}
+		ret = s.acquireTicket(agentID)
 
-		s.logger.Debug(s.ctx, "acquiring agent", slog.F("agent_id", agentID))
 		conn = codersdk.NewWorkspaceAgentConn(s.conn, codersdk.WorkspaceAgentConnOptions{
 			AgentID:   agentID,
 			CloseFunc: func() error { return codersdk.ErrSkipClose },
@@ -317,7 +362,6 @@ func (s *ServerTailnet) AgentConn(ctx context.Context, agentID uuid.UUID) (*code
 	reachable := conn.AwaitReachable(ctx)
 	if !reachable {
 		ret()
-		conn.Close()
 		return nil, nil, xerrors.New("agent is unreachable")
 	}
 
@@ -336,13 +380,11 @@ func (s *ServerTailnet) DialAgentNetConn(ctx context.Context, agentID uuid.UUID,
 	nc, err := conn.DialContext(ctx, network, addr)
 	if err != nil {
 		release()
-		conn.Close()
 		return nil, xerrors.Errorf("dial context: %w", err)
 	}
 
 	return &netConnCloser{Conn: nc, close: func() {
 		release()
-		conn.Close()
 	}}, err
 }
 

diff --git a/coderd/tailnet_test.go b/coderd/tailnet_test.go
@@ -14,6 +14,7 @@ import (
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace"
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
@@ -232,6 +233,7 @@ func setupAgent(t *testing.T, agentAddresses []netip.Prefix) (uuid.UUID, agent.A
 		manifest.DERPMap,
 		func(context.Context) (tailnet.MultiAgentConn, error) { return coord.ServeMultiAgent(uuid.New()), nil },
 		cache,
+		trace.NewNoopTracerProvider(),
 	)
 	require.NoError(t, err)
 

diff --git a/enterprise/wsproxy/wsproxy.go b/enterprise/wsproxy/wsproxy.go
@@ -250,6 +250,7 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 			connInfo.DERPMap,
 			s.DialCoordinator,
 			wsconncache.New(s.DialWorkspaceAgent, 0),
+			s.TracerProvider,
 		)
 		if err != nil {
 			return nil, xerrors.Errorf("create server tailnet: %w", err)

diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -2,6 +2,7 @@ package tailnet
 
 import (
 	"context"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"net"
@@ -19,7 +20,6 @@ import (
 	"golang.org/x/xerrors"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
-	"tailscale.com/hostinfo"
 	"tailscale.com/ipn/ipnstate"
 	"tailscale.com/net/connstats"
 	"tailscale.com/net/dns"
@@ -67,6 +67,7 @@ func init() {
 }
 
 type Options struct {
+	ID         uuid.UUID
 	Addresses  []netip.Prefix
 	DERPMap    *tailcfg.DERPMap
 	DERPHeader *http.Header
@@ -78,6 +79,18 @@ type Options struct {
 	ListenPort     uint16
 }
 
+// NodeID creates a Tailscale NodeID from the last 8 bytes of a UUID. It ensures
+// the returned NodeID is always positive.
+func NodeID(uid uuid.UUID) tailcfg.NodeID {
+	id := int64(binary.BigEndian.Uint64(uid[8:]))
+
+	// ensure id is positive
+	y := id >> 63
+	id = (id ^ y) - y
+
+	return tailcfg.NodeID(id)
+}
+
 // NewConn constructs a new Wireguard server that will accept connections from the addresses provided.
 func NewConn(options *Options) (conn *Conn, err error) {
 	if options == nil {
@@ -126,13 +139,23 @@ func NewConn(options *Options) (conn *Conn, err error) {
 			Caps: []filter.CapMatch{},
 		}},
 	}
-	nodeID, err := cryptorand.Int63()
-	if err != nil {
-		return nil, xerrors.Errorf("generate node id: %w", err)
+
+	var nodeID tailcfg.NodeID
+
+	// If we're provided with a UUID, use it to populate our node ID.
+	if options.ID != uuid.Nil {
+		nodeID = NodeID(options.ID)
+	} else {
+		uid, err := cryptorand.Int63()
+		if err != nil {
+			return nil, xerrors.Errorf("generate node id: %w", err)
+		}
+		nodeID = tailcfg.NodeID(uid)
 	}
+
 	// This is used by functions below to identify the node via key
 	netMap.SelfNode = &tailcfg.Node{
-		ID:         tailcfg.NodeID(nodeID),
+		ID:         nodeID,
 		Key:        nodePublicKey,
 		Addresses:  options.Addresses,
 		AllowedIPs: options.Addresses,
@@ -488,7 +511,7 @@ func (c *Conn) UpdateNodes(nodes []*Node, replacePeers bool) error {
 			AllowedIPs: node.AllowedIPs,
 			Endpoints:  node.Endpoints,
 			DERP:       fmt.Sprintf("%s:%d", tailcfg.DerpMagicIP, node.PreferredDERP),
-			Hostinfo:   hostinfo.New().View(),
+			Hostinfo:   (&tailcfg.Hostinfo{}).View(),
 		}
 		if c.blockEndpoints {
 			peerNode.Endpoints = nil
@@ -512,6 +535,56 @@ func (c *Conn) UpdateNodes(nodes []*Node, replacePeers bool) error {
 	return nil
 }
 
+// PeerSelector is used to select a peer from within a Tailnet.
+type PeerSelector struct {
+	ID tailcfg.NodeID
+	IP netip.Prefix
+}
+
+func (c *Conn) RemovePeer(selector PeerSelector) (deleted bool, err error) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	if c.isClosed() {
+		return false, xerrors.New("connection closed")
+	}
+
+	deleted = false
+	for _, peer := range c.peerMap {
+		if peer.ID == selector.ID {
+			delete(c.peerMap, peer.ID)
+			deleted = true
+			break
+		}
+
+		for _, peerIP := range peer.Addresses {
+			if peerIP.Bits() == selector.IP.Bits() && peerIP.Addr().Compare(selector.IP.Addr()) == 0 {
+				delete(c.peerMap, peer.ID)
+				deleted = true
+				break
+			}
+		}
+	}
+	if !deleted {
+		return false, nil
+	}
+
+	c.netMap.Peers = make([]*tailcfg.Node, 0, len(c.peerMap))
+	for _, peer := range c.peerMap {
+		c.netMap.Peers = append(c.netMap.Peers, peer.Clone())
+	}
+
+	netMapCopy := *c.netMap
+	c.logger.Debug(context.Background(), "updating network map")
+	c.wireguardEngine.SetNetworkMap(&netMapCopy)
+	err = c.reconfig()
+	if err != nil {
+		return false, xerrors.Errorf("reconfig: %w", err)
+	}
+
+	return true, nil
+}
+
 func (c *Conn) reconfig() error {
 	cfg, err := nmcfg.WGCfg(c.netMap, Logger(c.logger.Named("net.wgconfig")), netmap.AllowSingleHosts, "")
 	if err != nil {