feat: Add WebSocket transport support for AMQP 1.0 connections (#78)

* feat: Add WebSocket transport support for AMQP 1.0
* chore: revert changes to ExtractWithoutPassword helper

Author: vedanthnyk25

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/Azure/go-amqp v1.5.1
 	github.com/golang-jwt/jwt/v5 v5.2.2
 	github.com/google/uuid v1.6.0
+	github.com/gorilla/websocket v1.5.3
 	github.com/onsi/ginkgo/v2 v2.23.4
 	github.com/onsi/gomega v1.38.0
 )

go.sum

@@ -1,7 +1,3 @@
-github.com/Azure/go-amqp v1.4.0 h1:Xj3caqi4comOF/L1Uc5iuBxR/pB6KumejC01YQOqOR4=
-github.com/Azure/go-amqp v1.4.0/go.mod h1:vZAogwdrkbyK3Mla8m/CxSc/aKdnTZ4IbPxl51Y5WZE=
-github.com/Azure/go-amqp v1.5.0 h1:GRiQK1VhrNFbyx5VlmI6BsA1FCp27W5rb9kxOZScnTo=
-github.com/Azure/go-amqp v1.5.0/go.mod h1:vZAogwdrkbyK3Mla8m/CxSc/aKdnTZ4IbPxl51Y5WZE=
 github.com/Azure/go-amqp v1.5.1 h1:WyiPTz2C3zVvDL7RLAqwWdeoYhMtX62MZzQoP09fzsU=
 github.com/Azure/go-amqp v1.5.1/go.mod h1:vZAogwdrkbyK3Mla8m/CxSc/aKdnTZ4IbPxl51Y5WZE=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
@@ -20,6 +16,8 @@ github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 h1:BHT72Gu3keYf3ZEu2J
 github.com/google/pprof v0.0.0-20250403155104-27863c87afa6/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=

pkg/rabbitmqamqp/amqp_connection.go

@@ -3,15 +3,19 @@ package rabbitmqamqp
 import (
 	"context"
 	"crypto/tls"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"math/rand"
+	"net/http"
+	"net/url"
 	"sync"
 	"sync/atomic"
 	"time"
 
 	"github.com/Azure/go-amqp"
 	"github.com/google/uuid"
+	"github.com/gorilla/websocket"
 )
 
 var ErrConnectionClosed = errors.New("connection is closed")
@@ -453,10 +457,45 @@ func (a *AmqpConnection) open(ctx context.Context, address string, connOptions *
 		TLSConfig:    connOptions.TLSConfig,
 		WriteTimeout: connOptions.WriteTimeout,
 	}
-	azureConnection, err = amqp.Dial(ctx, address, amqpLiteConnOptions)
-	if err != nil && (connOptions.TLSConfig != nil || uri.Scheme == AMQPS) {
-		Error("Failed to open TLS connection", fmt.Sprintf("%s://%s", uri.Scheme, uri.Host), err, "ID", connOptions.Id)
-		return fmt.Errorf("failed to open TLS connection: %w", err)
+
+	u, err := url.Parse(address)
+	if err != nil {
+		return err
+	}
+
+	if u.Scheme == "ws" || u.Scheme == "wss" {
+
+		wsAddress, wsHeaders, err := sanitizeWebSocketURL(address)
+		if err != nil {
+			Error("Failed to sanitize websocket URL", "url", ExtractWithoutPassword(address), "error", err, "ID", connOptions.Id)
+			return fmt.Errorf("failed to sanitize websocket URL: %w", err)
+		}
+
+		// Create a WebSocket dialer
+		dialer := websocket.DefaultDialer
+		if u.Scheme == "wss" && connOptions.TLSConfig != nil {
+			dialer.TLSClientConfig = connOptions.TLSConfig
+		}
+
+		// Dial the WebSocket server
+		wsConn, _, err := dialer.Dial(wsAddress, wsHeaders)
+		if err != nil {
+			Error("Failed to open a websocket connection", "url", ExtractWithoutPassword(wsAddress), "error", err, "ID", connOptions.Id)
+			return fmt.Errorf("failed to open a websocket connection: %w", err)
+		}
+
+		// Wrap the WebSocket connection in a WebSocketConn
+		neConn := NewWebSocketConn(wsConn)
+		azureConnection, err = amqp.NewConn(ctx, neConn, amqpLiteConnOptions)
+		if err != nil {
+			Error("Failed to open AMQP over WebSocket connection", "url", ExtractWithoutPassword(address), "error", err, "ID", connOptions.Id)
+		}
+	} else {
+		azureConnection, err = amqp.Dial(ctx, address, amqpLiteConnOptions)
+		if err != nil && (connOptions.TLSConfig != nil || uri.Scheme == AMQPS) {
+			Error("Failed to open TLS connection", fmt.Sprintf("%s://%s", uri.Scheme, uri.Host), err, "ID", connOptions.Id)
+			return fmt.Errorf("failed to open TLS connection: %w", err)
+		}
 	}
 	if err != nil {
 		Error("Failed to open connection", "url", ExtractWithoutPassword(address), "error", err, "ID", connOptions.Id)
@@ -701,3 +740,36 @@ func (a *AmqpConnection) RefreshToken(background context.Context, token string)
 }
 
 //*** end management section ***
+
+// sanitizeWebSocketURL ensures the URL is correctly formatted for the Gorilla websocket dialer.
+func sanitizeWebSocketURL(rawURL string) (string, http.Header, error) {
+	u, err := url.Parse(rawURL)
+	if err != nil {
+		return "", nil, err
+	}
+
+	if u.Scheme != "ws" && u.Scheme != "wss" {
+		return "", nil, fmt.Errorf("invalid websocket scheme: %s", u.Scheme)
+	}
+
+	// Prepare Headers for Auth
+	headers := http.Header{}
+	if u.User != nil {
+		username := u.User.Username()
+		password, _ := u.User.Password()
+
+		// Construct Basic Auth Header manually
+		auth := base64.StdEncoding.EncodeToString([]byte(username + ":" + password))
+		headers.Add("Authorization", "Basic "+auth)
+
+		u.User = nil
+	}
+
+	if u.Path == "" {
+		u.Path = "/"
+	} else if u.Path[0] != '/' {
+		u.Path = "/" + u.Path
+	}
+
+	return u.String(), headers, nil
+}

pkg/rabbitmqamqp/transport_websocket.go

@@ -0,0 +1,74 @@
+package rabbitmqamqp
+
+import (
+	"io"
+	"net"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+type WebSocketConn struct {
+	conn   *websocket.Conn
+	reader io.Reader
+}
+
+func NewWebSocketConn(conn *websocket.Conn) *WebSocketConn {
+	return &WebSocketConn{conn: conn}
+}
+
+func (ws *WebSocketConn) Read(b []byte) (n int, err error) {
+	if ws.reader == nil {
+		messageType, reader, err := ws.conn.NextReader()
+		if err != nil {
+			return 0, err
+		}
+		if messageType != websocket.BinaryMessage {
+			return 0, io.ErrUnexpectedEOF
+		}
+		ws.reader = reader
+	}
+	n, err = ws.reader.Read(b)
+	if err == io.EOF {
+		ws.reader = nil
+		err = nil
+	}
+	return n, err
+}
+
+func (ws *WebSocketConn) Write(b []byte) (n int, err error) {
+	err = ws.conn.WriteMessage(websocket.BinaryMessage, b)
+	if err != nil {
+		return 0, err
+	}
+
+	return len(b), nil
+}
+
+func (ws *WebSocketConn) Close() error {
+	return ws.conn.Close()
+}
+
+func (ws *WebSocketConn) LocalAddr() net.Addr {
+	return ws.conn.LocalAddr()
+}
+
+func (ws *WebSocketConn) RemoteAddr() net.Addr {
+	return ws.conn.RemoteAddr()
+}
+
+func (ws *WebSocketConn) SetDeadline(t time.Time) error {
+	err := ws.conn.SetReadDeadline(t)
+	if err != nil {
+		return err
+	}
+	return ws.conn.SetWriteDeadline(t)
+}
+
+func (ws *WebSocketConn) SetReadDeadline(t time.Time) error {
+	return ws.conn.SetReadDeadline(t)
+}
+
+func (ws *WebSocketConn) SetWriteDeadline(t time.Time) error {
+	return ws.conn.SetWriteDeadline(t)
+}

pkg/rabbitmqamqp/uri.go

@@ -13,13 +13,15 @@ import (
 )
 
 var (
-	errURIScheme     = errors.New("AMQP scheme must be either 'amqp://' or 'amqps://'")
+	errURIScheme     = errors.New("AMQP scheme must be either 'amqp://', 'amqps://', 'ws://' or 'wss://'")
 	errURIWhitespace = errors.New("URI must not contain whitespace")
 )
 
 var schemePorts = map[string]int{
 	"amqp":  5672,
 	"amqps": 5671,
+	"wss":   15675,
+	"ws":    15675,
 }
 
 var defaultURI = URI{
@@ -117,7 +119,6 @@ func ParseURI(uri string) (URI, error) {
 }
 
 // Extract the Uri by omitting the password
-
 func ExtractWithoutPassword(addr string) string {
 	u, err := ParseURI(addr)
 	if err != nil {