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

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 {

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)