Add the ability to block endpoints

kylecarbs · kylecarbs · commit de5b13b38079 · 2022-10-13T21:35:41.000Z
diff --git a/enterprise/coderd/coderdenttest/coderdenttest.go b/enterprise/coderd/coderdenttest/coderdenttest.go
@@ -68,7 +68,7 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
 		BrowserOnly:                options.BrowserOnly,
 		SCIMAPIKey:                 options.SCIMAPIKey,
 		DERPServerRelayAddress:     oop.AccessURL.String(),
-		DERPServerRegionID:         1,
+		DERPServerRegionID:         oop.DERPMap.RegionIDs()[0],
 		ReplicaID:                  uuid.New(),
 		UserWorkspaceQuota:         options.UserWorkspaceQuota,
 		Options:                    oop,
diff --git a/enterprise/coderd/replicas_test.go b/enterprise/coderd/replicas_test.go
@@ -3,7 +3,6 @@ package coderd_test
 import (
 	"context"
 	"testing"
-	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -13,38 +12,63 @@ import (
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/coderd/database/dbtestutil"
 	"github.com/coder/coder/enterprise/coderd/coderdenttest"
+	"github.com/coder/coder/testutil"
 )
 
 func TestReplicas(t *testing.T) {
 	t.Parallel()
-	db, pubsub := dbtestutil.NewDB(t)
-	firstClient := coderdenttest.New(t, &coderdenttest.Options{
-		Options: &coderdtest.Options{
-			IncludeProvisionerDaemon: true,
-			Database:                 db,
-			Pubsub:                   pubsub,
-		},
+	t.Run("WarningsWithoutLicense", func(t *testing.T) {
+		t.Parallel()
+		db, pubsub := dbtestutil.NewDB(t)
+		firstClient := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				IncludeProvisionerDaemon: true,
+				Database:                 db,
+				Pubsub:                   pubsub,
+			},
+		})
+		_ = coderdtest.CreateFirstUser(t, firstClient)
+		secondClient := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+		})
+		secondClient.SessionToken = firstClient.SessionToken
+		ents, err := secondClient.Entitlements(context.Background())
+		require.NoError(t, err)
+		require.Len(t, ents.Warnings, 1)
 	})
-	firstUser := coderdtest.CreateFirstUser(t, firstClient)
-	coderdenttest.AddLicense(t, firstClient, coderdenttest.LicenseOptions{
-		HighAvailability: true,
-	})
-
-	secondClient := coderdenttest.New(t, &coderdenttest.Options{
-		Options: &coderdtest.Options{
-			Database: db,
-			Pubsub:   pubsub,
-		},
-	})
-	secondClient.SessionToken = firstClient.SessionToken
+	t.Run("ConnectAcrossMultiple", func(t *testing.T) {
+		t.Parallel()
+		db, pubsub := dbtestutil.NewDB(t)
+		firstClient := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				IncludeProvisionerDaemon: true,
+				Database:                 db,
+				Pubsub:                   pubsub,
+			},
+		})
+		firstUser := coderdtest.CreateFirstUser(t, firstClient)
+		coderdenttest.AddLicense(t, firstClient, coderdenttest.LicenseOptions{
+			HighAvailability: true,
+		})
 
-	agentID := setupWorkspaceAgent(t, firstClient, firstUser)
-	conn, err := secondClient.DialWorkspaceAgentTailnet(context.Background(), slogtest.Make(t, nil).Leveled(slog.LevelDebug), agentID)
-	require.NoError(t, err)
-	require.Eventually(t, func() bool {
-		_, err = conn.Ping()
-		return err == nil
-	}, 10*time.Second, 250*time.Millisecond)
+		secondClient := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+		})
+		secondClient.SessionToken = firstClient.SessionToken
 
-	_ = conn.Close()
+		agentID := setupWorkspaceAgent(t, firstClient, firstUser)
+		conn, err := secondClient.DialWorkspaceAgentTailnet(context.Background(), slogtest.Make(t, nil).Leveled(slog.LevelDebug), agentID)
+		require.NoError(t, err)
+		require.Eventually(t, func() bool {
+			_, err = conn.Ping()
+			return err == nil
+		}, testutil.WaitShort, testutil.IntervalFast)
+		_ = conn.Close()
+	})
 }
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -48,7 +48,10 @@ type Options struct {
 	Addresses []netip.Prefix
 	DERPMap   *tailcfg.DERPMap
 
-	Logger slog.Logger
+	// BlockEndpoints specifies whether P2P endpoints are blocked.
+	// If so, only DERPs can establish connections.
+	BlockEndpoints bool
+	Logger         slog.Logger
 }
 
 // NewConn constructs a new Wireguard server that will accept connections from the addresses provided.
@@ -175,6 +178,7 @@ func NewConn(options *Options) (*Conn, error) {
 	wireguardEngine.SetFilter(filter.New(netMap.PacketFilter, localIPs, logIPs, nil, Logger(options.Logger.Named("packet-filter"))))
 	dialContext, dialCancel := context.WithCancel(context.Background())
 	server := &Conn{
+		blockEndpoints:   options.BlockEndpoints,
 		dialContext:      dialContext,
 		dialCancel:       dialCancel,
 		closed:           make(chan struct{}),
@@ -240,11 +244,12 @@ func IP() netip.Addr {
 
 // Conn is an actively listening Wireguard connection.
 type Conn struct {
-	dialContext context.Context
-	dialCancel  context.CancelFunc
-	mutex       sync.Mutex
-	closed      chan struct{}
-	logger      slog.Logger
+	dialContext    context.Context
+	dialCancel     context.CancelFunc
+	mutex          sync.Mutex
+	closed         chan struct{}
+	logger         slog.Logger
+	blockEndpoints bool
 
 	dialer             *tsdial.Dialer
 	tunDevice          *tstun.Wrapper
@@ -429,6 +434,9 @@ func (c *Conn) sendNode() {
 		PreferredDERP: c.lastPreferredDERP,
 		DERPLatency:   c.lastDERPLatency,
 	}
+	if c.blockEndpoints {
+		node.Endpoints = nil
+	}
 	nodeCallback := c.nodeCallback
 	if nodeCallback == nil {
 		return
