| // Copyright 2019 The Fuchsia Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| package iomisc |
| |
| import ( |
| "bytes" |
| "context" |
| "errors" |
| "io" |
| |
| "go.fuchsia.dev/fuchsia/tools/lib/logger" |
| ) |
| |
| // MatchingReader is an io.Reader implementation that wraps another such |
| // implementation. It reads only up until one of the sequences has been read consecutively. |
| type MatchingReader struct { |
| r io.Reader |
| toMatch [][]byte |
| progress []int |
| matchIdx int |
| } |
| |
| // NewMatchingReader returns a MatchingReader that matches any of toMatch. |
| func NewMatchingReader(reader io.Reader, toMatch [][]byte) *MatchingReader { |
| return &MatchingReader{ |
| r: reader, |
| toMatch: toMatch, |
| progress: make([]int, len(toMatch)), |
| matchIdx: -1, |
| } |
| } |
| |
| // Match returns the first match among the bytes read, or nil if there |
| // has yet to be a match. |
| func (m *MatchingReader) Match() []byte { |
| if m.matchIdx >= 0 { |
| return m.toMatch[m.matchIdx] |
| } |
| return nil |
| } |
| |
| // Read reads from the underlying reader and checks whether the pattern |
| // has been matched among the bytes read. Once a match has been found, |
| // subsequent reads will return an io.EOF. |
| func (m *MatchingReader) Read(p []byte) (int, error) { |
| if m.matchIdx >= 0 { |
| return 0, io.EOF |
| } |
| n, err := m.r.Read(p) |
| p = p[:n] |
| for i, tm := range m.toMatch { |
| for j := 0; j < n; j++ { |
| remainingToMatch := tm[m.progress[i]:] |
| relevantP := p[j:min(len(remainingToMatch)+j, len(p))] |
| if bytes.HasPrefix(remainingToMatch, relevantP) { |
| m.progress[i] += len(relevantP) |
| if m.progress[i] == len(tm) { |
| m.matchIdx = i |
| } |
| break |
| } else { |
| m.progress[i] = 0 |
| } |
| } |
| if m.matchIdx >= 0 { |
| return n, io.EOF |
| } |
| } |
| return n, err |
| } |
| |
| // ReadUntilMatch reads from a MatchingReader until a match has been read |
| // and returns the match. |
| // Checks ctx for cancellation only between calls to m.Read(), so cancellation |
| // will not be noticed if m.Read() blocks. |
| // See https://github.com/golang/go/issues/20280 for discussion of similar issues. |
| func ReadUntilMatch(ctx context.Context, m *MatchingReader) ([]byte, error) { |
| // buf size considerations: smaller => more responsive to ctx cancellation, |
| // larger => less CPU overhead. |
| buf := make([]byte, 1024) |
| lastReadSize := 0 |
| for ctx.Err() == nil { |
| var err error |
| lastReadSize, err = m.Read(buf) |
| if err == nil { |
| continue |
| } |
| if errors.Is(err, io.EOF) { |
| if match := m.Match(); match != nil { |
| return match, nil |
| } |
| } |
| return nil, err |
| } |
| |
| // If we time out, it is helpful to see the last bytes processed. |
| logger.Debugf(ctx, "ReadUntilMatch: last %d bytes read before cancellation: %q", lastReadSize, buf[:lastReadSize]) |
| |
| return nil, ctx.Err() |
| } |
| |
| func min(a, b int) int { |
| if a <= b { |
| return a |
| } |
| return b |
| } |