Refactorings for simplify readability (#35)

internal/controller/rpc.go

@@ -3,6 +3,9 @@ package controller
 import (
 	v1pkg "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/cirruslabs/orchard/rpc"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/protobuf/types/known/emptypb"
+
 	//nolint:staticcheck // https://github.com/mitchellh/go-grpc-net-conn/pull/1
 	"github.com/golang/protobuf/proto"
 	grpc_net_conn "github.com/mitchellh/go-grpc-net-conn"
@@ -10,24 +13,18 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-func (controller *Controller) Watch(stream rpc.Controller_WatchServer) error {
+func (controller *Controller) Watch(_ *emptypb.Empty, stream rpc.Controller_WatchServer) error {
 	if !controller.authorizeGRPC(stream.Context(), v1pkg.ServiceAccountRoleWorker) {
 		return status.Errorf(codes.Unauthenticated, "auth failed")
 	}
 
-	// The first message is always an initialization
-	watchFromWorker, err := stream.Recv()
-	if err != nil {
-		return err
+	workerUIDMetadataValue := metadata.ValueFromIncomingContext(stream.Context(), rpc.MetadataWorkerUIDKey)
+	if len(workerUIDMetadataValue) == 0 {
+		return status.Errorf(codes.InvalidArgument, "no worker UID in metadata")
 	}
 
-	initAction, ok := watchFromWorker.Action.(*rpc.WatchFromWorker_InitAction)
-	if !ok {
-		return status.Errorf(codes.FailedPrecondition, "expected an initialization message")
-	}
-
-	// Subscribe to rendez-vous requests from the API for this worker
-	workerCh, cancel := controller.watcher.Subscribe(stream.Context(), initAction.InitAction.WorkerUid)
+	workerUID := workerUIDMetadataValue[0]
+	workerCh, cancel := controller.workerNotifier.Register(stream.Context(), workerUID)
 	defer cancel()
 
 	for {
@@ -47,39 +44,25 @@ func (controller *Controller) PortForward(stream rpc.Controller_PortForwardServe
 		return status.Errorf(codes.Unauthenticated, "auth failed")
 	}
 
-	// The first message is always an initialization
-	portForwardFromWorker, err := stream.Recv()
-	if err != nil {
-		return err
+	sessionMetadataValue := metadata.ValueFromIncomingContext(stream.Context(), rpc.MetadataWorkerPortForwardingSessionKey)
+	if len(sessionMetadataValue) == 0 {
+		return status.Errorf(codes.InvalidArgument, "no session in metadata")
 	}
 
-	initAction, ok := portForwardFromWorker.Action.(*rpc.PortForwardFromWorker_InitAction)
-	if !ok {
-		return status.Errorf(codes.FailedPrecondition, "expected an initialization message")
-	}
-
-	// Satisfy the rendez-vous request
 	conn := &grpc_net_conn.Conn{
-		Stream: stream,
-		Request: &rpc.PortForwardFromController{
-			Action: &rpc.PortForwardFromController_DataAction{
-				DataAction: &rpc.PortForwardFromController_Data{},
-			},
-		},
-		Response: &rpc.PortForwardFromWorker{
-			Action: &rpc.PortForwardFromWorker_DataAction{
-				DataAction: &rpc.PortForwardFromWorker_Data{},
-			},
-		},
+		Stream:   stream,
+		Request:  &rpc.PortForwardData{},
+		Response: &rpc.PortForwardData{},
 		Encode: grpc_net_conn.SimpleEncoder(func(message proto.Message) *[]byte {
-			return &message.(*rpc.PortForwardFromController).Action.(*rpc.PortForwardFromController_DataAction).DataAction.Data
+			return &message.(*rpc.PortForwardData).Data
 		}),
 		Decode: grpc_net_conn.SimpleDecoder(func(message proto.Message) *[]byte {
-			return &message.(*rpc.PortForwardFromWorker).Action.(*rpc.PortForwardFromWorker_DataAction).DataAction.Data
+			return &message.(*rpc.PortForwardData).Data
 		}),
 	}
 
-	proxyCtx, err := controller.proxy.Respond(initAction.InitAction.Token, conn)
+	// make proxy aware of the connection
+	proxyCtx, err := controller.proxy.Respond(sessionMetadataValue[0], conn)
 	if err != nil {
 		return err
 	}