Adjust to use interface to cast headers

kylecarbs · kylecarbs · commit aa99c6c9024c · 2023-03-21T17:41:40.000Z
diff --git a/cli/root.go b/cli/root.go
@@ -334,22 +334,16 @@ 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
-	client.DERPHeader = &http.Header{}
-	for header, value := range transport.headers {
-		client.DERPHeader.Set(header, value)
-	}
-
 	return client, nil
 }
 
@@ -661,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/client.go b/codersdk/client.go
@@ -79,8 +79,6 @@ 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
@@ -143,7 +143,13 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, opti
 	}
 
 	ip := tailnet.IP()
-	header := c.DERPHeader.Clone()
+	var header http.Header
+	headerTransport, ok := c.HTTPClient.Transport.(interface {
+		Header() http.Header
+	})
+	if ok {
+		header = headerTransport.Header().Clone()
+	}
 	conn, err := tailnet.NewConn(&tailnet.Options{
 		Addresses:      []netip.Prefix{netip.PrefixFrom(ip, 128)},
 		DERPMap:        connInfo.DERPMap,
diff --git a/tailnet/conn.go b/tailnet/conn.go
@@ -161,7 +161,6 @@ 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())
 	}

Original file line number	Diff line number	Diff line change
`@@ -334,22 +334,16 @@ 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`		`-`
`347`	`346`	`client.HTTPClient.Transport = transport`
`348`		`- client.DERPHeader = &http.Header{}`
`349`		`- for header, value := range transport.headers {`
`350`		`- client.DERPHeader.Set(header, value)`
`351`		`- }`
`352`		`-`
`353`	`347`	`return client, nil`
`354`	`348`	`}`
`355`	`349`
`@@ -661,12 +655,18 @@ func checkWarnings(cmd cobra.Command, client codersdk.Client) error {`
`661`	`655`
`662`	`656`	`type headerTransport struct {`
`663`	`657`	`transport http.RoundTripper`
`664`		`- headers map[string]string`
	`658`	`+ header http.Header`
	`659`	`+}`
	`660`	`+`
	`661`	`+func (h *headerTransport) Header() http.Header {`
	`662`	`+ return h.header.Clone()`
`665`	`663`	`}`
`666`	`664`
`667`	`665`	`func (h headerTransport) RoundTrip(req http.Request) (*http.Response, error) {`
`668`		`- for k, v := range h.headers {`
`669`		`- req.Header.Add(k, v)`
	`666`	`+ for k, v := range h.header {`
	`667`	`+ for _, vv := range v {`
	`668`	`+ req.Header.Add(k, vv)`
	`669`	`+ }`
`670`	`670`	`}`
`671`	`671`	`return h.transport.RoundTrip(req)`
`672`	`672`	`}`
Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,6 @@ func NewConn(options Options) (conn Conn, err error) {`
`161`	`161`	`if !ok {`
`162`	`162`	`return nil, xerrors.New("get wireguard internals")`
`163`	`163`	`}`
`164`		`-`
`165`	`164`	`if options.DERPHeader != nil {`
`166`	`165`	`magicConn.SetDERPHeader(options.DERPHeader.Clone())`
`167`	`166`	`}`