From a35cdff6c828a737b4dfac9d1961acf5dee1930f Mon Sep 17 00:00:00 2001 From: fiatjaf Date: Fri, 21 Jun 2024 09:32:53 -0300 Subject: [PATCH] problem: connection can be nil when writing and we're panic'ing. --- connection.go | 8 +++++++- relay.go | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/connection.go b/connection.go index 86b8ad6..da1161d 100644 --- a/connection.go +++ b/connection.go @@ -105,7 +105,13 @@ func NewConnection(ctx context.Context, url string, requestHeader http.Header, t }, nil } -func (c *Connection) WriteMessage(data []byte) error { +func (c *Connection) WriteMessage(ctx context.Context, data []byte) error { + select { + case <-ctx.Done(): + return errors.New("context canceled") + default: + } + if c.msgStateW.IsCompressed() && c.enableCompression { c.flateWriter.Reset(c.writer) if _, err := io.Copy(c.flateWriter, bytes.NewReader(data)); err != nil { diff --git a/relay.go b/relay.go index c624b27..dd6bff0 100644 --- a/relay.go +++ b/relay.go @@ -193,7 +193,7 @@ func (r *Relay) ConnectWithTLS(ctx context.Context, tlsConfig *tls.Config) error } case writeRequest := <-r.writeQueue: // all write requests will go through this to prevent races - if err := r.Connection.WriteMessage(writeRequest.msg); err != nil { + if err := r.Connection.WriteMessage(r.connectionContext, writeRequest.msg); err != nil { writeRequest.answer <- err } close(writeRequest.answer)