From eeb685f68fbeff63e6eaa748ac8aef91215f1b6e Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 16 Mar 2023 18:39:25 -0700
Subject: [PATCH] websocket: don't set a WSS multiaddr for accepted unencrypted
 conns (#2199)

---
 p2p/transport/websocket/listener.go       |  8 ++++++--
 p2p/transport/websocket/websocket_test.go | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/p2p/transport/websocket/listener.go b/p2p/transport/websocket/listener.go
index 8f2f02c31..ab9a73f8a 100644
--- a/p2p/transport/websocket/listener.go
+++ b/p2p/transport/websocket/listener.go
@@ -20,6 +20,9 @@ import (
 type listener struct {
 	nl     net.Listener
 	server http.Server
+	// The Go standard library sets the http.Server.TLSConfig no matter if this is a WS or WSS,
+	// so we can't rely on checking if server.TLSConfig is set.
+	isWss bool
 
 	laddr ma.Multiaddr
 
@@ -81,6 +84,7 @@ func newListener(a ma.Multiaddr, tlsConf *tls.Config) (*listener, error) {
 	}
 	ln.server = http.Server{Handler: ln}
 	if parsed.isWSS {
+		ln.isWss = true
 		ln.server.TLSConfig = tlsConf
 	}
 	return ln, nil
@@ -88,7 +92,7 @@ func newListener(a ma.Multiaddr, tlsConf *tls.Config) (*listener, error) {
 
 func (l *listener) serve() {
 	defer close(l.closed)
-	if l.server.TLSConfig == nil {
+	if !l.isWss {
 		l.server.Serve(l.nl)
 	} else {
 		l.server.ServeTLS(l.nl, "", "")
@@ -97,7 +101,7 @@ func (l *listener) serve() {
 
 func (l *listener) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	scheme := "ws"
-	if l.server.TLSConfig != nil {
+	if l.isWss {
 		scheme = "wss"
 	}
 
diff --git a/p2p/transport/websocket/websocket_test.go b/p2p/transport/websocket/websocket_test.go
index 0c7fc5f97..10878397c 100644
--- a/p2p/transport/websocket/websocket_test.go
+++ b/p2p/transport/websocket/websocket_test.go
@@ -308,6 +308,16 @@ func TestWebsocketTransport(t *testing.T) {
 	ttransport.SubtestTransport(t, ta, tb, "/ip4/127.0.0.1/tcp/0/ws", "peerA")
 }
 
+func isWSS(addr ma.Multiaddr) bool {
+	if _, err := addr.ValueForProtocol(ma.P_WSS); err == nil {
+		return true
+	}
+	if _, err := addr.ValueForProtocol(ma.P_WS); err == nil {
+		return false
+	}
+	panic("not a WebSocket address")
+}
+
 func connectAndExchangeData(t *testing.T, laddr ma.Multiaddr, secure bool) {
 	var opts []Option
 	var tlsConf *tls.Config
@@ -339,6 +349,8 @@ func connectAndExchangeData(t *testing.T, laddr ma.Multiaddr, secure bool) {
 		require.NoError(t, err)
 		c, err := tpt.Dial(context.Background(), l.Multiaddr(), server)
 		require.NoError(t, err)
+		require.Equal(t, secure, isWSS(c.LocalMultiaddr()))
+		require.Equal(t, secure, isWSS(c.RemoteMultiaddr()))
 		str, err := c.OpenStream(context.Background())
 		require.NoError(t, err)
 		defer str.Close()
@@ -349,6 +361,8 @@ func connectAndExchangeData(t *testing.T, laddr ma.Multiaddr, secure bool) {
 	c, err := l.Accept()
 	require.NoError(t, err)
 	defer c.Close()
+	require.Equal(t, secure, isWSS(c.LocalMultiaddr()))
+	require.Equal(t, secure, isWSS(c.RemoteMultiaddr()))
 	str, err := c.AcceptStream()
 	require.NoError(t, err)
 	defer str.Close()