highperfocused/go-nostr
1
0
Fork 0
mirror of https://github.com/nbd-wtf/go-nostr.git synced 2025-07-01 02:51:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

check if messages are compressed on receive

Browse Source
This commit is contained in:
Marc Tarnutzer
2023-05-06 01:48:01 +02:00
parent c86e907142
commit 69b9d82bb1
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
connection.go
View File

@ -26,6 +26,7 @@ type Connection struct {
reader *wsutil.Reader reader *wsutil.Reader
flateWriter *wsflate.Writer flateWriter *wsflate.Writer
writer *wsutil.Writer writer *wsutil.Writer
msgState *wsflate.MessageState
mutex sync.Mutex mutex sync.Mutex
} }
@ -67,6 +68,7 @@ func NewConnection(ctx context.Context, url string, requestHeader http.Header) (
Source: conn, Source: conn,
State: state, State: state,
OnIntermediate: controlHandler, OnIntermediate: controlHandler,
CheckUTF8: false,
Extensions: []wsutil.RecvExtension{ Extensions: []wsutil.RecvExtension{
&msgState, &msgState,
}, },
@ -98,6 +100,7 @@ func NewConnection(ctx context.Context, url string, requestHeader http.Header) (
flateReader: flateReader, flateReader: flateReader,
reader: reader, reader: reader,
flateWriter: flateWriter, flateWriter: flateWriter,
msgState: &msgState,
writer: writer, writer: writer,
}, nil }, nil
} }
@ -106,7 +109,7 @@ func (c *Connection) WriteJSON(v any) error {
c.mutex.Lock() c.mutex.Lock()
defer c.mutex.Unlock() defer c.mutex.Unlock()
if c.enableCompression { if c.enableCompression && c.msgState.IsCompressed() {
c.flateWriter.Reset(c.writer) c.flateWriter.Reset(c.writer)
if err := json.NewEncoder(c.flateWriter).Encode(v); err != nil { if err := json.NewEncoder(c.flateWriter).Encode(v); err != nil {
return fmt.Errorf("failed to encode json: %w", err) return fmt.Errorf("failed to encode json: %w", err)
@ -141,7 +144,7 @@ func (c *Connection) WriteMessage(data []byte) error {
c.mutex.Lock() c.mutex.Lock()
defer c.mutex.Unlock() defer c.mutex.Unlock()
if c.enableCompression { if c.msgState.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 {
return fmt.Errorf("failed to write message: %w", err) return fmt.Errorf("failed to write message: %w", err)
@ -194,7 +197,7 @@ func (c *Connection) ReadMessage(ctx context.Context) ([]byte, error) {
} }
buf := new(bytes.Buffer) buf := new(bytes.Buffer)
if c.enableCompression { if c.msgState.IsCompressed() && c.enableCompression {
c.flateReader.Reset(c.reader) c.flateReader.Reset(c.reader)
if _, err := io.Copy(buf, c.flateReader); err != nil { if _, err := io.Copy(buf, c.flateReader); err != nil {
return nil, fmt.Errorf("failed to read message: %w", err) return nil, fmt.Errorf("failed to read message: %w", err)