go

Reading first byte from a stream in Golang


I have the following function in Golang:


func ReadFromStdinIfAvailable(cmd *cobra.Command, args []string) ([]byte, error) {
    if len(args) == 0 {
        r := bufio.NewReader(os.Stdin)
        firstByte, err := r.Peek(1)
        if err != nil {
            return nil, err
        }

        if len(firstByte) == 0 {
            cmd.Println(userstrings.NoFilenameProvidedErrorString)
            return nil, errors.New(userstrings.NoFilenameProvidedErrorString)
        }

        var bytesResult []byte
        scanner := bufio.NewScanner(r)
        for scanner.Scan() {
            // Read from stdin - allocate a buffer to read into
            bytesResult = append(bytesResult, scanner.Bytes()...)
        }
        return bytesResult, nil
    }
    return nil, errors.New(userstrings.NoStdInProvidedErrorString)
}

It's a very small amount of data (<10 MB to be sure). However, the data should be there from the first read - if it's not, there's an error.

However, when i do this, it hangs on scanner.Scan() which is, to be expected, it's a stream.

How can I peek into the os.Stdin stream to see if there's any data there before I start reading it?


Solution

  • The answer is to read the first n bytes from the stream, then create a new io.Reader that has the same data you've already read followed by the remaining data in the stream.

    Here is a playground link showing this example.

    reader := strings.NewReader("i am the very model of a modern major general")
    
    buf := make([]byte, 10) // Read first 10 bytes
    _, _ = io.ReadAtLeast(reader, buf, 10)
    
    fmt.Println(string(buf))
    
    // Make a new stream that concats the data you read with what is remaining
    newReader := io.MultiReader(bytes.NewReader(buf), reader)
    data, _ := io.ReadAll(newReader)
    
    fmt.Println(string(data))