Merge branch 'fix/read-socket-response-in-chunks' into 'main'

Fix/read socket response in chunks

See merge request hectorjsmith/fail2ban-prometheus-exporter!24
This commit is contained in:
Hector 2021-08-29 15:05:39 +00:00
commit 617d711ecf

View file

@ -1,6 +1,7 @@
package socket package socket
import ( import (
"bufio"
"bytes" "bytes"
"fmt" "fmt"
"github.com/nlpodyssey/gopickle/pickle" "github.com/nlpodyssey/gopickle/pickle"
@ -9,7 +10,7 @@ import (
const ( const (
commandTerminator = "<F2B_END_COMMAND>" commandTerminator = "<F2B_END_COMMAND>"
pingCommand = "ping" pingCommand = "ping"
socketReadBufferSize = 10000 socketReadBufferSize = 1024
) )
func (s *Fail2BanSocket) sendCommand(command []string) (interface{}, error) { func (s *Fail2BanSocket) sendCommand(command []string) (interface{}, error) {
@ -33,13 +34,23 @@ func (s *Fail2BanSocket) write(command []string) error {
} }
func (s *Fail2BanSocket) read() (interface{}, error) { func (s *Fail2BanSocket) read() (interface{}, error) {
reader := bufio.NewReader(s.socket)
data := []byte{}
for {
buf := make([]byte, socketReadBufferSize) buf := make([]byte, socketReadBufferSize)
_, err := s.socket.Read(buf) _, err := reader.Read(buf)
if err != nil { if err != nil {
return nil, err return nil, err
} }
data = append(data, buf...)
containsTerminator := bytes.Contains(data, []byte(commandTerminator))
if containsTerminator {
break
}
}
bufReader := bytes.NewReader(buf) bufReader := bytes.NewReader(data)
unpickler := pickle.NewUnpickler(bufReader) unpickler := pickle.NewUnpickler(bufReader)
unpickler.FindClass = func(module, name string) (interface{}, error) { unpickler.FindClass = func(module, name string) (interface{}, error) {