adding async recording encryption with gRPC multipart uploader

Author: eriktate

api/client/client.go

@@ -24,6 +24,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"iter"
 	"log/slog"
 	"net"
 	"slices"
@@ -88,6 +89,7 @@ import (
 	oktapb "github.com/gravitational/teleport/api/gen/proto/go/teleport/okta/v1"
 	pluginspb "github.com/gravitational/teleport/api/gen/proto/go/teleport/plugins/v1"
 	presencepb "github.com/gravitational/teleport/api/gen/proto/go/teleport/presence/v1"
+	recordingencryptionv1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/recordingencryption/v1"
 	resourceusagepb "github.com/gravitational/teleport/api/gen/proto/go/teleport/resourceusage/v1"
 	samlidppb "github.com/gravitational/teleport/api/gen/proto/go/teleport/samlidp/v1"
 	secreportsv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/secreports/v1"
@@ -125,6 +127,7 @@ type AuthServiceClient struct {
 	auditlogpb.AuditLogServiceClient
 	userpreferencespb.UserPreferencesServiceClient
 	notificationsv1pb.NotificationServiceClient
+	recordingencryptionv1pb.RecordingEncryptionServiceClient
 }
 
 // Client is a gRPC Client that connects to a Teleport Auth server either
@@ -531,10 +534,11 @@ func (c *Client) dialGRPC(ctx context.Context, addr string) error {
 
 	c.conn = conn
 	c.grpc = AuthServiceClient{
-		AuthServiceClient:            proto.NewAuthServiceClient(c.conn),
-		AuditLogServiceClient:        auditlogpb.NewAuditLogServiceClient(c.conn),
-		UserPreferencesServiceClient: userpreferencespb.NewUserPreferencesServiceClient(c.conn),
-		NotificationServiceClient:    notificationsv1pb.NewNotificationServiceClient(c.conn),
+		AuthServiceClient:                proto.NewAuthServiceClient(c.conn),
+		AuditLogServiceClient:            auditlogpb.NewAuditLogServiceClient(c.conn),
+		UserPreferencesServiceClient:     userpreferencespb.NewUserPreferencesServiceClient(c.conn),
+		NotificationServiceClient:        notificationsv1pb.NewNotificationServiceClient(c.conn),
+		RecordingEncryptionServiceClient: recordingencryptionv1pb.NewRecordingEncryptionServiceClient(c.conn),
 	}
 	c.JoinServiceClient = NewJoinServiceClient(proto.NewJoinServiceClient(c.conn))
 
@@ -2454,6 +2458,47 @@ func (c *Client) StreamSessionEvents(ctx context.Context, sessionID string, star
 	return ch, e
 }
 
+// UploadEncryptedRecording streams encrypted recording parts to the auth
+// server to be saved in long term storage.
+func (c *Client) UploadEncryptedRecording(ctx context.Context, sessionID string, parts iter.Seq2[[]byte, error]) error {
+	createRes, err := c.grpc.CreateUpload(ctx, &recordingencryptionv1pb.CreateUploadRequest{
+		SessionId: sessionID,
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	var uploadedParts []*recordingencryptionv1pb.UploadPartResponse
+	var partNumber int64
+	for part, err := range parts {
+		defer func() {
+			partNumber++
+		}()
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		uploadRes, err := c.grpc.UploadPart(ctx, &recordingencryptionv1pb.UploadPartRequest{
+			Upload:     createRes.Upload,
+			PartNumber: partNumber,
+			Part:       part,
+		})
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		uploadedParts = append(uploadedParts, uploadRes)
+	}
+
+	if _, err := c.grpc.CompleteUpload(ctx, &recordingencryptionv1pb.CompleteUploadRequest{
+		Upload: createRes.Upload,
+		Parts:  uploadedParts,
+	}); err != nil {
+		return trace.Wrap(err)
+	}
+
+	return nil
+}
+
 // SearchEvents allows searching for events with a full pagination support.
 func (c *Client) SearchEvents(ctx context.Context, fromUTC, toUTC time.Time, namespace string, eventTypes []string, limit int, order types.EventOrder, startKey string) ([]events.AuditEvent, string, error) {
 	request := &proto.GetEventsRequest{

constants.go

@@ -304,6 +304,9 @@ const (
 	// ComponentMCP represents the MCP server handler.
 	ComponentMCP = "mcp"
 
+	// ComponentRecordingEncryption represents recording encryption
+	ComponentRecordingEncryption = "recording-encryption"
+
 	// VerboseLogsEnvVar forces all logs to be verbose (down to DEBUG level)
 	VerboseLogsEnvVar = "TELEPORT_DEBUG"
 

lib/auth/authclient/clt.go

@@ -406,6 +406,10 @@ func (c *Client) ListWindowsDesktopServices(ctx context.Context, req types.ListW
 	return nil, trace.NotImplemented(notImplementedMessage)
 }
 
+func (c *Client) GetMultipartUploader() events.MultipartUploader {
+	return nil
+}
+
 const (
 	// UserTokenTypeResetPasswordInvite is a token type used for the UI invite flow that
 	// allows users to change their password and set second factor (if enabled).

lib/auth/grpcserver.go

@@ -68,6 +68,7 @@ import (
 	mfav1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/mfa/v1"
 	notificationsv1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/notifications/v1"
 	presencev1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/presence/v1"
+	recordingencryptionv1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/recordingencryption/v1"
 	scopedaccessv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/scopes/access/v1"
 	scopedjoiningv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/scopes/joining/v1"
 	secreportsv1pb "github.com/gravitational/teleport/api/gen/proto/go/teleport/secreports/v1"
@@ -106,6 +107,7 @@ import (
 	"github.com/gravitational/teleport/lib/auth/machineid/workloadidentityv1"
 	"github.com/gravitational/teleport/lib/auth/notifications/notificationsv1"
 	"github.com/gravitational/teleport/lib/auth/presence/presencev1"
+	"github.com/gravitational/teleport/lib/auth/recordingencryption/recordingencryptionv1"
 	scopedaccess "github.com/gravitational/teleport/lib/auth/scopes/access"
 	scopedjoining "github.com/gravitational/teleport/lib/auth/scopes/joining"
 	"github.com/gravitational/teleport/lib/auth/secreports/secreportsv1"
@@ -5603,6 +5605,12 @@ func NewGRPCServer(cfg GRPCServerConfig) (*GRPCServer, error) {
 	}
 	userloginstatev1pb.RegisterUserLoginStateServiceServer(server, userLoginStateServer)
 
+	recordingEncryptionService, err := recordingencryptionv1.NewService(recordingencryptionv1.ServiceConfig{
+		Uploader: cfg.AuditLog.GetMultipartUploader(),
+		Logger:   cfg.AuthServer.logger.With(teleport.ComponentKey, teleport.ComponentRecordingEncryption),
+	})
+	recordingencryptionv1pb.RegisterRecordingEncryptionServiceServer(server, recordingEncryptionService)
+
 	clusterConfigService, err := clusterconfigv1.NewService(clusterconfigv1.ServiceConfig{
 		Cache:      cfg.AuthServer.Cache,
 		Backend:    cfg.AuthServer.Services,

lib/auth/init.go

@@ -390,6 +390,9 @@ type InitConfig struct {
 
 	// VnetConfigService manages the VNet config resource.
 	VnetConfigService services.VnetConfigService
+
+	// MultipartUploader handles multipart uploads.
+	MultipartUploader events.MultipartUploader
 }
 
 // Init instantiates and configures an instance of AuthServer

lib/auth/recordingencryption/recordingencryptionv1/service.go

@@ -0,0 +1,131 @@
+package recordingencryptionv1
+
+import (
+	"bytes"
+	"context"
+	"log/slog"
+
+	"github.com/gravitational/trace"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"github.com/gravitational/teleport"
+
+	"github.com/gravitational/teleport/lib/events"
+	"github.com/gravitational/teleport/lib/session"
+
+	recordingencryptionv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/recordingencryption/v1"
+)
+
+// ServiceConfig captures everything a [Service] requires to fulfill requests.
+type ServiceConfig struct {
+	Logger   *slog.Logger
+	Uploader events.MultipartUploader
+}
+
+// NewService returns a new [Service] based on the given [ServiceConfig].
+func NewService(cfg ServiceConfig) (*Service, error) {
+	if cfg.Logger == nil {
+		cfg.Logger = slog.Default()
+	}
+
+	return &Service{
+		logger:   cfg.Logger.With("component", teleport.ComponentRecordingEncryption),
+		uploader: cfg.Uploader,
+	}, nil
+}
+
+// Service implements a gRPC server for interacting with encrypted recordings.
+type Service struct {
+	recordingencryptionv1.UnimplementedRecordingEncryptionServiceServer
+
+	logger   *slog.Logger
+	uploader events.MultipartUploader
+}
+
+func streamUploadAsProto(upload events.StreamUpload) *recordingencryptionv1.Upload {
+	return &recordingencryptionv1.Upload{
+		UploadId:    upload.ID,
+		SessionId:   upload.SessionID.String(),
+		InitiatedAt: timestamppb.New(upload.Initiated),
+	}
+}
+
+func protoAsStreamUpload(upload *recordingencryptionv1.Upload) (events.StreamUpload, error) {
+	sessionID, err := session.ParseID(upload.SessionId)
+	if err != nil {
+		return events.StreamUpload{}, trace.BadParameter("invalid session ID", err)
+	}
+
+	return events.StreamUpload{
+		ID:        upload.UploadId,
+		SessionID: *sessionID,
+		Initiated: upload.InitiatedAt.AsTime(),
+	}, nil
+}
+
+func protoAsStreamPart(part *recordingencryptionv1.UploadPartResponse) events.StreamPart {
+	return events.StreamPart{
+		Number:       part.PartNumber,
+		ETag:         part.ETag,
+		LastModified: part.LastModified.AsTime(),
+	}
+}
+
+// CreateUpload begins a multipart upload for an encrypted session recording.
+func (s *Service) CreateUpload(ctx context.Context, req *recordingencryptionv1.CreateUploadRequest) (*recordingencryptionv1.CreateUploadResponse, error) {
+	sessionID, err := session.ParseID(req.SessionId)
+	if err != nil {
+		return nil, trace.BadParameter("invalid session ID", err)
+	}
+
+	upload, err := s.uploader.CreateUpload(ctx, *sessionID)
+	if err != nil {
+		return nil, trace.Wrap(err, "creating encrypted recording upload")
+	}
+
+	return &recordingencryptionv1.CreateUploadResponse{
+		Upload: streamUploadAsProto(*upload),
+	}, nil
+}
+
+// UploadPart uploads an encrypted session recording part to the given upload ID.
+func (s *Service) UploadPart(ctx context.Context, req *recordingencryptionv1.UploadPartRequest) (*recordingencryptionv1.UploadPartResponse, error) {
+	upload, err := protoAsStreamUpload(req.Upload)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if err := s.uploader.ReserveUploadPart(ctx, upload, req.PartNumber); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	part := bytes.NewReader(req.Part)
+	streamPart, err := s.uploader.UploadPart(ctx, upload, req.PartNumber, part)
+	if err != nil {
+		return nil, trace.Wrap(err, "uploading encrypted recording part")
+	}
+
+	return &recordingencryptionv1.UploadPartResponse{
+		PartNumber: streamPart.Number,
+		ETag:       streamPart.ETag,
+	}, nil
+}
+
+// CompleteUpload marks a given encrypted session upload as complete.
+func (s *Service) CompleteUpload(ctx context.Context, req *recordingencryptionv1.CompleteUploadRequest) (*recordingencryptionv1.CompleteUploadResponse, error) {
+	upload, err := protoAsStreamUpload(req.Upload)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	parts := make([]events.StreamPart, len(req.Parts))
+	for idx, part := range req.Parts {
+		parts[idx] = protoAsStreamPart(part)
+	}
+
+	if err := s.uploader.CompleteUpload(ctx, upload, parts); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return &recordingencryptionv1.CompleteUploadResponse{}, nil
+}

lib/events/api.go

@@ -22,6 +22,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"iter"
 	"math"
 	"time"
 
@@ -1107,6 +1108,8 @@ type StreamEmitter interface {
 type AuditLogSessionStreamer interface {
 	AuditLogger
 	SessionStreamer
+	EncryptedRecordingUploader
+	GetMultipartUploader() MultipartUploader
 }
 
 // SessionStreamer supports streaming session chunks or events.
@@ -1121,6 +1124,12 @@ type SessionStreamer interface {
 	StreamSessionEvents(ctx context.Context, sessionID session.ID, startIndex int64) (chan apievents.AuditEvent, chan error)
 }
 
+// EncryptedRecordingUploader takes a session ID and a sequence of encrypted
+// recording parts and uploads an encrypted session recording.
+type EncryptedRecordingUploader interface {
+	UploadEncryptedRecording(ctx context.Context, sessionID string, parts iter.Seq2[[]byte, error]) error
+}
+
 type SearchEventsRequest struct {
 	// From is oldest date of returned events, can be zero.
 	From time.Time

lib/events/auditlog.go

@@ -19,10 +19,12 @@
 package events
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"io"
 	"io/fs"
+	"iter"
 	"log/slog"
 	"os"
 	"path/filepath"
@@ -635,6 +637,55 @@ func (l *AuditLog) StreamSessionEvents(ctx context.Context, sessionID session.ID
 	return c, e
 }
 
+// UploadEncryptedRecording uploads recording parts provided over a channel using the
+// configured MultipartUploader.
+func UploadEncryptedRecording(ctx context.Context, uploader MultipartUploader, sessionID string, parts iter.Seq2[[]byte, error]) error {
+	sessID, err := session.ParseID(sessionID)
+	if err != nil {
+		return trace.BadParameter("invalid session ID", err)
+	}
+	upload, err := uploader.CreateUpload(ctx, *sessID)
+	if err != nil {
+		return trace.Wrap(err, "creating upload")
+	}
+
+	var streamParts []StreamPart
+	var partNumber int64
+	for part, err := range parts {
+		defer func() {
+			partNumber++
+		}()
+
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		if err := uploader.ReserveUploadPart(ctx, *upload, partNumber); err != nil {
+			return trace.Wrap(err)
+		}
+
+		streamPart, err := uploader.UploadPart(ctx, *upload, partNumber, bytes.NewReader(part))
+		if err != nil {
+			return trace.Wrap(err)
+		}
+		streamParts = append(streamParts, *streamPart)
+	}
+
+	return trace.Wrap(uploader.CompleteUpload(ctx, *upload, streamParts))
+}
+
+// UploadEncryptedRecording uploads encrypted recordings using the AuditLog's configured
+// UploadHandler.
+func (l *AuditLog) UploadEncryptedRecording(ctx context.Context, sessionID string, parts iter.Seq2[[]byte, error]) error {
+	return UploadEncryptedRecording(ctx, l.UploadHandler, sessionID, parts)
+}
+
+// GetMultipartUploader returns the upload handler used by the AuditLog for persisting
+// session recordings.
+func (l *AuditLog) GetMultipartUploader() MultipartUploader {
+	return l.UploadHandler
+}
+
 // getLocalLog returns the local (file based) AuditLogger.
 func (l *AuditLog) getLocalLog() AuditLogger {
 	l.RLock()