From db321b4e19dc6073ac9f8058dfe5153f139093dd Mon Sep 17 00:00:00 2001 From: Edoardo Spadolini Date: Thu, 20 Jul 2023 17:46:18 +0200 Subject: [PATCH] Azure Blob Storage for sessions (#29069) * Azure Blob Storage for sessions * use blob batch for part cleanup * address comments * require https/http/azblob/azblob-http scheme --- constants.go | 12 + go.mod | 2 +- go.sum | 5 +- lib/events/azsessions/azsessions.go | 578 +++++++++++++++++++++++ lib/events/azsessions/azsessions_test.go | 64 +++ lib/events/azsessions/utils.go | 51 ++ lib/service/service.go | 11 + 7 files changed, 720 insertions(+), 3 deletions(-) create mode 100644 lib/events/azsessions/azsessions.go create mode 100644 lib/events/azsessions/azsessions_test.go create mode 100644 lib/events/azsessions/utils.go diff --git a/constants.go b/constants.go index ddf21dcf323..363d545cac7 100644 --- a/constants.go +++ b/constants.go @@ -338,6 +338,10 @@ const ( // GCSTestURI turns on GCS tests GCSTestURI = "TEST_GCS_URI" + // AZBlobTestURI specifies the storage account URL to use for Azure Blob + // Storage tests. + AZBlobTestURI = "TEST_AZBLOB_URI" + // AWSRunTests turns on tests executed against AWS directly AWSRunTests = "TEST_AWS" @@ -371,6 +375,14 @@ const ( // SchemeGCS is used for Google Cloud Storage SchemeGCS = "gs" + // SchemeAZBlob is the Azure Blob Storage scheme, used as the scheme in the + // session storage URI to identify a storage account accessed over https. + SchemeAZBlob = "azblob" + + // SchemeAZBlobHTTP is the Azure Blob Storage scheme, used as the scheme in the + // session storage URI to identify a storage account accessed over http. + SchemeAZBlobHTTP = "azblob-http" + // LogsDir is a log subdirectory for events and logs LogsDir = "log" diff --git a/go.mod b/go.mod index 1f14edb8709..ec1ea984c5b 100644 --- a/go.mod +++ b/go.mod @@ -21,6 +21,7 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/redisenterprise/armredisenterprise v1.1.1 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/sql/armsql v1.1.0 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.1.0 + github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.1.0 github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 github.com/HdrHistogram/hdrhistogram-go v1.1.2 github.com/Microsoft/go-winio v0.6.1 @@ -207,7 +208,6 @@ require ( github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect github.com/99designs/keyring v1.2.2 // indirect github.com/Azure/azure-sdk-for-go/sdk/internal v1.3.0 // indirect - github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.0.0 // indirect github.com/Azure/go-ntlmssp v0.0.0-20221128193559-754e69321358 // indirect github.com/AzureAD/microsoft-authentication-library-for-go v1.0.0 // indirect github.com/BurntSushi/toml v1.1.0 // indirect diff --git a/go.sum b/go.sum index 31b2af1c94b..db6c309fd15 100644 --- a/go.sum +++ b/go.sum @@ -127,10 +127,11 @@ github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/redisenterprise/armredisen github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.0.0 h1:ECsQtyERDVz3NP3kvDOTLvbQhqWp/x9EsGKtb4ogUr8= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/sql/armsql v1.1.0 h1:G2MvNS98bjXD7Vks+psbTU/uBiBH7gicij12Xc8q6lM= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/sql/armsql v1.1.0/go.mod h1:f/IvRlQ/eFP31UXVUwh3BzTOOC2cEo6/u+7g9+KTzPk= +github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage v1.2.0 h1:Ma67P/GGprNwsslzEH6+Kb8nybI8jpDTm4Wmzu2ReK8= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.1.0 h1:pYhaMoTHP/zYIJGDA1sWsfyTDjdglaoYjIFMOEcL+/U= github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/subscription/armsubscription v1.1.0/go.mod h1:iLq8GwpQhj09gpI4EdELwifR9kHrb/Q0LThq6iQq9yY= -github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.0.0 h1:u/LLAOFgsMv7HmNL4Qufg58y+qElGOt5qv0z1mURkRY= -github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.0.0/go.mod h1:2e8rMJtl2+2j+HXbTBwnyGpm5Nou7KhvSfxOq8JpTag= +github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.1.0 h1:nVocQV40OQne5613EeLayJiRAJuKlBGy+m22qWG+WRg= +github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.1.0/go.mod h1:7QJP7dr2wznCMeqIrhMgWGf7XpAQnVrJqDm9nvV3Cu4= github.com/Azure/azure-service-bus-go v0.9.1/go.mod h1:yzBx6/BUGfjfeqbRZny9AQIbIe3AcV9WZbAdpkoXOa0= github.com/Azure/azure-storage-blob-go v0.8.0/go.mod h1:lPI3aLPpuLTeUwh1sViKXFxwl2B6teiRqI0deQUvsw0= github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOElx5B5HZ4hJQsoJ/PvUvKRhJHDQXO8P8= diff --git a/lib/events/azsessions/azsessions.go b/lib/events/azsessions/azsessions.go new file mode 100644 index 00000000000..e04bcd82a15 --- /dev/null +++ b/lib/events/azsessions/azsessions.go @@ -0,0 +1,578 @@ +// Copyright 2023 Gravitational, Inc +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package azsessions + +import ( + "context" + "encoding/base64" + "fmt" + "io" + "net/url" + "strconv" + "strings" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy" + "github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming" + "github.com/Azure/azure-sdk-for-go/sdk/azidentity" + "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob" + "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blockblob" + "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container" + "github.com/google/uuid" + "github.com/gravitational/trace" + "github.com/sirupsen/logrus" + "golang.org/x/exp/slices" + "golang.org/x/sync/errgroup" + + "github.com/gravitational/teleport" + "github.com/gravitational/teleport/lib/events" + "github.com/gravitational/teleport/lib/session" +) + +// clientIDParam is the parameter in the fragment that specifies the optional client ID. +const clientIDParam = "azure_client_id" + +// sessionContainerParam and inprogressContainerParam are the parameters in the +// fragment that specify the containers to use for finalized session recordings +// and in-progress data. +const ( + sessionContainerParam = "session_container" + inprogressContainerParam = "inprogress_container" +) + +// defaultSessionContainerName and defaultInprogressContainerName are the +// default container names for finalized session recordings and for in-progress +// data. +const ( + defaultSessionContainerName = "session" + defaultInprogressContainerName = "inprogress" +) + +// sessionName returns the name of the blob that contains the recording for a +// given session. +func sessionName(sid session.ID) string { + return sid.String() +} + +// uploadMarkerPrefix is the prefix of the names of the upload marker blobs. +// Listing blobs with this prefix will return an empty blob for each upload. +const uploadMarkerPrefix = "upload/" + +// uploadMarkerName returns the blob name for the marker for a given upload. +func uploadMarkerName(upload events.StreamUpload) string { + return fmt.Sprintf("%v%v/%v", uploadMarkerPrefix, upload.SessionID, upload.ID) +} + +// partPrefix returns the prefix for the upload part blobs for a given upload. +// Listing blobs with this prefix will return all the parts that currently make +// up the upload. +func partPrefix(upload events.StreamUpload) string { + return fmt.Sprintf("part/%v/%v/", upload.SessionID, upload.ID) +} + +// partName returns the name of the blob for a specific part in an upload. +func partName(upload events.StreamUpload, partNumber int64) string { + return fmt.Sprintf("%v%v", partPrefix(upload), partNumber) +} + +// field names used for logging +const ( + fieldSessionID = "session_id" + fieldUploadID = "upload_id" + fieldPartNumber = "part" + fieldPartCount = "parts" +) + +// Config is a struct of parameters to define the behavior of Handler. +type Config struct { + // ServiceURL is the URL for the storage account to use. + ServiceURL url.URL + + // SessionContainerName is the name of the container that stores finalized + // session recordings. Defaults to [defaultSessionContainerName]. + SessionContainerName string + + // InprogressContainerName is the name of the container that stores + // in-progress data that's yet to be finalized in a recording. Defaults to + // [defaultInprogressContainerName]. + InprogressContainerName string + + // ClientID, when set, defines the managed identity's client ID to use for + // authentication. + ClientID string + + // Log is the logger to use. If unset, it will default to the global logger + // with a component of "azblob". + Log logrus.FieldLogger +} + +// SetFromURL sets values in Config based on the passed in URL: the fragment of +// the URL is parsed as if it was made out of query parameters, which define +// options for ourselves, and then the remainder of the URL is set as the +// service URL. +func (c *Config) SetFromURL(u *url.URL) error { + if u == nil { + return nil + } + + c.ServiceURL = *u + + switch c.ServiceURL.Scheme { + case teleport.SchemeAZBlob: + c.ServiceURL.Scheme = "https" + case teleport.SchemeAZBlobHTTP: + c.ServiceURL.Scheme = "http" + case "http", "https": + default: + return trace.BadParameter("unsupported URL scheme %v", c.ServiceURL.Scheme) + } + + params, err := url.ParseQuery(c.ServiceURL.EscapedFragment()) + if err != nil { + return trace.Wrap(err) + } + c.ServiceURL.Fragment = "" + c.ServiceURL.RawFragment = "" + + c.ClientID = params.Get(clientIDParam) + c.SessionContainerName = params.Get(sessionContainerParam) + c.InprogressContainerName = params.Get(inprogressContainerParam) + + return nil +} + +func (c *Config) CheckAndSetDefaults() error { + if c.SessionContainerName == "" { + c.SessionContainerName = defaultSessionContainerName + } + + if c.InprogressContainerName == "" { + c.InprogressContainerName = defaultInprogressContainerName + } + + if c.Log == nil { + c.Log = logrus.WithField(trace.Component, "azblob") + } + + return nil +} + +func NewHandler(ctx context.Context, cfg Config) (*Handler, error) { + if err := cfg.CheckAndSetDefaults(); err != nil { + return nil, trace.Wrap(err) + } + + var cred azcore.TokenCredential + if cfg.ClientID != "" { + c, err := azidentity.NewManagedIdentityCredential(&azidentity.ManagedIdentityCredentialOptions{ + ID: azidentity.ClientID(cfg.ClientID), + }) + if err != nil { + return nil, trace.Wrap(err, "creating Azure managed identity credentials") + } + cred = c + } else { + c, err := azidentity.NewDefaultAzureCredential(nil) + if err != nil { + return nil, trace.Wrap(err, "creating default Azure credentials") + } + cred = c + } + + ensureContainer := func(name string) (*container.Client, error) { + containerURL := cfg.ServiceURL + containerURL.Path = name + + cntClient, err := container.NewClient(containerURL.String(), cred, nil) + if err != nil { + return nil, trace.Wrap(err) + } + + _, err = cErr(cntClient.GetProperties(ctx, nil)) + if err == nil { + return cntClient, nil + } + if !trace.IsNotFound(err) && !trace.IsAccessDenied(err) { + return nil, trace.Wrap(err) + } + + cfg.Log.WithError(err).Debugf("Failed to confirm that the %v container exists, attempting creation.", name) + // someone else might've created the container between GetProperties and + // Create, so we ignore AlreadyExists + _, err = cErr(cntClient.Create(ctx, nil)) + if err == nil || trace.IsAlreadyExists(err) { + return cntClient, nil + } + if trace.IsAccessDenied(err) { + // we might not have permissions to read the container or to create + // it, but we might have permissions to use it + cfg.Log.WithError(err).Warnf( + "Could not create the %v container, please ensure it exists or session recordings will not be stored correctly.", name) + return cntClient, nil + } + return nil, trace.Wrap(err) + } + + session, err := ensureContainer(cfg.SessionContainerName) + if err != nil { + return nil, trace.Wrap(err) + } + + inprogress, err := ensureContainer(cfg.InprogressContainerName) + if err != nil { + return nil, trace.Wrap(err) + } + + return &Handler{ + log: cfg.Log, + cred: cred, + session: session, + inprogress: inprogress, + }, nil +} + +// Handler is a MultipartHandler that stores data in Azure Blob Storage. +type Handler struct { + log logrus.FieldLogger + cred azcore.TokenCredential + session *container.Client + inprogress *container.Client +} + +var _ events.MultipartHandler = (*Handler)(nil) + +// sessionBlob returns a BlockBlobClient for the blob of the recording of the +// session. +func (h *Handler) sessionBlob(sessionID session.ID) *blockblob.Client { + return h.session.NewBlockBlobClient(sessionName(sessionID)) +} + +// uploadMarkerBlob returns a BlockBlobClient for the marker blob of the stream +// upload. +func (h *Handler) uploadMarkerBlob(upload events.StreamUpload) *blockblob.Client { + return h.inprogress.NewBlockBlobClient(uploadMarkerName(upload)) +} + +// partBlob returns a BlockBlobClient for the blob of the part of the specified +// upload, with the given part number. +func (h *Handler) partBlob(upload events.StreamUpload, partNumber int64) *blockblob.Client { + return h.inprogress.NewBlockBlobClient(partName(upload, partNumber)) +} + +// Upload implements [events.UploadHandler]. +func (h *Handler) Upload(ctx context.Context, sessionID session.ID, reader io.Reader) (string, error) { + sessionBlob := h.sessionBlob(sessionID) + + if _, err := cErr(sessionBlob.UploadStream(ctx, reader, &blockblob.UploadStreamOptions{ + AccessConditions: &blobDoesNotExist, + })); err != nil { + return "", trace.Wrap(err) + } + h.log.WithField(fieldSessionID, sessionID).Debug("Uploaded session.") + + return sessionBlob.URL(), nil +} + +// Download implements [events.UploadHandler]. +func (h *Handler) Download(ctx context.Context, sessionID session.ID, writerAt io.WriterAt) error { + resp, err := cErr(h.sessionBlob(sessionID).DownloadStream(ctx, nil)) + if err != nil { + return trace.Wrap(err) + } + + defer func() { + if err := resp.Body.Close(); err != nil { + h.log.WithError(err).WithField(fieldSessionID, sessionID).Warn("Error closing downloaded session blob.") + } + }() + + writer, ok := writerAt.(io.Writer) + if !ok { + writer = io.NewOffsetWriter(writerAt, 0) + } + + if _, err := io.Copy(writer, resp.Body); err != nil { + return trace.ConvertSystemError(cErr0(err)) + } + + h.log.WithField(fieldSessionID, sessionID).Debug("Downloaded session.") + return nil +} + +// CreateUpload implements [events.MultipartUploader]. +func (h *Handler) CreateUpload(ctx context.Context, sessionID session.ID) (*events.StreamUpload, error) { + upload := events.StreamUpload{ + ID: uuid.NewString(), + SessionID: sessionID, + } + + if _, err := cErr(h.uploadMarkerBlob(upload).Upload(ctx, nil, &blockblob.UploadOptions{ + AccessConditions: &blobDoesNotExist, + })); err != nil { + return nil, trace.Wrap(err) + } + h.log.WithField(fieldSessionID, sessionID).Debug("Created upload marker.") + + return &upload, nil +} + +// CompleteUpload implements [events.MultipartUploader] by composing the final +// session recording blob in the session container from the parts in the +// inprogress container, using the Put Block From URL API. Might take a little +// time, but doesn't require any data transfer. +func (h *Handler) CompleteUpload(ctx context.Context, upload events.StreamUpload, parts []events.StreamPart) error { + sessionBlob := h.sessionBlob(upload.SessionID) + + // TODO(espadolini): explore the possibility of using leases to get + // exclusive access while writing, and to guarantee that leftover parts are + // cleaned up before a new attempt + + parts = slices.Clone(parts) + slices.SortFunc(parts, func(a, b events.StreamPart) bool { return a.Number < b.Number }) + + partURLs := make([]string, 0, len(parts)) + for _, part := range parts { + partURLs = append(partURLs, h.partBlob(upload, part.Number).URL()) + } + + token, err := h.cred.GetToken(ctx, policy.TokenRequestOptions{ + Scopes: []string{"https://storage.azure.com/.default"}, + }) + if err != nil { + return trace.Wrap(err, "obtaining Azure authentication token") + } + copySourceAuthorization := "Bearer " + token.Token + stageOptions := &blockblob.StageBlockFromURLOptions{ + CopySourceAuthorization: ©SourceAuthorization, + } + + log := h.log.WithFields(logrus.Fields{ + fieldSessionID: upload.SessionID, + fieldUploadID: upload.ID, + }) + + eg, egCtx := errgroup.WithContext(ctx) + eg.SetLimit(5) // default parallelism as used by azblob.DoBatchTransfer + + log.WithField(fieldPartCount, len(parts)).Debug("Beginning upload completion.") + blockNames := make([]string, len(parts)) + // TODO(espadolini): use stable names (upload id, part number and then some + // hash maybe) to avoid re-staging parts more than once across multiple + // completion attempts? + for i := range parts { + ii := i + eg.Go(func() error { + // we use block names that are local to this function so we don't + // interact with other ongoing uploads; trick copied from + // (*BlockBlobClient).UploadBuffer and UploadFile + u := uuid.New() + blockNames[ii] = base64.StdEncoding.EncodeToString(u[:]) + + if _, err := cErr(sessionBlob.StageBlockFromURL(egCtx, blockNames[ii], partURLs[ii], stageOptions)); err != nil { + return trace.Wrap(err) + } + log.WithField(fieldPartNumber, ii).Debug("Staged part.") + return nil + }) + } + if err := eg.Wait(); err != nil { + return trace.Wrap(err) + } + + log.Debug("Committing part list.") + if _, err := cErr(sessionBlob.CommitBlockList(ctx, blockNames, &blockblob.CommitBlockListOptions{ + AccessConditions: &blobDoesNotExist, + })); err != nil { + if !trace.IsAlreadyExists(err) { + return trace.Wrap(err) + } + log.Warn("Session upload already exists, cleaning up marker.") + parts = nil // don't delete parts that we didn't persist + } else { + log.Debug("Completed session upload.") + } + + // TODO(espadolini): should the cleanup run in its own goroutine? What + // should the cancellation context for the cleanup be in that case? + if _, err := cErr(h.uploadMarkerBlob(upload).Delete(ctx, nil)); err != nil && !trace.IsNotFound(err) { + log.WithError(err).WithField(fieldPartCount, len(parts)).Warn("Failed to clean up upload marker.") + return nil + } + + const batchSize = 256 // https://learn.microsoft.com/en-us/rest/api/storageservices/blob-batch + for i := 0; i < len(parts); i += batchSize { + batch, err := cErr(h.inprogress.NewBatchBuilder()) + if err != nil { + return trace.Wrap(err) + } + + m := batchSize + if len(parts[i:]) < batchSize { + m = len(parts[i:]) + } + + for _, part := range parts[i : i+m] { + if err := batch.Delete(partName(upload, part.Number), nil); err != nil { + return trace.Wrap(err) + } + } + + resp, err := cErr(h.inprogress.SubmitBatch(ctx, batch, nil)) + if err != nil { + log.WithField(fieldPartNumber, parts[i].Number).WithError(err).Warn("Failed to clean up part batch.") + continue + } + + errs := 0 + for _, r := range resp.Responses { + if r.Error != nil { + err = r.Error + errs++ + } + } + if errs > 0 { + log.WithFields(logrus.Fields{ + fieldPartNumber: parts[i].Number, + "errors": errs, + "last_error": err, + }).Warn("Failed to clean up part batch.") + } + } + + return nil +} + +// ReserveUploadPart implements [events.MultipartUploader]. +func (*Handler) ReserveUploadPart(ctx context.Context, upload events.StreamUpload, partNumber int64) error { + return nil +} + +// UploadPart implements [events.MultipartUploader]. +func (h *Handler) UploadPart(ctx context.Context, upload events.StreamUpload, partNumber int64, partBody io.ReadSeeker) (*events.StreamPart, error) { + partBlob := h.partBlob(upload, partNumber) + + // our parts are just over 5 MiB (events.MinUploadPartSizeBytes) so we can + // upload them in one shot + if _, err := cErr(partBlob.Upload(ctx, streaming.NopCloser(partBody), nil)); err != nil { + return nil, trace.Wrap(err) + } + h.log.WithFields(logrus.Fields{ + fieldSessionID: upload.SessionID, + fieldUploadID: upload.ID, + fieldPartNumber: partNumber, + }).Debug("Uploaded part.") + + return &events.StreamPart{Number: partNumber}, nil +} + +// ListParts implements [events.MultipartUploader]. +func (h *Handler) ListParts(ctx context.Context, upload events.StreamUpload) ([]events.StreamPart, error) { + prefix := partPrefix(upload) + + var parts []events.StreamPart + pager := h.inprogress.NewListBlobsFlatPager(&azblob.ListBlobsFlatOptions{ + Prefix: &prefix, + }) + for pager.More() { + resp, err := cErr(pager.NextPage(ctx)) + if err != nil { + return nil, trace.Wrap(err) + } + + if resp.Segment == nil { + continue + } + parts = slices.Grow(parts, len(resp.Segment.BlobItems)) + for _, b := range resp.Segment.BlobItems { + if b == nil || + b.Name == nil || + !strings.HasPrefix(*b.Name, prefix) { + continue + } + + pn := strings.TrimPrefix(*b.Name, prefix) + partNumber, err := strconv.ParseInt(pn, 10, 64) + if err != nil { + continue + } + + parts = append(parts, events.StreamPart{Number: partNumber}) + } + } + + slices.SortFunc(parts, func(a, b events.StreamPart) bool { return a.Number < b.Number }) + + return parts, nil +} + +// ListUploads implements [events.MultipartUploader]. +func (h *Handler) ListUploads(ctx context.Context) ([]events.StreamUpload, error) { + prefix := uploadMarkerPrefix + var uploads []events.StreamUpload + + pager := h.inprogress.NewListBlobsFlatPager(&azblob.ListBlobsFlatOptions{ + Prefix: &prefix, + }) + for pager.More() { + r, err := cErr(pager.NextPage(ctx)) + if err != nil { + return nil, trace.Wrap(err) + } + + if r.Segment == nil { + continue + } + uploads = slices.Grow(uploads, len(r.Segment.BlobItems)) + for _, b := range r.Segment.BlobItems { + if b == nil || + b.Name == nil || + !strings.HasPrefix(*b.Name, prefix) || + b.Properties == nil || + b.Properties.CreationTime == nil { + continue + } + + name := strings.TrimPrefix(*b.Name, prefix) + sid, uid, ok := strings.Cut(name, "/") + if !ok { + continue + } + if _, err := session.ParseID(sid); err != nil { + continue + } + if _, err := uuid.Parse(uid); err != nil { + continue + } + + uploads = append(uploads, events.StreamUpload{ + ID: uid, + SessionID: session.ID(sid), + Initiated: *b.Properties.CreationTime, + }) + } + } + + slices.SortFunc(uploads, func(a, b events.StreamUpload) bool { return a.Initiated.Before(b.Initiated) }) + + return uploads, nil +} + +// GetUploadMetadata implements [events.MultipartUploader]. +func (h *Handler) GetUploadMetadata(sessionID session.ID) events.UploadMetadata { + return events.UploadMetadata{ + URL: h.sessionBlob(sessionID).URL(), + SessionID: sessionID, + } +} diff --git a/lib/events/azsessions/azsessions_test.go b/lib/events/azsessions/azsessions_test.go new file mode 100644 index 00000000000..88f8eb2819b --- /dev/null +++ b/lib/events/azsessions/azsessions_test.go @@ -0,0 +1,64 @@ +// Copyright 2023 Gravitational, Inc +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package azsessions + +import ( + "context" + "net/url" + "os" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/gravitational/teleport" + "github.com/gravitational/teleport/lib/events/test" + "github.com/gravitational/teleport/lib/utils" +) + +func TestMain(m *testing.M) { + utils.InitLoggerForTests() + os.Exit(m.Run()) +} + +// TestStreams runs the standard events test suite over azsessions, if a +// configuration URL is specified in the appropriate envvar. +func TestStreams(t *testing.T) { + ctx := context.Background() + + envURL := os.Getenv(teleport.AZBlobTestURI) + if envURL == "" { + t.Skipf("Skipping azsessions tests as %q is not set.", teleport.AZBlobTestURI) + } + + u, err := url.Parse(envURL) + require.NoError(t, err) + + var config Config + err = config.SetFromURL(u) + require.NoError(t, err) + + handler, err := NewHandler(ctx, config) + require.Nil(t, err) + + t.Run("StreamManyParts", func(t *testing.T) { + test.StreamManyParts(t, handler) + }) + t.Run("UploadDownload", func(t *testing.T) { + test.UploadDownload(t, handler) + }) + t.Run("DownloadNotFound", func(t *testing.T) { + test.DownloadNotFound(t, handler) + }) +} diff --git a/lib/events/azsessions/utils.go b/lib/events/azsessions/utils.go new file mode 100644 index 00000000000..5febc0e4fb7 --- /dev/null +++ b/lib/events/azsessions/utils.go @@ -0,0 +1,51 @@ +// Copyright 2023 Gravitational, Inc +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package azsessions + +import ( + "errors" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore" + "github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blob" + "github.com/gravitational/trace" +) + +var eTagAny = azcore.ETagAny + +var blobDoesNotExist = blob.AccessConditions{ + ModifiedAccessConditions: &blob.ModifiedAccessConditions{ + IfNoneMatch: &eTagAny, + }, +} + +// cErr0 attempts to convert err to a meaningful trace error if it's a +// *azblob.StorageError; if it can't, it'll return the error, wrapped. +func cErr0(err error) error { + if err == nil { + return nil + } + + var stErr *azcore.ResponseError + if !errors.As(err, &stErr) || stErr == nil { + return trace.Wrap(err) + } + + return trace.WrapWithMessage(trace.ReadError(stErr.StatusCode, nil), stErr.ErrorCode) +} + +// cErr converts the error as in [cErr0], leaving the first argument untouched. +func cErr[T any](v T, err error) (T, error) { + return v, cErr0(err) +} diff --git a/lib/service/service.go b/lib/service/service.go index 5fa2fc49529..389d1bd41b9 100644 --- a/lib/service/service.go +++ b/lib/service/service.go @@ -89,6 +89,7 @@ import ( "github.com/gravitational/teleport/lib/defaults" "github.com/gravitational/teleport/lib/events" "github.com/gravitational/teleport/lib/events/athena" + "github.com/gravitational/teleport/lib/events/azsessions" "github.com/gravitational/teleport/lib/events/dynamoevents" "github.com/gravitational/teleport/lib/events/filesessions" "github.com/gravitational/teleport/lib/events/firestoreevents" @@ -1370,6 +1371,16 @@ func initAuthUploadHandler(ctx context.Context, auditConfig types.ClusterAuditCo return nil, trace.Wrap(err) } return handler, nil + case teleport.SchemeAZBlob, teleport.SchemeAZBlobHTTP: + var config azsessions.Config + if err := config.SetFromURL(uri); err != nil { + return nil, trace.Wrap(err) + } + handler, err := azsessions.NewHandler(ctx, config) + if err != nil { + return nil, trace.Wrap(err) + } + return handler, nil case teleport.SchemeFile: if err := os.MkdirAll(uri.Path, teleport.SharedDirMode); err != nil { return nil, trace.ConvertSystemError(err)