Browse Source

Add parameter to indicate ReadBytesToWriter to pass size

master
Jaeha Choi 2 months ago
parent
commit
fc824cf8d6
  1. 10
      util/util.go
  2. 4
      util/util_test.go

10
util/util.go

@ -86,7 +86,7 @@ func ReadBytes(reader io.Reader) (b []byte, err error) {
// First four bytes of reader should be uint32 size of the message,
// represented in big endian.
// Common usage for this function is to read from net.Conn, and write to temp file.
func ReadBytesToWriter(reader io.Reader, writer io.Writer) (n int, err error) {
func ReadBytesToWriter(reader io.Reader, writer io.Writer, writeWithSize bool) (n int, err error) {
// Read message size
size, err := readSize(reader)
if err != nil {
@ -94,6 +94,14 @@ func ReadBytesToWriter(reader io.Reader, writer io.Writer) (n int, err error) {
return 0, err
}
if writeWithSize {
err := writeSize(writer, size)
if err != nil {
log.Debug(err)
return 0, err
}
}
totalReceived, err := readWrite(reader, writer, size)
if err != nil {
return totalReceived, err

4
util/util_test.go

@ -1120,7 +1120,7 @@ func TestReadBytesTemp(t *testing.T) {
buf.Write(testByte)
testByte = []byte("test")
buf.Write(testByte)
temp, err := ReadBytesToWriter(&buf, &output)
temp, err := ReadBytesToWriter(&buf, &output, false)
if err != nil || temp != 4100 {
t.Error(err)
return
@ -1204,7 +1204,7 @@ func BenchmarkReadBytesTemp(b *testing.B) {
}
buf.Write(testByte)
buf.Write([]byte("test"))
temp, err := ReadBytesToWriter(&buf, &output)
temp, err := ReadBytesToWriter(&buf, &output, false)
if err != nil || temp != 4100 {
b.Error(err)
return

Loading…
Cancel
Save