COCOS-256 - Progress bar on downloads (ultravioletrs#290)

* add progress bar for downloads

Signed-off-by: Sammy Oina <[email protected]>

* better error handling

Signed-off-by: Sammy Oina <[email protected]>

* fix test and refactor

Signed-off-by: Sammy Oina <[email protected]>

* fix failing test

Signed-off-by: Sammy Oina <[email protected]>

* add test coverage

Signed-off-by: Sammy Oina <[email protected]>

---------

Signed-off-by: Sammy Oina <[email protected]>

.github/workflows/checkproto.yaml

@@ -33,7 +33,7 @@ jobs:
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=27.3
+          PROTOC_VERSION=28.1
           PROTOC_GEN_VERSION=v1.34.2
           PROTOC_GRPC_VERSION=v1.4.0
 

.gitignore

@@ -1,14 +1,13 @@
 build
 
-build
 cmd/manager/img
 
 .cov
 
 *.pem
 
 dist/
-results.zip
+*.zip
 *.spec
 *.tar
 

agent/agent.proto

@@ -11,7 +11,7 @@ service AgentService {
   rpc Algo(stream AlgoRequest) returns (AlgoResponse) {}
   rpc Data(stream DataRequest) returns (DataResponse) {}
   rpc Result(ResultRequest) returns (stream ResultResponse) {}
-  rpc Attestation(AttestationRequest) returns (AttestationResponse) {}
+  rpc Attestation(AttestationRequest) returns (stream AttestationResponse) {}
 }
 
 message AlgoRequest {

agent/api/grpc/server.go

@@ -6,15 +6,20 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"io"
 
 	"github.com/go-kit/kit/transport/grpc"
 	"github.com/ultravioletrs/cocos/agent"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
-const bufferSize = 1024 * 1024
+const (
+	bufferSize  = 1024 * 1024
+	FileSizeKey = "file-size"
+)
 
 var _ agent.AgentServiceServer = (*grpcServer)(nil)
 
@@ -156,12 +161,16 @@ func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_
 	}
 	rr := res.(*agent.ResultResponse)
 
-	reusltBuffer := bytes.NewBuffer(rr.File)
+	if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: fmt.Sprint(len(rr.File))})); err != nil {
+		return status.Error(codes.Internal, err.Error())
+	}
+
+	resultBuffer := bytes.NewBuffer(rr.File)
 
 	buf := make([]byte, bufferSize)
 
 	for {
-		n, err := reusltBuffer.Read(buf)
+		n, err := resultBuffer.Read(buf)
 		if err == io.EOF {
 			break
 		}
@@ -177,11 +186,34 @@ func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_
 	return nil
 }
 
-func (s *grpcServer) Attestation(ctx context.Context, req *agent.AttestationRequest) (*agent.AttestationResponse, error) {
-	_, res, err := s.attestation.ServeGRPC(ctx, req)
+func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.AgentService_AttestationServer) error {
+	_, res, err := s.attestation.ServeGRPC(stream.Context(), req)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	rr := res.(*agent.AttestationResponse)
-	return rr, nil
+
+	if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: fmt.Sprint(len(rr.File))})); err != nil {
+		return status.Error(codes.Internal, err.Error())
+	}
+
+	attestationBuffer := bytes.NewBuffer(rr.File)
+
+	buf := make([]byte, bufferSize)
+
+	for {
+		n, err := attestationBuffer.Read(buf)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
+
+		if err := stream.Send(&agent.AttestationResponse{File: buf[:n]}); err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
+	}
+
+	return nil
 }

agent/api/grpc/server_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/mocks"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 type MockAgentService_AlgoServer struct {
@@ -64,11 +65,34 @@ func (m *MockAgentService_ResultServer) Context() context.Context {
 	return m.ctx
 }
 
+func (m *MockAgentService_ResultServer) SetHeader(metadata.MD) error {
+	return nil
+}
+
 func (m *MockAgentService_ResultServer) Send(resp *agent.ResultResponse) error {
 	args := m.Called(resp)
 	return args.Error(0)
 }
 
+type MockAgentService_AttestationServer struct {
+	grpc.ServerStream
+	mock.Mock
+	ctx context.Context
+}
+
+func (m *MockAgentService_AttestationServer) Context() context.Context {
+	return m.ctx
+}
+
+func (m *MockAgentService_AttestationServer) Send(resp *agent.AttestationResponse) error {
+	args := m.Called(resp)
+	return args.Error(0)
+}
+
+func (m *MockAgentService_AttestationServer) SetHeader(metadata.MD) error {
+	return nil
+}
+
 func TestAlgo(t *testing.T) {
 	mockService := new(mocks.Service)
 	server := NewServer(mockService)
@@ -124,12 +148,14 @@ func TestAttestation(t *testing.T) {
 	mockService := new(mocks.Service)
 	server := NewServer(mockService)
 
+	mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
+	mockStream.On("Send", mock.AnythingOfType("*agent.AttestationResponse")).Return(nil)
+
 	reportData := [agent.ReportDataSize]byte{}
 	mockService.On("Attestation", mock.Anything, reportData).Return([]byte("attestation data"), nil)
 
-	resp, err := server.Attestation(context.Background(), &agent.AttestationRequest{ReportData: reportData[:]})
+	err := server.Attestation(&agent.AttestationRequest{ReportData: reportData[:]}, mockStream)
 	assert.NoError(t, err)
-	assert.Equal(t, []byte("attestation data"), resp.File)
 
 	mockService.AssertExpectations(t)
 }