problem: connection can be nil when writing and we're panic'ing.

This commit is contained in:
fiatjaf
2024-06-21 09:32:53 -03:00
parent 3862333dfa
commit a35cdff6c8
2 changed files with 8 additions and 2 deletions

View File

@ -105,7 +105,13 @@ func NewConnection(ctx context.Context, url string, requestHeader http.Header, t
}, nil }, 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 { if c.msgStateW.IsCompressed() && c.enableCompression {
c.flateWriter.Reset(c.writer) c.flateWriter.Reset(c.writer)
if _, err := io.Copy(c.flateWriter, bytes.NewReader(data)); err != nil { if _, err := io.Copy(c.flateWriter, bytes.NewReader(data)); err != nil {

View File

@ -193,7 +193,7 @@ func (r *Relay) ConnectWithTLS(ctx context.Context, tlsConfig *tls.Config) error
} }
case writeRequest := <-r.writeQueue: case writeRequest := <-r.writeQueue:
// all write requests will go through this to prevent races // 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 writeRequest.answer <- err
} }
close(writeRequest.answer) close(writeRequest.answer)