constellation/internal/osimage/uplosi/uplosiupload.go
Malte Poll 65ed286c7d
image: replicate nvidia Azure images in eastus2 (#3027)
Co-authored-by: derpsteb <ob@edgeless.systems>
2024-07-11 11:09:15 +02:00

287 lines
8.8 KiB
Go

/*
Copyright (c) Edgeless Systems GmbH
SPDX-License-Identifier: AGPL-3.0-only
*/
// package uplosi implements uploading os images using uplosi.
package uplosi
import (
"bytes"
"context"
_ "embed"
"errors"
"fmt"
"log/slog"
"os"
"os/exec"
"path/filepath"
"regexp"
"strings"
"time"
"github.com/BurntSushi/toml"
"github.com/edgelesssys/constellation/v2/internal/api/versionsapi"
"github.com/edgelesssys/constellation/v2/internal/cloud/cloudprovider"
"github.com/edgelesssys/constellation/v2/internal/osimage"
)
//go:embed uplosi.conf.in
var uplosiConfigTemplate string
const timestampFormat = "20060102150405"
// Uploader can upload os images using uplosi.
type Uploader struct {
uplosiPath string
log *slog.Logger
}
// New creates a new Uploader.
func New(uplosiPath string, log *slog.Logger) *Uploader {
return &Uploader{
uplosiPath: uplosiPath,
log: log,
}
}
// Upload uploads the given os image using uplosi.
func (u *Uploader) Upload(ctx context.Context, req *osimage.UploadRequest) ([]versionsapi.ImageInfoEntry, error) {
config, err := prepareUplosiConfig(req)
if err != nil {
return nil, err
}
workspace, err := prepareWorkspace(config)
if err != nil {
return nil, err
}
defer os.RemoveAll(workspace)
uplosiOutput, err := runUplosi(ctx, u.uplosiPath, workspace, req.ImagePath)
if err != nil {
return nil, err
}
return parseUplosiOutput(uplosiOutput, req.Provider, req.AttestationVariant)
}
func prepareUplosiConfig(req *osimage.UploadRequest) ([]byte, error) {
var config map[string]any
if _, err := toml.Decode(uplosiConfigTemplate, &config); err != nil {
return nil, err
}
imageVersionStr, err := imageVersion(req.Provider, req.Version, req.Timestamp)
if err != nil {
return nil, err
}
baseConfig := config["base"].(map[string]any)
awsConfig := baseConfig["aws"].(map[string]any)
azureConfig := baseConfig["azure"].(map[string]any)
gcpConfig := baseConfig["gcp"].(map[string]any)
openstackConfig := baseConfig["openstack"].(map[string]any)
baseConfig["imageVersion"] = imageVersionStr
baseConfig["provider"] = strings.ToLower(req.Provider.String())
extendAWSConfig(awsConfig, req.Version, req.AttestationVariant, req.Timestamp)
extendAzureConfig(azureConfig, req.Version, req.AttestationVariant, req.Timestamp)
extendGCPConfig(gcpConfig, req.Version, req.AttestationVariant)
extendOpenStackConfig(openstackConfig, req.Version, req.AttestationVariant)
buf := new(bytes.Buffer)
if err := toml.NewEncoder(buf).Encode(config); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
func prepareWorkspace(config []byte) (string, error) {
workspace, err := os.MkdirTemp("", "uplosi-")
if err != nil {
return "", err
}
// write config to workspace
configPath := filepath.Join(workspace, "uplosi.conf")
if err := os.WriteFile(configPath, config, 0o644); err != nil {
return "", err
}
return workspace, nil
}
func runUplosi(ctx context.Context, uplosiPath string, workspace string, rawImage string) ([]byte, error) {
imagePath, err := filepath.Abs(rawImage)
if err != nil {
return nil, err
}
uplosiCmd := exec.CommandContext(ctx, uplosiPath, "upload", imagePath)
uplosiCmd.Dir = workspace
uplosiCmd.Stderr = os.Stderr
return uplosiCmd.Output()
}
func parseUplosiOutput(output []byte, csp cloudprovider.Provider, attestationVariant string) ([]versionsapi.ImageInfoEntry, error) {
lines := strings.Split(string(output), "\n")
var imageReferences []versionsapi.ImageInfoEntry
for _, line := range lines {
if len(line) == 0 {
continue
}
var region, reference string
if csp == cloudprovider.AWS {
var err error
region, reference, err = awsParseAMIARN(line)
if err != nil {
return nil, err
}
} else {
reference = line
}
imageReferences = append(imageReferences, versionsapi.ImageInfoEntry{
CSP: strings.ToLower(csp.String()),
AttestationVariant: attestationVariant,
Reference: reference,
Region: region,
})
}
return imageReferences, nil
}
func imageVersion(csp cloudprovider.Provider, version versionsapi.Version, timestamp time.Time) (string, error) {
cleanSemver := strings.TrimPrefix(regexp.MustCompile(`^v\d+\.\d+\.\d+`).FindString(version.Version()), "v")
if csp != cloudprovider.Azure {
return cleanSemver, nil
}
switch {
case version.Stream() == "stable":
fallthrough
case version.Stream() == "debug" && version.Ref() == "-":
return cleanSemver, nil
}
formattedTime := timestamp.Format(timestampFormat)
if len(formattedTime) != len(timestampFormat) {
return "", errors.New("invalid timestamp")
}
// <year>.<month><day>.<time>
return formattedTime[:4] + "." + formattedTime[4:8] + "." + formattedTime[8:], nil
}
func extendAWSConfig(awsConfig map[string]any, version versionsapi.Version, attestationVariant string, timestamp time.Time) {
awsConfig["amiName"] = awsAMIName(version, attestationVariant, timestamp)
awsConfig["snapshotName"] = awsAMIName(version, attestationVariant, timestamp)
awsConfig["blobName"] = fmt.Sprintf("image-%s-%s-%s-%d.raw", version.Stream(), version.Version(), attestationVariant, timestamp.Unix())
}
func awsAMIName(version versionsapi.Version, attestationVariant string, timestamp time.Time) string {
if version.Stream() == "stable" {
return fmt.Sprintf("constellation-%s-%s", version.Version(), attestationVariant)
}
return fmt.Sprintf("constellation-%s-%s-%s-%s", version.Stream(), version.Version(), attestationVariant, timestamp.Format(timestampFormat))
}
func awsParseAMIARN(arn string) (region string, amiID string, retErr error) {
parts := strings.Split(arn, ":")
if len(parts) != 6 {
return "", "", fmt.Errorf("invalid ARN (expected 5 path components) %q", arn)
}
if parts[0] != "arn" {
return "", "", fmt.Errorf("invalid ARN (prefix mismatch) %q", arn)
}
if parts[1] != "aws" {
return "", "", fmt.Errorf("invalid ARN (provider mismatch) %q", arn)
}
if parts[2] != "ec2" {
return "", "", fmt.Errorf("invalid ARN (service mismatch) %q", arn)
}
resourceParts := strings.Split(parts[5], "/")
if len(resourceParts) != 2 {
return "", "", fmt.Errorf("invalid ARN (expected resource type/id) %q", arn)
}
if resourceParts[0] != "image" {
return "", "", fmt.Errorf("invalid ARN (resource type mismatch) %q", arn)
}
return parts[3], resourceParts[1], nil
}
func extendAzureConfig(azureConfig map[string]any, version versionsapi.Version, attestationVariant string, timestamp time.Time) {
azureConfig["replicationRegions"] = azureReplicationRegions(attestationVariant)
azureConfig["attestationVariant"] = attestationVariant
azureConfig["sharedImageGallery"] = azureGalleryName(version, attestationVariant)
azureConfig["imageDefinitionName"] = azureImageOffer(version)
azureConfig["offer"] = azureImageOffer(version)
formattedTime := timestamp.Format(timestampFormat)
azureConfig["diskName"] = fmt.Sprintf("constellation-%s-%s-%s", version.Stream(), formattedTime, attestationVariant)
}
func azureGalleryName(version versionsapi.Version, attestationVariant string) string {
var prefix string
switch version.Stream() {
case "stable":
prefix = "Constellation"
case "debug":
prefix = "Constellation_Debug"
default:
prefix = "Constellation_Testing"
}
var suffix string
switch attestationVariant {
case "azure-tdx":
suffix = "_TDX"
case "azure-sev-snp":
suffix = "_CVM"
}
return prefix + suffix
}
func azureImageOffer(version versionsapi.Version) string {
switch {
case version.Stream() == "stable":
return "constellation"
case version.Stream() == "debug" && version.Ref() == "-":
return version.Version()
}
return version.Ref() + "-" + version.Stream()
}
func azureReplicationRegions(attestationVariant string) []string {
switch attestationVariant {
case "azure-tdx":
return []string{"northeurope", "westeurope", "centralus", "eastus2"}
case "azure-sev-snp":
return []string{"northeurope", "westeurope", "germanywestcentral", "eastus", "eastus2", "westus", "southeastasia"}
}
return nil
}
func extendGCPConfig(gcpConfig map[string]any, version versionsapi.Version, attestationVariant string) {
gcpConfig["imageFamily"] = gcpImageFamily(version)
gcpConfig["imageName"] = gcpImageName(version, attestationVariant)
gcpConfig["blobName"] = gcpImageName(version, attestationVariant) + ".tar.gz"
}
func extendOpenStackConfig(openstackConfig map[string]any, version versionsapi.Version, attestationVariant string) {
openstackConfig["imageName"] = fmt.Sprintf("constellation-%s-%s-%s", version.Stream(), version.Version(), attestationVariant)
}
func gcpImageFamily(version versionsapi.Version) string {
if version.Stream() == "stable" {
return "constellation"
}
truncatedRef := version.Ref()
if len(version.Ref()) > 45 {
truncatedRef = version.Ref()[:45]
}
return "constellation-" + truncatedRef
}
func gcpImageName(version versionsapi.Version, attestationVariant string) string {
return strings.ReplaceAll(version.Version(), ".", "-") + "-" + attestationVariant + "-" + version.Stream()
}