feat: check if peer exists before proxying (#14)

Author: coadler

.golangci.yaml

@@ -211,7 +211,6 @@ linters:
     - asciicheck
     - bidichk
     - bodyclose
-    - deadcode
     - dogsled
     - errcheck
     - errname
@@ -255,4 +254,3 @@ linters:
     - typecheck
     - unconvert
     - unused
-    - varcheck

cmd/tunneld/main.go

@@ -123,9 +123,9 @@ func main() {
 				EnvVars: []string{"TUNNELD_TRACING_HONEYCOMB_TEAM"},
 			},
 			&cli.StringFlag{
-				Name:    "tracing-service-id",
-				Usage:   "The service ID to annotate all traces with that uniquely identifies this deployment.",
-				EnvVars: []string{"TUNNELD_TRACING_SERVICE_ID"},
+				Name:    "tracing-instance-id",
+				Usage:   "The instance ID to annotate all traces with that uniquely identifies this deployment.",
+				EnvVars: []string{"TUNNELD_TRACING_INSTANCE_ID"},
 			},
 		},
 		Action: runApp,
@@ -152,7 +152,7 @@ func runApp(ctx *cli.Context) error {
 		realIPHeader           = ctx.String("real-ip-header")
 		pprofListenAddress     = ctx.String("pprof-listen-address")
 		tracingHoneycombTeam   = ctx.String("tracing-honeycomb-team")
-		tracingServiceID       = ctx.String("tracing-service-id")
+		tracingInstanceID      = ctx.String("tracing-instance-id")
 	)
 	if baseURL == "" {
 		return xerrors.New("base-url is required. See --help for more information.")
@@ -185,7 +185,7 @@ func runApp(ctx *cli.Context) error {
 
 		// Create a new tracer provider with a batch span processor and the otlp
 		// exporter.
-		tp := newTraceProvider(exp, tracingServiceID)
+		tp := newTraceProvider(exp, tracingInstanceID)
 		otel.SetTracerProvider(tp)
 		otel.SetTextMapPropagator(
 			propagation.NewCompositeTextMapPropagator(

cmd/tunneld/tracing.go

@@ -26,11 +26,11 @@ func newHoneycombExporter(ctx context.Context, teamID string) (*otlptrace.Export
 	return otlptrace.New(ctx, client)
 }
 
-func newTraceProvider(exp *otlptrace.Exporter, serviceID string) *sdktrace.TracerProvider {
+func newTraceProvider(exp *otlptrace.Exporter, instanceID string) *sdktrace.TracerProvider {
 	rsc := resource.NewWithAttributes(
 		semconv.SchemaURL,
 		semconv.ServiceNameKey.String("WireguardTunnel"),
-		semconv.ServiceInstanceIDKey.String(serviceID),
+		semconv.ServiceInstanceIDKey.String(instanceID),
 		semconv.ServiceVersionKey.String(buildinfo.Version()),
 	)
 

go.mod

@@ -19,7 +19,7 @@ require (
 	golang.org/x/mod v0.8.0
 	golang.org/x/sync v0.1.0
 	golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2
-	golang.zx2c4.com/wireguard v0.0.0-20230223181233-21636207a675
+	golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230215201556-9c5414ab4bde
 	google.golang.org/grpc v1.53.0
 )
@@ -52,7 +52,7 @@ require (
 	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
 	golang.org/x/crypto v0.6.0 // indirect
 	golang.org/x/net v0.7.0 // indirect
-	golang.org/x/sys v0.5.0 // indirect
+	golang.org/x/sys v0.5.1-0.20230222185716-a3b23cc77e89 // indirect
 	golang.org/x/term v0.5.0 // indirect
 	golang.org/x/text v0.7.0 // indirect
 	golang.org/x/time v0.3.0 // indirect

go.sum

@@ -432,8 +432,8 @@ golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210510120138-977fb7262007/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210514084401-e8d321eab015/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.5.0 h1:MUK/U/4lj1t1oPg0HfuXDN/Z1wv31ZJ/YcPiGccS4DU=
-golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.5.1-0.20230222185716-a3b23cc77e89 h1:260HNjMTPDya+jq5AM1zZLgG9pv9GASPAGiEEJUbRg4=
+golang.org/x/sys v0.5.1-0.20230222185716-a3b23cc77e89/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.5.0 h1:n2a8QNdAb0sZNpU9R1ALUXBbY+w51fCQDN+7EdxNBsY=
 golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
@@ -508,8 +508,8 @@ golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 h1:H2TDz8ibqkAF6YGhCdN3j
 golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2/go.mod h1:K8+ghG5WaK9qNqU5K3HdILfMLy1f3aNYFI/wnl100a8=
 golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 h1:B82qJJgjvYKsXS9jeunTOisW56dUokqW/FOteYJJ/yg=
 golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2/go.mod h1:deeaetjYA+DHMHg+sMSMI58GrEteJUUzzw7en6TJQcI=
-golang.zx2c4.com/wireguard v0.0.0-20230223181233-21636207a675 h1:/J/RVnr7ng4fWPRH3xa4WtBJ1Jp+Auu4YNLmGiPv5QU=
-golang.zx2c4.com/wireguard v0.0.0-20230223181233-21636207a675/go.mod h1:whfbyDBt09xhCYQWtO2+3UVjlaq6/9hDZrjg2ZE6SyA=
+golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b h1:J1CaxgLerRR5lgx3wnr6L04cJFbWoceSK9JWBdglINo=
+golang.zx2c4.com/wireguard v0.0.0-20230325221338-052af4a8072b/go.mod h1:tqur9LnfstdR9ep2LaJT4lFUl0EjlHtge+gAjmsHUG4=
 golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230215201556-9c5414ab4bde h1:ybF7AMzIUikL9x4LgwEmzhXtzRpKNqngme1VGDWz+Nk=
 golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230215201556-9c5414ab4bde/go.mod h1:mQqgjkW8GQQcJQsbBvK890TKqUK1DfKWkuBGbOkuMHQ=
 google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE=

tunneld/api.go

@@ -165,10 +165,24 @@ func (api *API) registerClient(req tunnelsdk.ClientRegisterRequest) (tunnelsdk.C
 
 	ip, urls := api.WireguardPublicKeyToIPAndURLs(req.PublicKey, req.Version)
 
+	api.pkeyCacheMu.Lock()
+	api.pkeyCache[ip] = cachedPeer{
+		key:           req.PublicKey,
+		lastHandshake: time.Now(),
+	}
+	api.pkeyCacheMu.Unlock()
+
 	exists := true
 	if api.wgDevice.LookupPeer(req.PublicKey) == nil {
 		exists = false
 
+		api.pkeyCacheMu.Lock()
+		api.pkeyCache[ip] = cachedPeer{
+			key:           req.PublicKey,
+			lastHandshake: time.Now(),
+		}
+		api.pkeyCacheMu.Unlock()
+
 		err := api.wgDevice.IpcSet(fmt.Sprintf(`public_key=%x
 allowed_ip=%s/128`,
 			req.PublicKey,
@@ -186,6 +200,7 @@ allowed_ip=%s/128`,
 
 	return tunnelsdk.ClientRegisterResponse{
 		Version:         req.Version,
+		ReregisterWait:  api.PeerRegisterInterval,
 		TunnelURLs:      urlsStr,
 		ClientIP:        ip,
 		ServerEndpoint:  api.WireguardEndpoint,
@@ -205,6 +220,12 @@ func (api *API) handleTunnel(rw http.ResponseWriter, r *http.Request) {
 	subdomainParts := strings.Split(subdomain, "-")
 	user := subdomainParts[len(subdomainParts)-1]
 
+	span := trace.SpanFromContext(ctx)
+	span.SetAttributes(
+		attribute.Bool("proxy_request", true),
+		attribute.String("user", user),
+	)
+
 	ip, err := api.HostnameToWireguardIP(user)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, tunnelsdk.Response{
@@ -214,11 +235,17 @@ func (api *API) handleTunnel(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	span := trace.SpanFromContext(ctx)
-	span.SetAttributes(
-		attribute.Bool("proxy_request", true),
-		attribute.String("user", user),
-	)
+	api.pkeyCacheMu.RLock()
+	pkey, ok := api.pkeyCache[ip]
+	api.pkeyCacheMu.RUnlock()
+
+	if !ok || time.Since(pkey.lastHandshake) > api.PeerTimeout {
+		httpapi.Write(ctx, rw, http.StatusBadGateway, tunnelsdk.Response{
+			Message: "Peer is not connected.",
+			Detail:  "",
+		})
+		return
+	}
 
 	// The transport on the reverse proxy uses this ctx value to know which
 	// IP to dial. See tunneld.go.

tunneld/api_test.go

@@ -100,6 +100,7 @@ func Test_postClients(t *testing.T) {
 	require.Equal(t, td.WireguardServerIP, res.ServerIP)
 	require.Equal(t, td.WireguardKey.NoisePublicKey(), res.ServerPublicKey)
 	require.Equal(t, td.WireguardMTU, res.WireguardMTU)
+	require.Equal(t, td.PeerRegisterInterval, res.ReregisterWait)
 
 	// Register the same client again.
 	res2, err := client.ClientRegister(context.Background(), tunnelsdk.ClientRegisterRequest{

tunneld/options.go

@@ -19,8 +19,10 @@ import (
 )
 
 const (
-	DefaultWireguardMTU    = 1280
-	DefaultPeerDialTimeout = 10 * time.Second
+	DefaultWireguardMTU     = 1280
+	DefaultPeerDialTimeout  = 10 * time.Second
+	DefaultPeerPollDuration = 30 * time.Second
+	DefaultPeerTimeout      = 2 * time.Minute
 )
 
 var (
@@ -70,6 +72,12 @@ type Options struct {
 	// PeerDialTimeout is the timeout for dialing a peer on a request. Defaults
 	// to 10 seconds.
 	PeerDialTimeout time.Duration
+
+	// PeerRegisterInterval is how often the clients should re-register.
+	PeerRegisterInterval time.Duration
+
+	// PeerTimeout is how long the server will wait before removing the peer.
+	PeerTimeout time.Duration
 }
 
 // Validate checks that the options are valid and populates default values for
@@ -127,6 +135,18 @@ func (options *Options) Validate() error {
 	if options.PeerDialTimeout <= 0 {
 		options.PeerDialTimeout = DefaultPeerDialTimeout
 	}
+	if options.PeerRegisterInterval <= 0 {
+		options.PeerRegisterInterval = DefaultPeerPollDuration
+	}
+	if options.PeerTimeout <= 0 {
+		options.PeerTimeout = DefaultPeerTimeout
+	}
+	if options.PeerRegisterInterval >= options.PeerTimeout {
+		return xerrors.Errorf("PeerRegisterInterval(%s) must be less than PeerTimeout(%s)",
+			options.PeerRegisterInterval.String(),
+			options.PeerTimeout.String(),
+		)
+	}
 
 	return nil
 }

tunneld/options_test.go

@@ -41,6 +41,8 @@ func Test_Option(t *testing.T) {
 				WireguardNetworkPrefix: netip.MustParsePrefix("feed::1/64"),
 				RealIPHeader:           "X-Real-Ip",
 				PeerDialTimeout:        1 * time.Second,
+				PeerRegisterInterval:   time.Second,
+				PeerTimeout:            2 * time.Second,
 			}
 
 			clone := o

tunneld/tunneld.go

@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 	"net/netip"
+	"sync"
 	"time"
 
 	"go.opentelemetry.io/otel"
@@ -24,6 +25,14 @@ type API struct {
 	wgNet     *netstack.Net
 	wgDevice  *device.Device
 	transport *http.Transport
+
+	pkeyCacheMu sync.RWMutex
+	pkeyCache   map[netip.Addr]cachedPeer
+}
+
+type cachedPeer struct {
+	key           device.NoisePublicKey
+	lastHandshake time.Time
 }
 
 func New(options *Options) (*API, error) {
@@ -72,9 +81,10 @@ listen_port=%d`,
 	}
 
 	return &API{
-		Options:  options,
-		wgNet:    wgNet,
-		wgDevice: dev,
+		Options:   options,
+		wgNet:     wgNet,
+		wgDevice:  dev,
+		pkeyCache: make(map[netip.Addr]cachedPeer),
 		transport: &http.Transport{
 			DialContext: func(ctx context.Context, network, addr string) (nc net.Conn, err error) {
 				ctx, span := otel.GetTracerProvider().Tracer("").Start(ctx, "(http.Transport).DialContext")

tunneld/tunneld_test.go

@@ -2,6 +2,7 @@ package tunneld_test
 
 import (
 	"context"
+	"encoding/json"
 	"io"
 	"log"
 	"net"
@@ -337,6 +338,83 @@ func TestTimeout(t *testing.T) {
 	require.Equal(t, http.StatusBadGateway, res.StatusCode)
 }
 
+func TestPeerTimeout(t *testing.T) {
+	t.Parallel()
+
+	td, client := createTestTunneld(t, &tunneld.Options{
+		PeerTimeout:          time.Second,
+		PeerRegisterInterval: 100 * time.Millisecond,
+	})
+	require.NotNil(t, td)
+
+	// Start a tunnel.
+	key, err := tunnelsdk.GeneratePrivateKey()
+	require.NoError(t, err, "generate private key")
+	tunnel, err := client.LaunchTunnel(context.Background(), tunnelsdk.TunnelConfig{
+		Log: slogtest.
+			Make(t, &slogtest.Options{IgnoreErrors: true}).
+			Named("tunnel_client"),
+		PrivateKey: key,
+	})
+	require.NoError(t, err, "launch tunnel")
+	defer func() {
+		_ = tunnel.Close()
+		<-tunnel.Wait()
+	}()
+
+	require.NotNil(t, tunnel.URL)
+	require.Len(t, tunnel.OtherURLs, 1)
+	require.NotEqual(t, tunnel.URL.String(), tunnel.OtherURLs[0].String())
+
+	serveTunnel(t, tunnel)
+	waitForTunnelReady(t, client, tunnel)
+
+	// Successfully send a request to the peer.
+	{
+		u, err := tunnel.URL.Parse("/test/1")
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		res, err := client.Request(ctx, http.MethodGet, u.String(), nil)
+		if !assert.NoError(t, err) {
+			return
+		}
+		defer res.Body.Close()
+		assert.Equal(t, http.StatusOK, res.StatusCode)
+
+		body, err := io.ReadAll(res.Body)
+		require.NoError(t, err)
+		require.Equal(t, "hello world /test/1", string(body))
+	}
+
+	err = tunnel.Close()
+	require.NoError(t, err, "close tunnel")
+	<-tunnel.Wait()
+
+	time.Sleep(td.PeerTimeout)
+
+	// The correct error should be returned after the peer goes away.
+	{
+		u, err := tunnel.URL.Parse("/test/1")
+		require.NoError(t, err)
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		res, err := client.Request(ctx, http.MethodGet, u.String(), nil)
+		require.NoError(t, err)
+		defer res.Body.Close()
+
+		tres := tunnelsdk.Response{}
+		err = json.NewDecoder(res.Body).Decode(&tres)
+		require.NoError(t, err)
+
+		require.Equal(t, http.StatusBadGateway, res.StatusCode)
+		require.Equal(t, "Peer is not connected.", tres.Message)
+	}
+}
+
 func freeUDPPort(t *testing.T) uint16 {
 	t.Helper()
 

tunnelsdk/api.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"net/netip"
+	"time"
 
 	"golang.zx2c4.com/wireguard/device"
 )
@@ -20,7 +21,8 @@ type ClientRegisterRequest struct {
 }
 
 type ClientRegisterResponse struct {
-	Version TunnelVersion `json:"version"`
+	Version        TunnelVersion `json:"version"`
+	ReregisterWait time.Duration `json:"reregister_wait"`
 	// TunnelURLs contains a list of valid URLs that will be forwarded from the
 	// server to this tunnel client once connected. The first URL is the
 	// preferred URL, and the other URLs are provided for compatibility