Update client/crypto to use new helpers

Author: tscolari

client/crypto.go

@@ -22,6 +22,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
 
+	"github.com/dapr/go-sdk/client/internal/crypto"
 	commonv1pb "github.com/dapr/go-sdk/internal/proto/dapr/proto/common/v1"
 	runtimev1pb "github.com/dapr/go-sdk/internal/proto/dapr/proto/runtime/v1"
 )
@@ -48,9 +49,11 @@ func (c *GRPCClient) Encrypt(ctx context.Context, in io.Reader, opts EncryptOpti
 	}
 
 	// Use the context of the stream here.
-	return c.performCryptoOperation(
-		stream.Context(), stream,
-		in, opts,
+	return performCryptoOperation(
+		stream.Context(),
+		stream,
+		in,
+		opts,
 		&runtimev1pb.EncryptRequest{},
 		&runtimev1pb.EncryptResponse{},
 	)
@@ -72,15 +75,25 @@ func (c *GRPCClient) Decrypt(ctx context.Context, in io.Reader, opts DecryptOpti
 	}
 
 	// Use the context of the stream here.
-	return c.performCryptoOperation(
-		stream.Context(), stream,
-		in, opts,
+	return performCryptoOperation(
+		stream.Context(),
+		stream,
+		in,
+		opts,
 		&runtimev1pb.DecryptRequest{},
 		&runtimev1pb.DecryptResponse{},
 	)
 }
 
-func (c *GRPCClient) performCryptoOperation(ctx context.Context, stream grpc.ClientStream, in io.Reader, opts cryptoOperationOpts, reqProto runtimev1pb.CryptoRequests, resProto runtimev1pb.CryptoResponses) (io.Reader, error) {
+func performCryptoOperation[T runtimev1pb.DecryptRequest | runtimev1pb.EncryptRequest, Y runtimev1pb.DecryptResponse | runtimev1pb.EncryptResponse](
+	ctx context.Context,
+	stream grpc.ClientStream,
+	in io.Reader,
+	opts cryptoOperationOpts,
+	reqProto *T,
+	resProto *Y,
+) (io.Reader, error) {
+
 	var err error
 	// Pipe for writing the response
 	pr, pw := io.Pipe()
@@ -110,11 +123,11 @@ func (c *GRPCClient) performCryptoOperation(ctx context.Context, stream grpc.Cli
 
 			// First message only - add the options
 			if optsProto != nil {
-				reqProto.SetOptions(optsProto)
+				crypto.SetOptions(reqProto, optsProto)
 				optsProto = nil
 			} else {
 				// Reset the object so we can re-use it
-				reqProto.Reset()
+				crypto.Reset(reqProto)
 			}
 
 			n, err = in.Read(*reqBuf)
@@ -127,7 +140,7 @@ func (c *GRPCClient) performCryptoOperation(ctx context.Context, stream grpc.Cli
 
 			// Send the chunk if there's anything to send
 			if n > 0 {
-				reqProto.SetPayload(&commonv1pb.StreamPayload{
+				crypto.SetPayload(reqProto, &commonv1pb.StreamPayload{
 					Data: (*reqBuf)[:n],
 					Seq:  seq,
 				})
@@ -184,7 +197,7 @@ func (c *GRPCClient) performCryptoOperation(ctx context.Context, stream grpc.Cli
 			}
 
 			// Write the data, if any, into the pipe
-			payload = resProto.GetPayload()
+			payload = crypto.GetPayload(resProto)
 			if payload != nil {
 				if payload.GetSeq() != expectSeq {
 					pw.CloseWithError(fmt.Errorf("invalid sequence number in chunk: %d (expected: %d)", payload.GetSeq(), expectSeq))
@@ -205,7 +218,7 @@ func (c *GRPCClient) performCryptoOperation(ctx context.Context, stream grpc.Cli
 			}
 
 			// Reset the proto
-			resProto.Reset()
+			crypto.Reset(resProto)
 		}
 
 		// Close the writer of the pipe when done

client/crypto_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 
+	"github.com/dapr/go-sdk/client/internal/crypto"
 	commonv1 "github.com/dapr/go-sdk/internal/proto/dapr/proto/common/v1"
 	runtimev1pb "github.com/dapr/go-sdk/internal/proto/dapr/proto/runtime/v1"
 )
@@ -189,22 +190,24 @@ func TestDecrypt(t *testing.T) {
 /* --- Server methods --- */
 
 func (s *testDaprServer) EncryptAlpha1(stream runtimev1pb.Dapr_EncryptAlpha1Server) error {
-	return s.performCryptoOperation(
+	return testPerformCryptoOperation(
 		stream,
 		&runtimev1pb.EncryptRequest{},
 		&runtimev1pb.EncryptResponse{},
 	)
 }
 
 func (s *testDaprServer) DecryptAlpha1(stream runtimev1pb.Dapr_DecryptAlpha1Server) error {
-	return s.performCryptoOperation(
+	return testPerformCryptoOperation(
 		stream,
 		&runtimev1pb.DecryptRequest{},
 		&runtimev1pb.DecryptResponse{},
 	)
 }
 
-func (s *testDaprServer) performCryptoOperation(stream grpc.ServerStream, reqProto runtimev1pb.CryptoRequests, resProto runtimev1pb.CryptoResponses) error {
+func testPerformCryptoOperation[T runtimev1pb.DecryptRequest | runtimev1pb.EncryptRequest, Y runtimev1pb.DecryptResponse | runtimev1pb.EncryptResponse](
+	stream grpc.ServerStream, reqProto *T, resProto *Y,
+) error {
 	// This doesn't really encrypt or decrypt the data and just sends back whatever it receives
 	pr, pw := io.Pipe()
 
@@ -216,7 +219,7 @@ func (s *testDaprServer) performCryptoOperation(stream grpc.ServerStream, reqPro
 		)
 		first := true
 		for !done && stream.Context().Err() == nil {
-			reqProto.Reset()
+			crypto.Reset(reqProto)
 			err = stream.RecvMsg(reqProto)
 			if errors.Is(err, io.EOF) {
 				done = true
@@ -225,16 +228,16 @@ func (s *testDaprServer) performCryptoOperation(stream grpc.ServerStream, reqPro
 				return
 			}
 
-			if first && !reqProto.HasOptions() {
+			if first && !hasOptions(reqProto) {
 				pw.CloseWithError(errors.New("first message must have options"))
 				return
-			} else if !first && reqProto.HasOptions() {
+			} else if !first && hasOptions(reqProto) {
 				pw.CloseWithError(errors.New("messages after first must not have options"))
 				return
 			}
 			first = false
 
-			payload := reqProto.GetPayload()
+			payload := crypto.GetPayload(reqProto)
 			if payload != nil {
 				if payload.GetSeq() != expectSeq {
 					pw.CloseWithError(fmt.Errorf("invalid sequence number: %d (expected: %d)", payload.GetSeq(), expectSeq))
@@ -261,7 +264,7 @@ func (s *testDaprServer) performCryptoOperation(stream grpc.ServerStream, reqPro
 	)
 	buf := make([]byte, 2<<10)
 	for !done && stream.Context().Err() == nil {
-		resProto.Reset()
+		crypto.Reset(resProto)
 
 		n, err = pr.Read(buf)
 		if errors.Is(err, io.EOF) {
@@ -271,7 +274,7 @@ func (s *testDaprServer) performCryptoOperation(stream grpc.ServerStream, reqPro
 		}
 
 		if n > 0 {
-			resProto.SetPayload(&commonv1.StreamPayload{
+			crypto.SetPayload(resProto, &commonv1.StreamPayload{
 				Seq:  seq,
 				Data: buf[:n],
 			})
@@ -286,3 +289,18 @@ func (s *testDaprServer) performCryptoOperation(stream grpc.ServerStream, reqPro
 
 	return nil
 }
+
+func hasOptions[T runtimev1pb.DecryptRequest | runtimev1pb.EncryptRequest](msg *T) bool {
+	if msg == nil {
+		return false
+	}
+
+	switch r := any(msg).(type) {
+	case *runtimev1pb.EncryptRequest:
+		return r.Options != nil
+	case *runtimev1pb.DecryptRequest:
+		return r.Options != nil
+	}
+
+	return false
+}