reader.go

package netargv

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"iter"
	"strings"
	"time"
)

var (
	// ErrStreamDesynced is returned by ReadNext when a previous call left the
	// stream in an unrecoverable state. The Reader must be discarded.
	ErrStreamDesynced = errors.New("stream desynced: Reader is no longer usable")

	// ErrParseFailure is returned when a message header cannot be tokenised.
	ErrParseFailure = errors.New("message parse error")

	// ErrMissingPayloadLength is returned when the payload marker `--`` appears in a header without a following byte count.
	ErrMissingPayloadLength = errors.New("no payload length declared")

	// ErrInvalidPayloadLength is returned when the byte count after -- is not a valid non-negative integer.
	ErrInvalidPayloadLength = errors.New("payload length must be a non-negative integer")

	// ErrPayloadTooLarge is returned when the declared payload length exceeds the Reader's MaxPayloadSize.
	ErrPayloadTooLarge = errors.New("payload length exceeds maximum allowed size")
)

// timeoutSetter is satisfied by net.Conn and anything else that supports
// per-read deadline control.
type timeoutSetter interface {
	SetReadDeadline(time.Time) error
}

// Reader reads netargv messages from a buffered stream.
// It is not safe for concurrent use.
//
// If a context cancellation or deadline interrupts a read mid-stream, the
// Reader latches a fatal error and refuses all subsequent calls. The caller
// must discard the Reader and, if appropriate, close and redial the connection.
type Reader struct {
	src            *bufio.Reader
	raw            io.Reader
	fatalErr       error
	maxPayloadSize int
}

// NewReader constructs a Reader from any io.Reader.
func NewReader(r io.Reader) *Reader {
	i := &Reader{
		raw:            r,
		maxPayloadSize: 2 * 1024 * 1024, // 2MB
	}

	// Prevent double-buffering if the caller already provided a *bufio.Reader
	if br, ok := r.(*bufio.Reader); ok {
		i.src = br
		return i
	}

	i.src = bufio.NewReader(r)

	return i
}

// ReadNext reads the next message from the stream.
//
// ctx cancellation is honoured for both header reads and payload reads.
// For net.Conn sources, cancellation is implemented by setting a past read
// deadline on the connection (via context.AfterFunc), which causes the
// blocking Read to return immediately. For other io.Reader types that
// implement io.Closer, Close is called. For bare io.Readers that support
// neither, the read blocks until the underlying source returns.
//
// Returns io.EOF when the stream closes cleanly between messages.
// Returns ErrStreamDesynced if a previous read was interrupted mid-stream.
func (r *Reader) ReadNext(ctx context.Context) (Message, error) {
	if r.fatalErr != nil {
		return Message{}, fmt.Errorf("%w: %w", ErrStreamDesynced, r.fatalErr)
	}

	msg, err := r.doRead(ctx)
	if err != nil {
		// EOF between messages is clean — do not poison the Reader.
		if err != io.EOF {
			r.fatalErr = err
		}
		return Message{}, err
	}

	return msg, nil
}

// Itterate returns a push iterator over the messages in the stream.
// It calls ReadNext repeatedly and yields each message until an error occurs.
// The caller should check the error value on each iteration; io.EOF indicates
// the stream closed cleanly, while any other error signals a read failure.
//
// Example:
//
//	for msg, err := range r.Itterate(ctx) {
//		if err != nil {
//			// handle error
//			break
//		}
//		// use msg
//	}
func (r *Reader) Itterate(ctx context.Context) iter.Seq2[Message, error] {
	return func(yield func(Message, error) bool) {
		for {
			msg, err := r.ReadNext(ctx)
			if !yield(msg, err) {
				return
			}
			if err != nil {
				return
			}
		}
	}
}

func (r *Reader) doRead(ctx context.Context) (Message, error) {
	if err := ctx.Err(); err != nil {
		return Message{}, err
	}

	// Only install the cancellation hook when the context can actually be
	// cancelled. context.Background() and context.TODO() return a nil Done
	// channel, so AfterFunc would allocate a goroutine-tracking structure
	// that is never needed.
	if ctx.Done() != nil {
		stop := context.AfterFunc(ctx, func() {
			if ts, ok := r.raw.(timeoutSetter); ok {
				ts.SetReadDeadline(time.Now())
				return
			}
			if cl, ok := r.raw.(io.Closer); ok {
				cl.Close()
			}
		})
		defer stop()
	}

	lineBytes, err := r.readLogicalLine()
	if err != nil {
		return Message{}, r.mapErr(ctx, err)
	}

	msgType, args, flags, payloadLength, err := r.parseLine(string(lineBytes))
	if err != nil {
		return Message{}, err
	}

	m := Message{
		verb:  msgType,
		args:  args,
		flags: flags,
		raw:   lineBytes,
	}

	if payloadLength > 0 {
		if r.maxPayloadSize > 0 && payloadLength > r.maxPayloadSize {
			return Message{}, fmt.Errorf("%w: declared %d, max %d", ErrPayloadTooLarge, payloadLength, r.maxPayloadSize)
		}
		payload := make([]byte, payloadLength)
		if _, err := io.ReadFull(r.src, payload); err != nil {
			if err == io.EOF {
				err = io.ErrUnexpectedEOF // Force truncation to be treated as an error
			}
			return Message{}, r.mapErr(ctx, err)
		}
		m.payload = payload
	}

	return m, nil
}

// mapErr translates a read error into a context error when the context was
// responsible for interrupting the read.
func (r *Reader) mapErr(ctx context.Context, err error) error {
	if ctxErr := ctx.Err(); ctxErr != nil {
		return ctxErr
	}
	return err
}

// readLogicalLine reads from r.src until an unescaped newline, joining physical
// lines that end with a trailing backslash (\).
// The returned slice is only valid until the next ReadBytes call on r.src;
// callers that need to retain it must copy.
func (r *Reader) readLogicalLine() ([]byte, error) {
	lineBytes, err := r.src.ReadBytes('\n')
	if err != nil && err != io.EOF {
		return nil, err
	}

	// Trim trailing \r\n.
	lineBytes = trimNewline(lineBytes)

	// Fast path: no line continuation — return the trimmed slice directly.
	if len(lineBytes) == 0 || lineBytes[len(lineBytes)-1] != '\\' {
		if err == io.EOF {
			if len(lineBytes) == 0 {
				return nil, io.EOF
			}
		}
		return bytes.TrimSpace(lineBytes), nil
	}

	// Slow path: one or more physical lines joined by trailing backslash.
	var b bytes.Buffer
	for {
		// Append without the trailing backslash.
		b.Write(lineBytes[:len(lineBytes)-1])

		if err == io.EOF {
			return b.Bytes(), io.ErrUnexpectedEOF
		}

		lineBytes, err = r.src.ReadBytes('\n')
		if err != nil && err != io.EOF {
			return nil, err
		}
		lineBytes = trimNewline(lineBytes)

		if len(lineBytes) == 0 || lineBytes[len(lineBytes)-1] != '\\' {
			b.Write(lineBytes)
			// Any EOF while resolving a continuation is unexpected — the logical
			// line was never terminated.
			if err == io.EOF {
				return b.Bytes(), io.ErrUnexpectedEOF
			}
			break
		}
	}

	return bytes.TrimSpace(b.Bytes()), nil
}

// trimNewline strips a trailing \r and/or \n from b in-place.
func trimNewline(b []byte) []byte {
	for len(b) > 0 && (b[len(b)-1] == '\n' || b[len(b)-1] == '\r') {
		b = b[:len(b)-1]
	}
	return b
}

// parseLine tokenises a logical line into its constituent parts.
// Returns the message type, positional args, flags, payload byte length, and any error.
//
// A payload is declared by ending the header with -- followed by a decimal byte count:
//
//	upload --name=script.sh -- 1024
//
// The count is the exact number of raw bytes that immediately follow the header line.
// No delimiter is used; the receiver reads exactly that many bytes via io.ReadFull.
func (r *Reader) parseLine(line string) (string, []string, FlagSet, int, error) {
	parts, err := splitWords(line)
	if err != nil {
		return "", nil, nil, 0, fmt.Errorf("%w: %w", ErrParseFailure, err)
	}

	if len(parts) == 0 {
		return "", nil, nil, 0, nil
	}

	msgType := parts[0]
	var args []string
	flags := make(FlagSet)
	var payloadLength int

	for i := 1; i < len(parts); i++ {
		part := parts[i]

		if part == "--" {
			// Consume the next token as the byte count.
			if i+1 >= len(parts) {
				return "", nil, nil, 0, ErrMissingPayloadLength
			}
			i++
			n, err := parseInt(parts[i])
			if err != nil || n < 0 {
				return "", nil, nil, 0, fmt.Errorf("%w: %q", ErrInvalidPayloadLength, parts[i])
			}
			payloadLength = n
			continue
		}

		if strings.HasPrefix(part, "-") {
			flagStr := strings.TrimLeft(part, "-")
			kv := strings.SplitN(flagStr, "=", 2)
			key := kv[0]
			val := "true"

			if len(kv) > 1 {
				val = kv[1]
			}

			flags[key] = append(flags[key], val)
			continue
		}

		args = append(args, part)
	}

	return msgType, args, flags, payloadLength, nil
}

// parseInt parses a non-negative decimal integer from s.
func parseInt(s string) (int, error) {
	n := 0
	if len(s) == 0 {
		return 0, ErrInvalidPayloadLength
	}
	for _, c := range s {
		if c < '0' || c > '9' {
			return 0, ErrInvalidPayloadLength
		}
		n = n*10 + int(c-'0')
	}
	return n, nil
}