mcg/internal/client/reader.go

package client

import (
	"fmt"
	"net/textproto"
	"strconv"
	"strings"
)

type reader interface {
	ReadData(client *Client, args []string) (bool, []string, error)
}

type noopReader struct{}

func (r *noopReader) ReadData(client *Client, args []string) (bool, []string, error) {
	return true, args, nil
}

type emptyReader struct{}

func (r *emptyReader) ReadData(client *Client, args []string) (bool, []string, error) {
	_, err := client.read()

	return true, args, err
}

type listReader struct {
	key  string
	data []string
}

func (r *listReader) ReadData(client *Client, args []string) (bool, []string, error) {
	data, err := client.read()
	if err != nil {
		return true, args, err
	}

	r.data = []string{}
	for _, line := range data {
		if !strings.HasPrefix(line, r.key+": ") {
			return true, args, textproto.ProtocolError(fmt.Sprintf("Unexpected line: %s", line))
		}

		r.data = append(r.data, line[len(r.key)+2:])
	}

	return true, args, nil
}

func (r *listReader) List() []string {
	return r.data
}

type mapReader struct {
	data map[string]string
}

func (r *mapReader) ReadData(client *Client, args []string) (bool, []string, error) {
	data, err := client.read()
	if err != nil {
		return true, args, err
	}

	r.data = make(map[string]string)
	for _, line := range data {
		parts := strings.Split(line, ": ")
		if len(parts) < 2 {
			return true, args, textproto.ProtocolError(fmt.Sprintf("Unexpected line: %s", line))
		}
		r.data[parts[0]] = strings.Join(parts[1:], ": ")
	}

	return true, args, nil
}

func (r *mapReader) Dict() map[string]string {
	return r.data
}

type listMapReader struct {
	Delimiter string
	data      []map[string]string
}

func (r *listMapReader) ReadData(client *Client, args []string) (bool, []string, error) {
	data, err := client.read()
	if err != nil {
		return true, args, err
	}

	r.data = nil
	var entry map[string]string
	for _, line := range data {
		parts := strings.Split(line, ": ")
		if len(parts) < 2 {
			return true, args, textproto.ProtocolError(fmt.Sprintf("Unexpected line: %s", line))
		}

		if parts[0] == r.Delimiter {
			if entry != nil {
				r.data = append(r.data, entry)
			}
			entry = make(map[string]string)
		}
		entry[parts[0]] = strings.Join(parts[1:], ": ")
	}
	if entry != nil {
		r.data = append(r.data, entry)
	}

	return true, args, nil
}

func (r *listMapReader) DictList() []map[string]string {
	return r.data
}

type binaryReader struct {
	data   []byte
	size   int
	offset int
}

func (r *binaryReader) ReadData(client *Client, args []string) (bool, []string, error) {
	firstLine, dataComplete, firstLineErr := r.readFirstLine(client)
	if firstLineErr != nil {
		return true, args, firstLineErr
	}
	if dataComplete {
		return true, args, nil
	}

	totalSize, sizeErr := r.readSize(firstLine)
	if sizeErr != nil {
		return true, args, sizeErr
	}
	r.size = totalSize

	currentSize, bytesCountErr := r.readBytesCount(client)
	if bytesCountErr != nil {
		return true, args, bytesCountErr
	}

	r.initData(totalSize)
	bytesErr := r.readBytes(client, currentSize)
	if bytesErr != nil {
		return true, args, bytesErr
	}

	_, lineBreakErr := client.readLine()
	if lineBreakErr != nil {
		return true, args, lineBreakErr
	}
	_, lineCompletionErr := client.readLine()
	if lineCompletionErr != nil {
		return true, args, lineCompletionErr
	}

	complete := (r.offset >= r.size)
	newArgs := r.updateOffsetArgument(args)

	return complete, newArgs, nil
}

func (r *binaryReader) readFirstLine(client *Client) (line string, complete bool, err error) {
	line, lineErr := client.readLine()
	if lineErr != nil {
		return "", true, lineErr
	}

	// OK here means the binary data is complete
	if line == "OK" {
		return "", true, nil
	}

	return line, false, nil
}

func (r *binaryReader) readSize(line string) (int, error) {
	if !strings.HasPrefix(line, "size: ") {
		return 0, textproto.ProtocolError(fmt.Sprintf("Unexpected line: %s", line))
	}

	parsedSize, err := strconv.Atoi(line[6:])
	if err != nil {
		return 0, textproto.ProtocolError(fmt.Sprintf("Invalid size: %s", line[6:]))
	}

	return parsedSize, nil
}

func (r *binaryReader) readBytesCount(client *Client) (int, error) {
	line, lineErr := client.readLine()
	if lineErr != nil {
		return 0, lineErr
	}

	if !strings.HasPrefix(line, "binary: ") {
		return 0, textproto.ProtocolError(fmt.Sprintf("Unexpected line: %s", line))
	}

	count, countErr := strconv.Atoi(line[8:])
	if countErr != nil {
		return 0, textproto.ProtocolError(fmt.Sprintf("Invalid binary: %s", line[8:]))
	}

	return count, nil
}

func (r *binaryReader) initData(size int) {
	if r.data == nil {
		r.data = make([]byte, size)
	}
}

func (r *binaryReader) readBytes(client *Client, size int) error {
	bytesError := client.readBytes(r.data[r.offset : r.offset+size])
	if bytesError != nil {
		return bytesError
	}

	r.offset += size

	return nil
}

func (r *binaryReader) updateOffsetArgument(args []string) []string {
	args[len(args)-1] = strconv.Itoa(r.offset)

	return args
}

func (r *binaryReader) Binary() []byte {
	return r.data
}