From 35faff858ad65b23d683d77f20a9f7cb942bb8e6 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Sun, 30 Jul 2023 17:12:30 -0300
Subject: [PATCH] reuse buffer when reading messages from websocket.

---
 connection.go | 17 ++++++++---------
 relay.go      |  9 ++++++---
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/connection.go b/connection.go
index 445a43f..5385f35 100644
--- a/connection.go
+++ b/connection.go
@@ -121,23 +121,23 @@ func (c *Connection) WriteMessage(data []byte) error {
 	return nil
 }
 
-func (c *Connection) ReadMessage(ctx context.Context) ([]byte, error) {
+func (c *Connection) ReadMessage(ctx context.Context, buf io.Writer) error {
 	for {
 		select {
 		case <-ctx.Done():
-			return nil, errors.New("context canceled")
+			return errors.New("context canceled")
 		default:
 		}
 
 		h, err := c.reader.NextFrame()
 		if err != nil {
 			c.conn.Close()
-			return nil, fmt.Errorf("failed to advance frame: %w", err)
+			return fmt.Errorf("failed to advance frame: %w", err)
 		}
 
 		if h.OpCode.IsControl() {
 			if err := c.controlHandler(h, c.reader); err != nil {
-				return nil, fmt.Errorf("failed to handle control frame: %w", err)
+				return fmt.Errorf("failed to handle control frame: %w", err)
 			}
 		} else if h.OpCode == ws.OpBinary ||
 			h.OpCode == ws.OpText {
@@ -145,23 +145,22 @@ func (c *Connection) ReadMessage(ctx context.Context) ([]byte, error) {
 		}
 
 		if err := c.reader.Discard(); err != nil {
-			return nil, fmt.Errorf("failed to discard: %w", err)
+			return fmt.Errorf("failed to discard: %w", err)
 		}
 	}
 
-	buf := new(bytes.Buffer)
 	if c.msgState.IsCompressed() && c.enableCompression {
 		c.flateReader.Reset(c.reader)
 		if _, err := io.Copy(buf, c.flateReader); err != nil {
-			return nil, fmt.Errorf("failed to read message: %w", err)
+			return fmt.Errorf("failed to read message: %w", err)
 		}
 	} else {
 		if _, err := io.Copy(buf, c.reader); err != nil {
-			return nil, fmt.Errorf("failed to read message: %w", err)
+			return fmt.Errorf("failed to read message: %w", err)
 		}
 	}
 
-	return buf.Bytes(), nil
+	return nil
 }
 
 func (c *Connection) Close() error {
diff --git a/relay.go b/relay.go
index 68c18ac..57ee75d 100644
--- a/relay.go
+++ b/relay.go
@@ -1,6 +1,7 @@
 package nostr
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"log"
@@ -232,16 +233,18 @@ func (r *Relay) Connect(ctx context.Context) error {
 
 	// general message reader loop
 	go func() {
+		buf := new(bytes.Buffer)
+
 		for {
-			message, err := conn.ReadMessage(r.connectionContext)
-			if err != nil {
+			buf.Reset()
+			if err := conn.ReadMessage(r.connectionContext, buf); err != nil {
 				r.ConnectionError = err
 				r.Close()
 				break
 			}
 
+			message := buf.Bytes()
 			debugLogf("{%s} %v\n", r.URL, message)
-
 			envelope := ParseMessage(message)
 			if envelope == nil {
 				continue