从Golang中的io.Reader到io.Writer读取/复制一定数量的字节，或者如果超过一定字节数限制，则返回错误？

Question

I have an io.Reader in Golang and I want to double-check that the size of its data is below a predetermined maximum before or while running io.Copy() to save it to disk using io.Writer.

Since the file data in io.Reader could theoretically be quite large, I want to minimize memory usage and processing here if avoidable.
I don't think there's a function that's like io.CopyLessThanOrEqualToThisManyBytesOrReturnError(), but I did notice that io.ReadFull() can do the opposite to return an error if not enough bytes are there to fill the provided buffer.

Does anyone have a solution to this?

---

EDIT:
To clarify, copying a fraction of the data is not OK. It either needs to fail if it's over the threshold, or work if it's under.

Answer 1

Since io.Reader interface not knows anything about size or length of underlying data, there is only one solution for this problem:

You may use one buffer with maximum size nMax (predetermined maximum)+1 and in every call to Your CopyLessThanOrEqualToThisManyBytesOrReturnError function, inside this function read input and buffer it, and check for this buffer length, if it is less than or equal to nMax then do io.Write, otherwise return error:

const nMax = 5 // your predetermined maximum

func CopyLessThanOrEqualToThisManyBytesOrReturnError(r io.Reader, w io.Writer) error {
    var buf = make([]byte, nMax+1)
    nRead, e := io.ReadFull(r, buf)
    if nRead > 0 && nRead <= nMax {
        w.Write(buf[:nRead])
        return nil
    }
    if nRead > nMax {
        return fmt.Errorf("there is more data")
    }
    return e
}

---

Like this working sample code unsing string:

package main

import (
    "fmt"
    "io"
    "os"
    "strings"
)

const nMax = 5 // your predetermined maximum

func CopyLessThanOrEqualToThisManyBytesOrReturnError(r io.Reader, w io.Writer) error {
    var buf = make([]byte, nMax+1)
    nRead, e := io.ReadFull(r, buf)
    if nRead > 0 && nRead <= nMax {
        w.Write(buf[:nRead])
        return nil
    }
    if nRead > nMax {
        return fmt.Errorf("there is more data")
    }
    return e
}

func main() {
    r := strings.NewReader("123456789")
    err := CopyLessThanOrEqualToThisManyBytesOrReturnError(r, os.Stdout)
    if err != nil {
        fmt.Println(err) // there is more data
    }

    r = strings.NewReader("123
")
    err = CopyLessThanOrEqualToThisManyBytesOrReturnError(r, os.Stdout) // 123
    if err != nil {
        fmt.Println(err)
    }

    r = strings.NewReader("")
    err = CopyLessThanOrEqualToThisManyBytesOrReturnError(r, os.Stdout)
    if err != nil {
        fmt.Println(err) // EOF
    }
}

output:

there is more data
123
EOF

---

Working sample code, using files:

package main

import (
    "fmt"
    "io"
    "os"
)

const nMax = 5 // your predetermined maximum

func CopyLessThanOrEqualToThisManyBytesOrReturnError(r io.Reader, w io.Writer) error {
    var buf = make([]byte, nMax+1)
    nRead, e := io.ReadFull(r, buf)
    if nRead > 0 && nRead <= nMax {
        w.Write(buf[:nRead])
        return nil
    }
    if nRead > nMax {
        return fmt.Errorf("there is more data")
    }
    return e
}

func main() {
    r, err := os.Open("input.bin")
    if err != nil {
        panic(err)
    }
    defer r.Close()

    w, err := os.Create("output.bin")
    if err != nil {
        panic(err)
    }
    defer w.Close()

    err = CopyLessThanOrEqualToThisManyBytesOrReturnError(r, w)
    if err != nil {
        fmt.Println(err)
    }
    fmt.Println("Done.")
}

---

Working sample code, using []byte:

package main

import (
    "bytes"
    "fmt"
    "io"
)

const nMax = 5 // your predetermined maximum

func CopyLessThanOrEqualToThisManyBytesOrReturnError(r io.Reader, w io.Writer) error {
    var buf = make([]byte, nMax+1)
    nRead, e := io.ReadFull(r, buf)
    if nRead > 0 && nRead <= nMax {
        w.Write(buf[:nRead])
        return nil
    }
    if nRead > nMax {
        return fmt.Errorf("there is more data")
    }
    return e
}

func main() {
    bs := []byte{1, 2, 3, 4, 5}
    r := bytes.NewReader(bs)

    w := &bytes.Buffer{}

    err := CopyLessThanOrEqualToThisManyBytesOrReturnError(r, w)
    if err != nil {
        fmt.Println(err)
    }

    fmt.Println("Done.")
    fmt.Println(w.Bytes())
}

output:

Done.
[1 2 3 4 5]

Answer 2

You can use io.CopyN: https://golang.org/pkg/io/#CopyN which will return a nil error if it successfully copies exactly N bytes, or else an io.EOF (or possibly another error) if less than N bytes.

For example:

func copyMax(dst io.Writer, src io.Reader, n int64) error
    _, err := io.CopyN(dst, src, n)
    if err != nil {
        // Don't care if there's less available
        return nil
    }

    nextByte := make([]byte, 1)
    nRead, _ := io.ReadFull(src, nextByte)
    if nRead > 0 {
        // There's too much data
        return errors.New("Too much data")
    }

    return nil
}