Repository: open-telemetry/opamp-go (go)

Issue:
Follow HTTP redirects on websocket handshake failure
It would be useful to check for a valid redirect on handshake failure in the `opamp-go` client, and follow it to attempt to establish a websocket connection to the location specified by the redirect. This would allow more flexibility for operators who are supporting a large network of opamp servers. It is also a backwards-compatible change.

The websocket standard permits websocket servers to offer redirection to websocket clients. (See https://www.rfc-editor.org/rfc/rfc6455#section-4.2.2)

In particular,

> 3.  The server MAY redirect the client using a 3xx status code
       [[RFC2616](https://www.rfc-editor.org/rfc/rfc2616)].  Note that this step can happen together with, before,
       or after the optional authentication step described above.

The `gorilla/websocket` library does not implement this optional redirect following, but it does return a valid HTTP response when a handshake error occurs. (See https://github.com/gorilla/websocket/blob/main/client.go#L412)

What I am proposing specifically:

In https://github.com/open-telemetry/opamp-go, check for an HTTP response status `>= 300 && <400`. Parse the location header, and if valid, set `c.url` to be the URL from the location header. Then, call `tryConnectOnce` again.

I am happy to do the implementation work this would require, and have a changelog I can contribute in the near future.

Pull request:
Follow HTTP redirects after failed WS dials
This commit allows the opamp client to follow redirects after websocket handshake failures. Redirect following is not implemented by gorilla/websocket, but can be handled by inspecting the returned response object for 3xx status and Location header.

Closes #250 

Patch:
diff --git a/client/wsclient.go b/client/wsclient.go
index 268ffdd..db17999 100644
--- a/client/wsclient.go
+++ b/client/wsclient.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/url"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
@@ -35,6 +36,10 @@ type wsClient struct {
 
 	// The sender is responsible for sending portion of the OpAMP protocol.
 	sender *internal.WSSender
+
+	// last non-nil internal error that was encountered in the conn retry loop,
+	// currently used only for testing.
+	lastInternalErr atomic.Pointer[error]
 }
 
 // NewWebSocket creates a new OpAMP Client that uses WebSocket transport.
@@ -131,8 +136,27 @@ func (c *wsClient) tryConnectOnce(ctx context.Context) (err error, retryAfter sh
 			c.common.Callbacks.OnConnectFailed(ctx, err)
 		}
 		if resp != nil {
-			c.common.Logger.Errorf(ctx, "Server responded with status=%v", resp.Status)
 			duration := sharedinternal.ExtractRetryAfterHeader(resp)
+			if resp.StatusCode >= 300 && resp.StatusCode < 400 {
+				// very liberal handling of 3xx that largely ignores HTTP semantics
+				redirect, err := resp.Location()
+				if err != nil {
+					c.common.Logger.Errorf(ctx, "%d redirect, but no valid location: %s", resp.StatusCode, err)
+					return err, duration
+				}
+				// rewrite the scheme for the sake of tolerance
+				if redirect.Scheme == "http" {
+					redirect.Scheme = "ws"
+				} else if redirect.Scheme == "https" {
+					redirect.Scheme = "wss"
+				}
+				c.common.Logger.Debugf(ctx, "%d redirect to %s", resp.StatusCode, redirect)
+				// Set the URL to the redirect, so that it connects to it on the
+				// next cycle.
+				c.url = redirect
+			} else {
+				c.common.Logger.Errorf(ctx, "Server responded with status=%v", resp.Status)
+			}
 			return err, duration
 		}
 		return err, sharedinternal.OptionalDuration{Defined: false}
@@ -167,6 +191,7 @@ func (c *wsClient) ensureConnected(ctx context.Context) error {
 		case <-timer.C:
 			{
 				if err, retryAfter := c.tryConnectOnce(ctx); err != nil {
+					c.lastInternalErr.Store(&err)
 					if errors.Is(err, context.Canceled) {
 						c.common.Logger.Debugf(ctx, "Client is stopped, will not try anymore.")
 						return err


Test patch:
diff --git a/client/wsclient_test.go b/client/wsclient_test.go
index 7696e9e..d0b8b4e 100644
--- a/client/wsclient_test.go
+++ b/client/wsclient_test.go
@@ -3,6 +3,9 @@ package client
 import (
 	"context"
 	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -177,3 +180,85 @@ func TestVerifyWSCompress(t *testing.T) {
 		})
 	}
 }
+
+func redirectServer(to string, status int) *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		http.Redirect(w, req, to, http.StatusSeeOther)
+	}))
+}
+
+func errServer() *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		w.WriteHeader(302)
+	}))
+}
+
+func TestRedirectWS(t *testing.T) {
+	redirectee := internal.StartMockServer(t)
+	tests := []struct {
+		Name       string
+		Redirector *httptest.Server
+		ExpError   bool
+	}{
+		{
+			Name:       "redirect ws scheme",
+			Redirector: redirectServer("ws://"+redirectee.Endpoint, 302),
+		},
+		{
+			Name:       "redirect http scheme",
+			Redirector: redirectServer("http://"+redirectee.Endpoint, 302),
+		},
+		{
+			Name:       "missing location header",
+			Redirector: errServer(),
+			ExpError:   true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.Name, func(t *testing.T) {
+			var conn atomic.Value
+			redirectee.OnWSConnect = func(c *websocket.Conn) {
+				conn.Store(c)
+			}
+
+			// Start an OpAMP/WebSocket client.
+			var connected int64
+			var connectErr atomic.Value
+			settings := types.StartSettings{
+				Callbacks: types.CallbacksStruct{
+					OnConnectFunc: func(ctx context.Context) {
+						atomic.StoreInt64(&connected, 1)
+					},
+					OnConnectFailedFunc: func(ctx context.Context, err error) {
+						if err != websocket.ErrBadHandshake {
+							connectErr.Store(err)
+						}
+					},
+				},
+			}
+			reURL, err := url.Parse(test.Redirector.URL)
+			assert.NoError(t, err)
+			reURL.Scheme = "ws"
+			settings.OpAMPServerURL = reURL.String()
+			client := NewWebSocket(nil)
+			startClient(t, settings, client)
+
+			// Wait for connection to be established.
+			eventually(t, func() bool {
+				return conn.Load() != nil || connectErr.Load() != nil || client.lastInternalErr.Load() != nil
+			})
+			if test.ExpError {
+				if connectErr.Load() == nil && client.lastInternalErr.Load() == nil {
+					t.Error("expected non-nil error")
+				}
+			} else {
+				assert.True(t, connectErr.Load() == nil)
+			}
+
+			// Stop the client.
+			err = client.Stop(context.Background())
+			assert.NoError(t, err)
+		})
+	}
+}