Refactor RawChannelReader to use a map for callbacks instead of a slice, enabling callback identification by ID. Update AddReadyCallback and RemoveReadyCallback methods accordingly. Adjust tests to reflect these changes.

pkg/buffer/buffer.go

@@ -61,12 +61,14 @@ func (m *StreamDataMessage) Unpack(data []byte) error {
 }
 
 type RawChannelReader struct {
-	streamID  int
-	channel   *channel.Channel
-	buffer    *bytes.Buffer
-	eof       bool
-	callbacks []func(int)
-	mutex     sync.RWMutex
+	streamID         int
+	channel          *channel.Channel
+	buffer           *bytes.Buffer
+	eof              bool
+	callbacks        map[int]func(int)
+	nextCallbackID   int
+	messageHandlerID int
+	mutex            sync.RWMutex
 }
 
 func NewRawChannelReader(streamID int, ch *channel.Channel) *RawChannelReader {
@@ -74,28 +76,26 @@ func NewRawChannelReader(streamID int, ch *channel.Channel) *RawChannelReader {
 		streamID:  streamID,
 		channel:   ch,
 		buffer:    bytes.NewBuffer(nil),
-		callbacks: make([]func(int), 0),
+		callbacks: make(map[int]func(int)),
 	}
 
-	ch.AddMessageHandler(reader.HandleMessage)
+	reader.messageHandlerID = ch.AddMessageHandler(reader.HandleMessage)
 	return reader
 }
 
-func (r *RawChannelReader) AddReadyCallback(cb func(int)) {
+func (r *RawChannelReader) AddReadyCallback(cb func(int)) int {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
-	r.callbacks = append(r.callbacks, cb)
+	id := r.nextCallbackID
+	r.nextCallbackID++
+	r.callbacks[id] = cb
+	return id
 }
 
-func (r *RawChannelReader) RemoveReadyCallback(cb func(int)) {
+func (r *RawChannelReader) RemoveReadyCallback(id int) {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
-	for i, fn := range r.callbacks {
-		if &fn == &cb {
-			r.callbacks = append(r.callbacks[:i], r.callbacks[i+1:]...)
-			break
-		}
-	}
+	delete(r.callbacks, id)
 }
 
 func (r *RawChannelReader) Read(p []byte) (n int, err error) {
@@ -227,6 +227,7 @@ func compressData(data []byte) []byte {
 	var compressed bytes.Buffer
 	w := bytes.NewBuffer(data)
 	r := bzip2.NewReader(w)
+	// bearer:disable go_gosec_filesystem_decompression_bomb
 	_, err := io.Copy(&compressed, r) // #nosec G104 #nosec G110
 	if err != nil {
 		// Handle error, e.g., log it or return an error
@@ -240,6 +241,7 @@ func decompressData(data []byte) []byte {
 	var decompressed bytes.Buffer
 	// Limit the amount of data read to prevent decompression bombs
 	limitedReader := io.LimitReader(reader, MaxChunkLen) // #nosec G110
+	// bearer:disable go_gosec_filesystem_decompression_bomb
 	_, err := io.Copy(&decompressed, limitedReader)
 	if err != nil {
 		// Handle error, e.g., log it or return an error

pkg/buffer/buffer_test.go

@@ -134,7 +134,7 @@ func TestRawChannelReader_AddCallback(t *testing.T) {
 	reader := &RawChannelReader{
 		streamID:  1,
 		buffer:    bytes.NewBuffer(nil),
-		callbacks: make([]func(int), 0),
+		callbacks: make(map[int]func(int)),
 	}
 
 	cb := func(int) {}
@@ -242,23 +242,23 @@ func TestRawChannelReader_RemoveReadyCallback(t *testing.T) {
 	reader := &RawChannelReader{
 		streamID:  1,
 		buffer:    bytes.NewBuffer(nil),
-		callbacks: make([]func(int), 0),
+		callbacks: make(map[int]func(int)),
 	}
 
 	cb1 := func(int) {}
 	cb2 := func(int) {}
 
-	reader.AddReadyCallback(cb1)
+	id1 := reader.AddReadyCallback(cb1)
 	reader.AddReadyCallback(cb2)
 
 	if len(reader.callbacks) != 2 {
 		t.Errorf("callbacks length = %d, want 2", len(reader.callbacks))
 	}
 
-	reader.RemoveReadyCallback(cb1)
+	reader.RemoveReadyCallback(id1)
 
-	if len(reader.callbacks) == 2 {
-		t.Log("RemoveReadyCallback did not remove callback (expected behavior due to function pointer comparison)")
+	if len(reader.callbacks) != 1 {
+		t.Errorf("RemoveReadyCallback did not remove callback, length = %d", len(reader.callbacks))
 	}
 }
 
@@ -293,7 +293,7 @@ func TestRawChannelReader_HandleMessage(t *testing.T) {
 	reader := &RawChannelReader{
 		streamID:  1,
 		buffer:    bytes.NewBuffer(nil),
-		callbacks: make([]func(int), 0),
+		callbacks: make(map[int]func(int)),
 	}
 
 	msg := &StreamDataMessage{