ygggo/internal/remote/io.go

/*
 *  Copyright (c) 2026 Kaiy Ragur
 *
 *  This Source Code Form is subject to the terms of the Mozilla Public
 *  License, v. 2.0. If a copy of the MPL was not distributed with this
 *  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */
package remote

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"sync"

	pb "github.com/neilalexander/yggmail/internal/remote/types"
	"google.golang.org/protobuf/proto"
)

var MagicHeader = []byte{0x59, 0x47, 0x47, 0x21} // "YGG!"
var MemoryLimit uint32 = 1024 * 1024 * 2         // 2 Mb

type Payload struct {
	R    io.ReadCloser
	Size uint32
}

func (p *Payload) Close() error {
	if p == nil || p.R == nil {
		return nil
	}
	return p.R.Close()
}

type Tempfile struct {
	io.ReadSeeker
	file *os.File
	size uint32
}

func (f *Tempfile) Len() int64 {
	return int64(f.size)
}

func (f *Tempfile) Close() error {
	if f.file != nil {
		f.file.Close()
		return os.Remove(f.file.Name())
	}
	return nil
}

func ReadRawData(reader io.Reader) (io.ReadCloser, int, error) {
	peekBuf := new(bytes.Buffer)
	peekBuf.Grow(int(MemoryLimit))

	limited := io.LimitReader(reader, int64(MemoryLimit)+1)
	written, err := peekBuf.ReadFrom(limited)
	if err != nil && err != io.EOF {
		return nil, 0, fmt.Errorf("failed to read data: %w", err)
	}
	if written < int64(MemoryLimit) {
		return &Tempfile{
			ReadSeeker: bytes.NewReader(peekBuf.Bytes()),
			size:       uint32(peekBuf.Len())}, peekBuf.Len(), nil
	}

	// Sew data that we readed and not-read
	reader = io.MultiReader(peekBuf, reader)
	// Read large mail to temp file
	payload, err := os.CreateTemp("", "yggmail-*")
	if err != nil {
		return nil, 0, fmt.Errorf("failed to create temp-data file: %w", err)
	}
	size, err := io.Copy(payload, reader)
	payload.Seek(0, 0)

	return &Tempfile{payload, payload, uint32(size)}, int(size), nil
}

// func ReadRawData(r io.Reader, rawLen uint32) (io.ReadCloser, error) {
// 	if rawLen == 0 {
// 		return nil, nil
// 	}
//
// 	if rawLen <= MemoryLimit {
// 		// Write in memory
// 		buf := bytes.NewBuffer(make([]byte, 0, rawLen))
// 		_, err := io.CopyN(buf, r, int64(rawLen))
// 		return io.NopCloser(buf), err
// 	} else {
// 		// Write to tmp
// 		tmp, _ := os.CreateTemp("", "yggmail-payload-*")
// 		_, err := io.CopyN(tmp, r, int64(rawLen))
// 		tmp.Seek(0, 0) // Return in ahead
// 		return &Tempfile{File: tmp}, err
// 	}
// }

//
// Frame = Magic + uint32(Len Protobuff Meta) + uint32(Len Payload) + Protobuff + Payload
//
// Magic Header
//
// len(Protobuff)
// len()
//
// Protobuff Metadata
// Payload 	 Rawdata

// # SendFrame
//
// Send Magic + uint32(Len Protobuff Meta) + uint32(Len Payload) + Protobuff + Payload
func readWithMagic(br *bufio.Reader, maxJunkBytes int) error {
	magicIdx := 0

	for range maxJunkBytes {
		b, err := br.ReadByte()
		if err != nil {
			return err
		}

		if b == MagicHeader[magicIdx] {
			magicIdx++
			if magicIdx == len(MagicHeader) {
				return nil // Found
			}
		} else {
			if b == MagicHeader[0] {
				magicIdx = 1
			} else {
				magicIdx = 0
			}
		}
	}
	return fmt.Errorf("sync failed: magic header not found")
}

// # ReadFrame
//
// Magic + uint32(Len Protobuff Meta) + uint32(Len Payload) + Protobuff + Payload
//
// Return Response / Request Reader
//
// RawData Reader

var marshalPool = sync.Pool{
	New: func() interface{} {
		return make([]byte, 0, 1024*64) // Начинаем с 64КБ
	},
}

type limitedReadCloser struct {
	io.Reader
	io.Closer
}

func ReadFrame(br *bufio.Reader, msg proto.Message) (proto.Message, *Payload, error) {
	if err := readWithMagic(br, 1024); err != nil {
		return nil, nil, err
	}

	var lengths [8]byte
	if _, err := io.ReadFull(br, lengths[:]); err != nil {
		return nil, nil, err
	}

	protoLen := binary.BigEndian.Uint32(lengths[0:4])
	rawLen := binary.BigEndian.Uint32(lengths[4:8])

	protoData := make([]byte, protoLen)
	if _, err := io.ReadFull(br, protoData); err != nil {
		return nil, nil, err
	}

	if err := proto.Unmarshal(protoData, msg); err != nil {
		return nil, nil, err
	}

	var payload *Payload
	if rawLen > 0 {
		payload = &Payload{
			Size: rawLen,
		}

		lr := io.LimitReader(br, int64(rawLen))
		if rc, ok := lr.(io.ReadCloser); ok {
			payload.R = &limitedReadCloser{lr, rc}
		} else {
			payload.R = io.NopCloser(lr)
		}

	}
	return msg, payload, nil
}

func SendFrame(w io.Writer, msg proto.Message, payload *Payload) error {
	if payload != nil && payload.R != nil {
		defer payload.R.Close()
	}

	buf := marshalPool.Get().([]byte)
	buf = buf[:0]

	opts := proto.MarshalOptions{}
	data, err := opts.MarshalAppend(buf, msg)
	if err != nil {
		marshalPool.Put(buf)
		return err
	}
	defer marshalPool.Put(data)

	pSize := uint32(0)
	if payload != nil {
		pSize = payload.Size
	}

	var header [12]byte
	copy(header[0:4], MagicHeader)
	binary.BigEndian.PutUint32(header[4:8], uint32(len(data)))
	binary.BigEndian.PutUint32(header[8:12], pSize)

	if _, err := w.Write(header[:]); err != nil {
		return err
	}
	if _, err := w.Write(data); err != nil {
		return err
	}

	if payload != nil && payload.Size > 0 {
		if _, err := io.CopyN(w, payload.R, int64(pSize)); err != nil {
			return err
		}
	}
	return nil
}

func ReadRequest(r *bufio.Reader) (*pb.Request, *Payload, error) {
	req := &pb.Request{}
	_, payload, err := ReadFrame(r, req)
	if err != nil {
		return nil, nil, fmt.Errorf("ReadRequest: %w", err)
	}

	if payload == nil {
		return req, nil, nil
	}
	payload.R, _, err = ReadRawData(payload.R)
	if err != nil {
		return nil, nil, fmt.Errorf("Error read payload")
	}
	return req, payload, nil
}

func ReadResponse(r *bufio.Reader) (*pb.Response, *Payload, error) {
	resp := &pb.Response{}
	_, payload, err := ReadFrame(r, resp)
	if err != nil {
		return nil, nil, fmt.Errorf("ReadRequest: %w", err)
	}

	// Unmarshal Protobuff metadata
	// Read File
	if payload == nil {
		return resp, nil, nil
	}
	payload.R, _, err = ReadRawData(payload.R)
	if err != nil {
		return nil, nil, fmt.Errorf("Error read payload")
	}
	return resp, payload, nil
}