feat: allow DERP headers to be set

coder · kylecarbs · Mar 21, 2023 · Mar 12, 2023 · Mar 13, 2023 · Mar 14, 2023
commit bbcc4a96287ebb87aaa47936a8f399751286cd20
diff --git a/cli/root.go b/cli/root.go
@@ -46,17 +46,19 @@ var (
 )
 
 const (
-	varURL              = "url"
-	varToken            = "token"
-	varAgentToken       = "agent-token"
-	varAgentURL         = "agent-url"
-	varHeader           = "header"
-	varNoOpen           = "no-open"
-	varNoVersionCheck   = "no-version-warning"
-	varNoFeatureWarning = "no-feature-warning"
-	varForceTty         = "force-tty"
-	varVerbose          = "verbose"
-	notLoggedInMessage  = "You are not logged in. Try logging in using 'coder login <url>'."
+	varURL                  = "url"
+	varToken                = "token"
+	varAgentToken           = "agent-token"
+	varAgentURL             = "agent-url"
+	varHeader               = "header"
+	varDerpHeader           = "derp-header"
+	varDerpHeaderUseDefault = "derp-header-use-default"
+	varNoOpen               = "no-open"
+	varNoVersionCheck       = "no-version-warning"
+	varNoFeatureWarning     = "no-feature-warning"
+	varForceTty             = "force-tty"
+	varVerbose              = "verbose"
+	notLoggedInMessage      = "You are not logged in. Try logging in using 'coder login <url>'."
 
 	envNoVersionCheck   = "CODER_NO_VERSION_WARNING"
 	envNoFeatureWarning = "CODER_NO_FEATURE_WARNING"
@@ -170,6 +172,8 @@ func Root(subcommands []*cobra.Command) *cobra.Command {
 	_ = cmd.PersistentFlags().MarkHidden(varAgentURL)
 	cliflag.String(cmd.PersistentFlags(), config.FlagName, "", "CODER_CONFIG_DIR", config.DefaultDir(), "Path to the global `coder` config directory.")
 	cliflag.StringArray(cmd.PersistentFlags(), varHeader, "", "CODER_HEADER", []string{}, "HTTP headers added to all requests. Provide as \"Key=Value\"")
+	cliflag.StringArray(cmd.PersistentFlags(), varDerpHeader, "", "CODER_DERP_HEADER", []string{}, "HTTP headers added to all DERP requests. Provide as \"Key=Value\"")
+	cliflag.Bool(cmd.PersistentFlags(), varDerpHeaderUseDefault, "", "CODER_DERP_HEADER_USE_DEFAULT", false, "Use default HTTP headers for all DERP requests.")
 	cmd.PersistentFlags().Bool(varForceTty, false, "Force the `coder` command to run as if connected to a TTY.")
 	_ = cmd.PersistentFlags().MarkHidden(varForceTty)
 	cmd.PersistentFlags().Bool(varNoOpen, false, "Block automatically opening URLs in the browser.")
@@ -332,19 +336,61 @@ func createUnauthenticatedClient(cmd *cobra.Command, serverURL *url.URL) (*coder
 	if err != nil {
 		return nil, err
 	}
+
+	headerMap, err := parseHeaderString(headers)
+	if err != nil {
+		return nil, err
+	}
+
 	transport := &headerTransport{
 		transport: http.DefaultTransport,
-		headers:   map[string]string{},
+		headers:   headerMap,
 	}
+
+	client.HTTPClient.Transport = transport
+
+	derpHeaders, err := cmd.Flags().GetStringArray(varDerpHeader)
+	if err != nil {
+		return nil, err
+	}
+
+	derpHeaderMap, err := parseHeaderString(derpHeaders)
+	if err != nil {
+		return nil, err
+	}
+
+	client.DERPHeader = &http.Header{}
+
+	derpHeadersUseDefault, err := cmd.Flags().GetBool(varDerpHeaderUseDefault)
+	if err != nil {
+		return nil, err
+	}
+
+	if derpHeadersUseDefault {
+		for header, value := range headerMap {
+			client.DERPHeader.Set(header, value)
+		}
+	}
+
+	for header, value := range derpHeaderMap {
+		client.DERPHeader.Set(header, value)
+	}
+
+	return client, nil
+}
+
+func parseHeaderString(headers []string) (map[string]string, error) {
+	headerMap := map[string]string{}
+
 	for _, header := range headers {
 		parts := strings.SplitN(header, "=", 2)
 		if len(parts) < 2 {
 			return nil, xerrors.Errorf("split header %q had less than two parts", header)
 		}
-		transport.headers[parts[0]] = parts[1]
+		headerMap[parts[0]] = parts[1]
 	}
-	client.HTTPClient.Transport = transport
-	return client, nil
+
+	return headerMap, nil
 }
 
 // createAgentClient returns a new client from the command context.

diff --git a/codersdk/client.go b/codersdk/client.go
@@ -79,6 +79,8 @@ type Client struct {
 	HTTPClient *http.Client
 	URL        *url.URL
 
+	DERPHeader *http.Header
+
 	// Logger is optionally provided to log requests.
 	// Method, URL, and response code will be logged by default.
 	Logger slog.Logger

diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -146,6 +146,7 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, opti
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:      []netip.Prefix{netip.PrefixFrom(ip, 128)},
 		DERPMap:        connInfo.DERPMap,
+		DERPHeader:     c.DERPHeader,
 		Logger:         options.Logger,
 		BlockEndpoints: options.BlockEndpoints,
 	})

diff --git a/go.mod b/go.mod
@@ -40,7 +40,7 @@ replace github.com/tcnksm/go-httpstat => github.com/kylecarbs/go-httpstat v0.0.0
 
 // There are a few minor changes we make to Tailscale that we're slowly upstreaming. Compare here:
 // https://github.com/tailscale/tailscale/compare/main...coder:tailscale:main
-replace tailscale.com => github.com/coder/tailscale v1.1.1-0.20230307022319-1e5e724a3949
+replace tailscale.com => github.com/JoshVee/tailscale v0.0.0-20230312075737-979756badabf
 
 // Switch to our fork that imports fixes from http://github.com/tailscale/ssh.
 // See: https://github.com/coder/coder/issues/3371

diff --git a/go.sum b/go.sum
@@ -114,6 +114,8 @@ github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym
 github.com/ClickHouse/clickhouse-go v1.4.3/go.mod h1:EaI/sW7Azgz9UATzd5ZdZHRUhHgv5+JMS9NSr2smCJI=
 github.com/DATA-DOG/go-sqlmock v1.5.0 h1:Shsta01QNfFxHCfpW6YH2STWB0MudeXXEWMr20OEh60=
 github.com/Djarvur/go-err113 v0.0.0-20210108212216-aea10b59be24/go.mod h1:4UJr5HIiMZrwgkSPdsjy2uOQExX/WEILpIrO9UPGuXs=
+github.com/JoshVee/tailscale v0.0.0-20230312075737-979756badabf h1:4r5bosAHmI1IlivX9aHZmWlrgz206xmm+iE3zBmA6l4=
+github.com/JoshVee/tailscale v0.0.0-20230312075737-979756badabf/go.mod h1:jpg+77g19FpXL43U1VoIqoSg1K/Vh5CVxycGldQ8KhA=
 github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc=
 github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE=
 github.com/Masterminds/goutils v1.1.0/go.mod h1:8cTjp+g8YejhMuvIA5y2vz3BpJxksy863GQaJW2MFNU=
@@ -376,8 +378,6 @@ github.com/coder/retry v1.3.1-0.20230210155434-e90a2e1e091d h1:09JG37IgTB6n3ouX9
 github.com/coder/retry v1.3.1-0.20230210155434-e90a2e1e091d/go.mod h1:r+1J5i/989wt6CUeNSuvFKKA9hHuKKPMxdzDbTuvwwk=
 github.com/coder/ssh v0.0.0-20220811105153-fcea99919338 h1:tN5GKFT68YLVzJoA8AHuiMNJ0qlhoD3pGN3JY9gxSko=
 github.com/coder/ssh v0.0.0-20220811105153-fcea99919338/go.mod h1:ZSS+CUoKHDrqVakTfTWUlKSr9MtMFkC4UvtQKD7O914=
-github.com/coder/tailscale v1.1.1-0.20230307022319-1e5e724a3949 h1:8WfMfRTDaEpnmhCJWfFQ7JHz19GyP+EgFgLGu5ngdek=
-github.com/coder/tailscale v1.1.1-0.20230307022319-1e5e724a3949/go.mod h1:jpg+77g19FpXL43U1VoIqoSg1K/Vh5CVxycGldQ8KhA=
 github.com/coder/terraform-provider-coder v0.6.15 h1:Llvh4RwxSQ/goy7ToTOeHf3tdEz+79qbyOh61hNnJs0=
 github.com/coder/terraform-provider-coder v0.6.15/go.mod h1:UIfU3bYNeSzJJvHyJ30tEKjD6Z9utloI+HUM/7n94CY=
 github.com/containerd/aufs v0.0.0-20200908144142-dab0cbea06f4/go.mod h1:nukgQABAEopAHvB6j7cnP5zJ+/3aVcE7hCYqvIwAHyE=

diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"net/http"
 	"net/netip"
 	"reflect"
 	"strconv"
@@ -50,8 +51,9 @@ func init() {
 }
 
 type Options struct {
-	Addresses []netip.Prefix
-	DERPMap   *tailcfg.DERPMap
+	Addresses  []netip.Prefix
+	DERPMap    *tailcfg.DERPMap
+	DERPHeader *http.Header
 
 	// BlockEndpoints specifies whether P2P endpoints are blocked.
 	// If so, only DERPs can establish connections.
@@ -160,6 +162,10 @@ func NewConn(options *Options) (conn *Conn, err error) {
 		return nil, xerrors.New("get wireguard internals")
 	}
 
+	if options.DERPHeader != nil {
+		magicConn.SetDERPHeader(options.DERPHeader.Clone())
+	}
+
 	// Update the keys for the magic connection!
 	err = magicConn.SetPrivateKey(nodePrivateKey)
 	if err != nil {