mirror of
https://github.com/coder/coder.git
synced 2025-07-06 15:41:45 +00:00
This adds the ability for `TunnelAuth` to also authorize incoming wireguard node IPs, preventing agents from reporting anything other than their static IP generated from the agent ID.
351 lines
8.4 KiB
Go
351 lines
8.4 KiB
Go
package tailnettest
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"fmt"
|
|
"html"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/exp/slices"
|
|
"tailscale.com/derp"
|
|
"tailscale.com/derp/derphttp"
|
|
"tailscale.com/net/stun/stuntest"
|
|
"tailscale.com/tailcfg"
|
|
"tailscale.com/types/key"
|
|
tslogger "tailscale.com/types/logger"
|
|
"tailscale.com/types/nettype"
|
|
|
|
"cdr.dev/slog/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/tailnet"
|
|
"github.com/coder/coder/v2/tailnet/proto"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
//go:generate mockgen -destination ./multiagentmock.go -package tailnettest github.com/coder/coder/v2/tailnet MultiAgentConn
|
|
//go:generate mockgen -destination ./coordinatormock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinator
|
|
//go:generate mockgen -destination ./coordinateemock.go -package tailnettest github.com/coder/coder/v2/tailnet Coordinatee
|
|
|
|
type derpAndSTUNCfg struct {
|
|
DisableSTUN bool
|
|
DERPIsEmbedded bool
|
|
}
|
|
|
|
type DERPAndStunOption func(cfg *derpAndSTUNCfg)
|
|
|
|
func DisableSTUN(cfg *derpAndSTUNCfg) {
|
|
cfg.DisableSTUN = true
|
|
}
|
|
|
|
func DERPIsEmbedded(cfg *derpAndSTUNCfg) {
|
|
cfg.DERPIsEmbedded = true
|
|
}
|
|
|
|
// RunDERPAndSTUN creates a DERP mapping for tests.
|
|
func RunDERPAndSTUN(t *testing.T, opts ...DERPAndStunOption) (*tailcfg.DERPMap, *derp.Server) {
|
|
cfg := new(derpAndSTUNCfg)
|
|
for _, o := range opts {
|
|
o(cfg)
|
|
}
|
|
logf := tailnet.Logger(slogtest.Make(t, nil))
|
|
d := derp.NewServer(key.NewNode(), logf)
|
|
server := httptest.NewUnstartedServer(derphttp.Handler(d))
|
|
server.Config.ErrorLog = tslogger.StdLogger(logf)
|
|
server.Config.TLSNextProto = make(map[string]func(*http.Server, *tls.Conn, http.Handler))
|
|
server.StartTLS()
|
|
t.Cleanup(func() {
|
|
server.CloseClientConnections()
|
|
server.Close()
|
|
d.Close()
|
|
})
|
|
|
|
stunPort := -1
|
|
if !cfg.DisableSTUN {
|
|
stunAddr, stunCleanup := stuntest.ServeWithPacketListener(t, nettype.Std{})
|
|
t.Cleanup(func() {
|
|
stunCleanup()
|
|
})
|
|
stunPort = stunAddr.Port
|
|
}
|
|
tcpAddr, ok := server.Listener.Addr().(*net.TCPAddr)
|
|
if !ok {
|
|
t.FailNow()
|
|
}
|
|
|
|
return &tailcfg.DERPMap{
|
|
Regions: map[int]*tailcfg.DERPRegion{
|
|
1: {
|
|
RegionID: 1,
|
|
RegionCode: "test",
|
|
RegionName: "Test",
|
|
Nodes: []*tailcfg.DERPNode{
|
|
{
|
|
Name: "t2",
|
|
RegionID: 1,
|
|
IPv4: "127.0.0.1",
|
|
IPv6: "none",
|
|
STUNPort: stunPort,
|
|
DERPPort: tcpAddr.Port,
|
|
InsecureForTests: true,
|
|
},
|
|
},
|
|
EmbeddedRelay: cfg.DERPIsEmbedded,
|
|
},
|
|
},
|
|
}, d
|
|
}
|
|
|
|
// RunDERPOnlyWebSockets creates a DERP mapping for tests that
|
|
// only allows WebSockets through it. Many proxies do not support
|
|
// upgrading DERP, so this is a good fallback.
|
|
func RunDERPOnlyWebSockets(t *testing.T) *tailcfg.DERPMap {
|
|
logf := tailnet.Logger(slogtest.Make(t, nil))
|
|
d := derp.NewServer(key.NewNode(), logf)
|
|
handler := derphttp.Handler(d)
|
|
var closeFunc func()
|
|
handler, closeFunc = tailnet.WithWebsocketSupport(d, handler)
|
|
server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/derp" {
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = w.Write([]byte("hello"))
|
|
return
|
|
}
|
|
if r.Header.Get("Upgrade") != "websocket" {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
_, _ = w.Write([]byte(fmt.Sprintf(`Invalid "Upgrade" header: %s`, html.EscapeString(r.Header.Get("Upgrade")))))
|
|
return
|
|
}
|
|
handler.ServeHTTP(w, r)
|
|
}))
|
|
server.Config.ErrorLog = tslogger.StdLogger(logf)
|
|
server.Config.TLSNextProto = make(map[string]func(*http.Server, *tls.Conn, http.Handler))
|
|
server.StartTLS()
|
|
t.Cleanup(func() {
|
|
server.CloseClientConnections()
|
|
server.Close()
|
|
closeFunc()
|
|
d.Close()
|
|
})
|
|
|
|
tcpAddr, ok := server.Listener.Addr().(*net.TCPAddr)
|
|
if !ok {
|
|
t.FailNow()
|
|
}
|
|
|
|
return &tailcfg.DERPMap{
|
|
Regions: map[int]*tailcfg.DERPRegion{
|
|
1: {
|
|
RegionID: 1,
|
|
RegionCode: "test",
|
|
RegionName: "Test",
|
|
Nodes: []*tailcfg.DERPNode{
|
|
{
|
|
Name: "t1",
|
|
RegionID: 1,
|
|
IPv4: "127.0.0.1",
|
|
IPv6: "none",
|
|
DERPPort: tcpAddr.Port,
|
|
InsecureForTests: true,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
type TestMultiAgent struct {
|
|
t testing.TB
|
|
id uuid.UUID
|
|
a tailnet.MultiAgentConn
|
|
nodeKey []byte
|
|
discoKey string
|
|
}
|
|
|
|
func NewTestMultiAgent(t testing.TB, coord tailnet.Coordinator) *TestMultiAgent {
|
|
nk, err := key.NewNode().Public().MarshalBinary()
|
|
require.NoError(t, err)
|
|
dk, err := key.NewDisco().Public().MarshalText()
|
|
require.NoError(t, err)
|
|
m := &TestMultiAgent{t: t, id: uuid.New(), nodeKey: nk, discoKey: string(dk)}
|
|
m.a = coord.ServeMultiAgent(m.id)
|
|
return m
|
|
}
|
|
|
|
func (m *TestMultiAgent) SendNodeWithDERP(d int32) {
|
|
m.t.Helper()
|
|
err := m.a.UpdateSelf(&proto.Node{
|
|
Key: m.nodeKey,
|
|
Disco: m.discoKey,
|
|
PreferredDerp: d,
|
|
})
|
|
require.NoError(m.t, err)
|
|
}
|
|
|
|
func (m *TestMultiAgent) Close() {
|
|
m.t.Helper()
|
|
err := m.a.Close()
|
|
require.NoError(m.t, err)
|
|
}
|
|
|
|
func (m *TestMultiAgent) RequireSubscribeAgent(id uuid.UUID) {
|
|
m.t.Helper()
|
|
err := m.a.SubscribeAgent(id)
|
|
require.NoError(m.t, err)
|
|
}
|
|
|
|
func (m *TestMultiAgent) RequireUnsubscribeAgent(id uuid.UUID) {
|
|
m.t.Helper()
|
|
err := m.a.UnsubscribeAgent(id)
|
|
require.NoError(m.t, err)
|
|
}
|
|
|
|
func (m *TestMultiAgent) RequireEventuallyHasDERPs(ctx context.Context, expected ...int) {
|
|
m.t.Helper()
|
|
for {
|
|
resp, ok := m.a.NextUpdate(ctx)
|
|
require.True(m.t, ok)
|
|
nodes, err := tailnet.OnlyNodeUpdates(resp)
|
|
require.NoError(m.t, err)
|
|
if len(nodes) != len(expected) {
|
|
m.t.Logf("expected %d, got %d nodes", len(expected), len(nodes))
|
|
continue
|
|
}
|
|
|
|
derps := make([]int, 0, len(nodes))
|
|
for _, n := range nodes {
|
|
derps = append(derps, n.PreferredDERP)
|
|
}
|
|
for _, e := range expected {
|
|
if !slices.Contains(derps, e) {
|
|
m.t.Logf("expected DERP %d to be in %v", e, derps)
|
|
continue
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *TestMultiAgent) RequireNeverHasDERPs(ctx context.Context, expected ...int) {
|
|
m.t.Helper()
|
|
for {
|
|
resp, ok := m.a.NextUpdate(ctx)
|
|
if !ok {
|
|
return
|
|
}
|
|
nodes, err := tailnet.OnlyNodeUpdates(resp)
|
|
require.NoError(m.t, err)
|
|
if len(nodes) != len(expected) {
|
|
m.t.Logf("expected %d, got %d nodes", len(expected), len(nodes))
|
|
continue
|
|
}
|
|
|
|
derps := make([]int, 0, len(nodes))
|
|
for _, n := range nodes {
|
|
derps = append(derps, n.PreferredDERP)
|
|
}
|
|
for _, e := range expected {
|
|
if !slices.Contains(derps, e) {
|
|
m.t.Logf("expected DERP %d to be in %v", e, derps)
|
|
continue
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *TestMultiAgent) RequireEventuallyClosed(ctx context.Context) {
|
|
m.t.Helper()
|
|
tkr := time.NewTicker(testutil.IntervalFast)
|
|
defer tkr.Stop()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
m.t.Fatal("timeout")
|
|
return // unhittable
|
|
case <-tkr.C:
|
|
if m.a.IsClosed() {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
type FakeCoordinator struct {
|
|
CoordinateCalls chan *FakeCoordinate
|
|
ServeClientCalls chan *FakeServeClient
|
|
}
|
|
|
|
func (*FakeCoordinator) ServeHTTPDebug(http.ResponseWriter, *http.Request) {
|
|
panic("unimplemented")
|
|
}
|
|
|
|
func (*FakeCoordinator) Node(uuid.UUID) *tailnet.Node {
|
|
panic("unimplemented")
|
|
}
|
|
|
|
func (f *FakeCoordinator) ServeClient(conn net.Conn, id uuid.UUID, agent uuid.UUID) error {
|
|
errCh := make(chan error)
|
|
f.ServeClientCalls <- &FakeServeClient{
|
|
Conn: conn,
|
|
ID: id,
|
|
Agent: agent,
|
|
ErrCh: errCh,
|
|
}
|
|
return <-errCh
|
|
}
|
|
|
|
func (*FakeCoordinator) ServeAgent(net.Conn, uuid.UUID, string) error {
|
|
panic("unimplemented")
|
|
}
|
|
|
|
func (*FakeCoordinator) Close() error {
|
|
panic("unimplemented")
|
|
}
|
|
|
|
func (*FakeCoordinator) ServeMultiAgent(uuid.UUID) tailnet.MultiAgentConn {
|
|
panic("unimplemented")
|
|
}
|
|
|
|
func (f *FakeCoordinator) Coordinate(ctx context.Context, id uuid.UUID, name string, a tailnet.CoordinateeAuth) (chan<- *proto.CoordinateRequest, <-chan *proto.CoordinateResponse) {
|
|
reqs := make(chan *proto.CoordinateRequest, 100)
|
|
resps := make(chan *proto.CoordinateResponse, 100)
|
|
f.CoordinateCalls <- &FakeCoordinate{
|
|
Ctx: ctx,
|
|
ID: id,
|
|
Name: name,
|
|
Auth: a,
|
|
Reqs: reqs,
|
|
Resps: resps,
|
|
}
|
|
return reqs, resps
|
|
}
|
|
|
|
func NewFakeCoordinator() *FakeCoordinator {
|
|
return &FakeCoordinator{
|
|
CoordinateCalls: make(chan *FakeCoordinate, 100),
|
|
ServeClientCalls: make(chan *FakeServeClient, 100),
|
|
}
|
|
}
|
|
|
|
type FakeCoordinate struct {
|
|
Ctx context.Context
|
|
ID uuid.UUID
|
|
Name string
|
|
Auth tailnet.CoordinateeAuth
|
|
Reqs chan *proto.CoordinateRequest
|
|
Resps chan *proto.CoordinateResponse
|
|
}
|
|
|
|
type FakeServeClient struct {
|
|
Conn net.Conn
|
|
ID uuid.UUID
|
|
Agent uuid.UUID
|
|
ErrCh chan error
|
|
}
|