fix: add custom headers for the upgrade request (#64)

* fix: add custom headers for the upgrade request
Signed-off-by: Ying Wang <yingwang@us.ibm.com>

Author: yingwang-us

fluent/client/ws_client.go

@@ -95,6 +95,7 @@ type DefaultWSConnectionFactory struct {
 	URL       string
 	AuthInfo  *IAMAuthInfo
 	TLSConfig *tls.Config
+	Header    http.Header
 }
 
 func (wcf *DefaultWSConnectionFactory) New() (ext.Conn, error) {
@@ -103,6 +104,13 @@ func (wcf *DefaultWSConnectionFactory) New() (ext.Conn, error) {
 		header = http.Header{}
 	)
 
+	// set additional custom headers. here we do not validate
+	// header names and values. Caller should make sure the
+	// headers provided are not conflict with protocols
+	if wcf.Header != nil {
+		header = wcf.Header
+	}
+
 	if wcf.AuthInfo != nil && len(wcf.AuthInfo.IAMToken()) > 0 {
 		header.Add(AuthorizationHeader, wcf.AuthInfo.IAMToken())
 	}

fluent/client/ws_client_test.go

@@ -64,6 +64,7 @@ var _ = Describe("DefaultWSConnectionFactory", func() {
 		svr               *httptest.Server
 		ch                chan struct{}
 		useTLS, testError bool
+		testHeaders       http.Header
 		customErr         *client.WSConnError
 	)
 
@@ -78,6 +79,11 @@ var _ = Describe("DefaultWSConnectionFactory", func() {
 			header := r.Header.Get(fclient.AuthorizationHeader)
 			Expect(header).To(Equal("oi"))
 
+			for k := range testHeaders {
+				v := r.Header.Get(k)
+				Expect(v).To(Equal(testHeaders[k][0]))
+			}
+
 			svrConnection, err := ws.NewConnection(wc, svrOpts)
 			if err != nil {
 				Fail("broke")
@@ -111,6 +117,7 @@ var _ = Describe("DefaultWSConnectionFactory", func() {
 
 	AfterEach(func() {
 		svr.Close()
+		testHeaders = nil
 	})
 
 	It("sends auth headers", func() {
@@ -131,6 +138,31 @@ var _ = Describe("DefaultWSConnectionFactory", func() {
 		Expect(cli.Disconnect()).ToNot(HaveOccurred())
 	})
 
+	It("sends auth headers with additional header", func() {
+		u := "ws" + strings.TrimPrefix(svr.URL, "http")
+
+		testHeaders = http.Header{
+			"User-Agent": []string{"xxxx:1.0.5"}, // user agent
+			"X-a":        []string{""},           // empty value
+			"X-b":        []string{"value"},      // some string value
+		}
+
+		cli := fclient.NewWS(client.WSConnectionOptions{
+			Factory: &client.DefaultWSConnectionFactory{
+				URL: u,
+				TLSConfig: &tls.Config{
+					InsecureSkipVerify: true,
+				},
+				AuthInfo: NewIAMAuthInfo("oi"),
+				Header:   testHeaders,
+			},
+		})
+
+		Expect(cli.Connect()).ToNot(HaveOccurred())
+		Eventually(ch).Should(Receive())
+		Expect(cli.Disconnect()).ToNot(HaveOccurred())
+	})
+
 	When("sends wrong url, expects error", func() {
 
 		BeforeEach(func() {