internal/pkg/rsmsg/serialization.go

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the \"License\");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an \"AS IS\" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package rsmsg provides serialization for ResultStore messages.
package rsmsg

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"sync"

	"google.golang.org/protobuf/proto"

	rspb "google.golang.org/genproto/googleapis/devtools/resultstore/v2"
)

// UploadRequestsToBytes serializes a slice of UploadRequests into a byte slice,
// with each message prefixed by its length.
func UploadRequestsToBytes(reqs []*rspb.UploadRequest) ([]byte, error) {
	reqChan := make(chan *rspb.UploadRequest)

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for _, req := range reqs {
			reqChan <- req
		}
		close(reqChan)
	}()
	defer wg.Wait()

	var buf bytes.Buffer
	if err := WriteUploadRequests(&buf, reqChan); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// BytesToUploadRequests deserializes a byte slice into a slice of UploadRequests.
func BytesToUploadRequests(ctx context.Context, data []byte) ([]*rspb.UploadRequest, error) {
	gotReqs := make([]*rspb.UploadRequest, 0)
	if len(data) == 0 {
		return gotReqs, nil
	}
	reqChan := make(chan *rspb.UploadRequest)
	errChan := make(chan error, 1)

	buf := bytes.NewBuffer(data)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		errChan <- ReadUploadRequests(ctx, buf, reqChan)
	}()

	for req := range reqChan {
		gotReqs = append(gotReqs, req)
	}

	wg.Wait()
	if err := <-errChan; err != nil {
		return nil, err
	}

	return gotReqs, nil
}

// WriteUploadRequests writes a stream of UploadRequests from a channel to a writer,
// prefixing each message with its length.
func WriteUploadRequests(output io.Writer, reqs <-chan *rspb.UploadRequest) error {
	for req := range reqs {
		if err := WriteUploadRequest(output, req); err != nil {
			return err
		}
	}
	return nil
}

// WriteUploadRequest writes a single UploadRequest to a writer,
// prefixing the message with its length.
func WriteUploadRequest(output io.Writer, req *rspb.UploadRequest) error {
	wireBytes, err := proto.Marshal(req)
	if err != nil {
		return err
	}
	if err := binary.Write(output, binary.LittleEndian, uint32(len(wireBytes))); err != nil {
		return err
	}
	if _, err := output.Write(wireBytes); err != nil {
		return err
	}
	return nil
}

// ReadUploadRequest reads a single length-prefixed UploadRequest from a reader.
func ReadUploadRequest(input io.Reader) (*rspb.UploadRequest, error) {
	var messageLen uint32
	if err := binary.Read(input, binary.LittleEndian, &messageLen); err != nil {
		return nil, err
	}
	messagePayload := make([]byte, messageLen)
	if _, err := io.ReadFull(input, messagePayload); err != nil {
		return nil, err
	}
	req := &rspb.UploadRequest{}
	if err := proto.Unmarshal(messagePayload, req); err != nil {
		return nil, err
	}
	return req, nil
}

// ReadUploadRequests reads length-prefixed UploadRequests from a reader and sends
// them to a channel. It returns if the reader is exhausted, an error occurs,
// or the context is cancelled.
func ReadUploadRequests(ctx context.Context, input io.Reader, reqs chan<- *rspb.UploadRequest) error {
	defer close(reqs)
	for {
		req, err := ReadUploadRequest(input)
		if err == io.EOF {
			break
		}
		if err != nil {
			return fmt.Errorf("Error reading message: %v", err)
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case reqs <- req:
		}
	}
	return nil
}