peek.go 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. package util
  2. import (
  3. "bytes"
  4. "io"
  5. "strings"
  6. )
  7. // PeekedReadCloser is a ReadCloser that allows peeking into a stream and buffering it in memory.
  8. // It can be instantiated using the Peek function. After a stream has been peeked, it can still be fully
  9. // read by reading the PeekedReadCloser. It first drained from the memory buffer, and then from the remaining
  10. // underlying reader.
  11. type PeekedReadCloser struct {
  12. PeekedBytes []byte
  13. LimitReached bool
  14. peeked io.Reader
  15. underlying io.ReadCloser
  16. closed bool
  17. }
  18. // Peek reads the underlying ReadCloser into memory up until the limit and returns a PeekedReadCloser
  19. func Peek(underlying io.ReadCloser, limit int) (*PeekedReadCloser, error) {
  20. if underlying == nil {
  21. underlying = io.NopCloser(strings.NewReader(""))
  22. }
  23. peeked := make([]byte, limit)
  24. read, err := io.ReadFull(underlying, peeked)
  25. if err != nil && err != io.ErrUnexpectedEOF && err != io.EOF {
  26. return nil, err
  27. }
  28. return &PeekedReadCloser{
  29. PeekedBytes: peeked[:read],
  30. LimitReached: read == limit,
  31. underlying: underlying,
  32. peeked: bytes.NewReader(peeked[:read]),
  33. closed: false,
  34. }, nil
  35. }
  36. // Read reads from the peeked bytes and then from the underlying stream
  37. func (r *PeekedReadCloser) Read(p []byte) (n int, err error) {
  38. if r.closed {
  39. return 0, io.EOF
  40. }
  41. n, err = r.peeked.Read(p)
  42. if err == io.EOF {
  43. return r.underlying.Read(p)
  44. } else if err != nil {
  45. return 0, err
  46. }
  47. return
  48. }
  49. // Close closes the underlying stream
  50. func (r *PeekedReadCloser) Close() error {
  51. if r.closed {
  52. return io.EOF
  53. }
  54. r.closed = true
  55. return r.underlying.Close()
  56. }