/* * 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 }