constellation/debugd/internal/filetransfer/chunkstream.go
2023-01-25 09:58:56 +01:00

48 lines
1.0 KiB
Go

/*
Copyright (c) Edgeless Systems GmbH
SPDX-License-Identifier: AGPL-3.0-only
*/
package filetransfer
import (
"errors"
pb "github.com/edgelesssys/constellation/v2/debugd/service"
)
// recvChunkStream is a wrapper around a RecvFilesStream that only returns chunks.
type recvChunkStream struct {
stream RecvFilesStream
}
// Recv receives a FileTransferMessage and returns the chunk.
func (s *recvChunkStream) Recv() (*pb.Chunk, error) {
msg, err := s.stream.Recv()
if err != nil {
return nil, err
}
chunk := msg.GetChunk()
if chunk == nil {
return nil, errors.New("expected chunk")
}
return chunk, nil
}
// sendChunkStream is a wrapper around a SendFilesStream that wraps chunks for every message.
type sendChunkStream struct {
stream SendFilesStream
}
// Send wraps the given chunk in a FileTransferMessage and sends it.
func (s *sendChunkStream) Send(chunk *pb.Chunk) error {
chunkMessage := &pb.FileTransferMessage_Chunk{
Chunk: chunk,
}
message := &pb.FileTransferMessage{
Kind: chunkMessage,
}
return s.stream.Send(message)
}