From 2c5f8019c9d071f6d6508815e7afe4525b8ded9c Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 17 Mar 2023 13:34:25 +1300
Subject: [PATCH] websocket: don't set a WSS multiaddr for accepted unencrypted
 conns

---
 p2p/transport/websocket/listener.go       |  6 +++++-
 p2p/transport/websocket/websocket_test.go | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/p2p/transport/websocket/listener.go b/p2p/transport/websocket/listener.go
index 70dbce0fea..6d0fca5af0 100644
--- a/p2p/transport/websocket/listener.go
+++ b/p2p/transport/websocket/listener.go
@@ -19,6 +19,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
 
@@ -80,6 +83,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
@@ -96,7 +100,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 0c7fc5f974..10878397cf 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()