feat: allow DERP headers to be set (#6572)

JoshVee · kylecarbs · web-flow · commit 97f77c450710 · 2023-03-21T18:43:20.000Z
* feat: allow DERP headers to be set

* chore: remove custom flag

* Clone DERP header on client create

* Adjust to use interface to cast headers

---------

Co-authored-by: Kyle Carberry &lt;kyle@carberry.com&gt;
diff --git a/cli/root.go b/cli/root.go
@@ -334,14 +334,14 @@ func createUnauthenticatedClient(cmd *cobra.Command, serverURL *url.URL) (*coder
 	}
 	transport := &headerTransport{
 		transport: http.DefaultTransport,
-		headers:   map[string]string{},
+		header:    http.Header{},
 	}
 	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]
+		transport.header.Add(parts[0], parts[1])
 	}
 	client.HTTPClient.Transport = transport
 	return client, nil
@@ -655,12 +655,18 @@ func checkWarnings(cmd *cobra.Command, client *codersdk.Client) error {
 
 type headerTransport struct {
 	transport http.RoundTripper
-	headers   map[string]string
+	header    http.Header
+}
+
+func (h *headerTransport) Header() http.Header {
+	return h.header.Clone()
 }
 
 func (h *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	for k, v := range h.headers {
-		req.Header.Add(k, v)
+	for k, v := range h.header {
+		for _, vv := range v {
+			req.Header.Add(k, vv)
+		}
 	}
 	return h.transport.RoundTrip(req)
 }
diff --git a/cli/scaletest.go b/cli/scaletest.go
@@ -330,8 +330,8 @@ func scaletestCleanup() *cobra.Command {
 			client.HTTPClient = &http.Client{
 				Transport: &headerTransport{
 					transport: http.DefaultTransport,
-					headers: map[string]string{
-						codersdk.BypassRatelimitHeader: "true",
+					header: map[string][]string{
+						codersdk.BypassRatelimitHeader: {"true"},
 					},
 				},
 			}
@@ -515,8 +515,8 @@ It is recommended that all rate limits are disabled on the server before running
 			client.HTTPClient = &http.Client{
 				Transport: &headerTransport{
 					transport: http.DefaultTransport,
-					headers: map[string]string{
-						codersdk.BypassRatelimitHeader: "true",
+					header: map[string][]string{
+						codersdk.BypassRatelimitHeader: {"true"},
 					},
 				},
 			}
diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -143,9 +143,17 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, opti
 	}
 
 	ip := tailnet.IP()
+	var header http.Header
+	headerTransport, ok := c.HTTPClient.Transport.(interface {
+		Header() http.Header
+	})
+	if ok {
+		header = headerTransport.Header()
+	}
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:      []netip.Prefix{netip.PrefixFrom(ip, 128)},
 		DERPMap:        connInfo.DERPMap,
+		DERPHeader:     &header,
 		Logger:         options.Logger,
 		BlockEndpoints: options.BlockEndpoints,
 	})
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.
@@ -159,6 +161,9 @@ func NewConn(options *Options) (conn *Conn, err error) {
 	if !ok {
 		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)

Original file line number	Diff line number	Diff line change
`@@ -334,14 +334,14 @@ func createUnauthenticatedClient(cmd cobra.Command, serverURL url.URL) (*coder`
`334`	`334`	`}`
`335`	`335`	`transport := &headerTransport{`
`336`	`336`	`transport: http.DefaultTransport,`
`337`		`- headers: map[string]string{},`
	`337`	`+ header: http.Header{},`
`338`	`338`	`}`
`339`	`339`	`for _, header := range headers {`
`340`	`340`	`parts := strings.SplitN(header, "=", 2)`
`341`	`341`	`if len(parts) < 2 {`
`342`	`342`	`return nil, xerrors.Errorf("split header %q had less than two parts", header)`
`343`	`343`	`}`
`344`		`- transport.headers[parts[0]] = parts[1]`
	`344`	`+ transport.header.Add(parts[0], parts[1])`
`345`	`345`	`}`
`346`	`346`	`client.HTTPClient.Transport = transport`
`347`	`347`	`return client, nil`
`@@ -655,12 +655,18 @@ func checkWarnings(cmd cobra.Command, client codersdk.Client) error {`
`655`	`655`
`656`	`656`	`type headerTransport struct {`
`657`	`657`	`transport http.RoundTripper`
`658`		`- headers map[string]string`
	`658`	`+ header http.Header`
	`659`	`+}`
	`660`	`+`
	`661`	`+func (h *headerTransport) Header() http.Header {`
	`662`	`+ return h.header.Clone()`
`659`	`663`	`}`
`660`	`664`
`661`	`665`	`func (h headerTransport) RoundTrip(req http.Request) (*http.Response, error) {`
`662`		`- for k, v := range h.headers {`
`663`		`- req.Header.Add(k, v)`
	`666`	`+ for k, v := range h.header {`
	`667`	`+ for _, vv := range v {`
	`668`	`+ req.Header.Add(k, vv)`
	`669`	`+ }`
`664`	`670`	`}`
`665`	`671`	`return h.transport.RoundTrip(req)`
`666`	`672`	`}`