Don't hardcode the registry URL to authenticate to, use the API respo…

…nse (#1525)

There were a few paths where the registry to log in to was hardcoded
based on the environment (but not all of them -- for instance the Push
path using the go client library was already using this field to auth against!)

This change makes the Push and Pull flow use the pre-existing
`ImageRepository` flow in the Deploy(ment) response from Astro.

This also removes a bit of duplicated code -- various places were
stripping of the `Bearer ` prefix from the password before calling the
Push function, but that function itself did that.

Co-authored-by: Neel Dalsania <neel.dalsania@astronomer.io>

.golangci.yml

@@ -49,7 +49,7 @@ linters-settings:
       mnd:
         # don't include the "operation" and "assign"
         checks: [argument, case, condition, return]
-        ignored-numbers: 10,64  # numbers used by strconv
+        ignored-numbers: 2,10,64  # numbers used by strconv
   govet:
     check-shadowing: true
     settings:
@@ -136,6 +136,7 @@ issues:
         - typecheck
         - revive
         - unused
+        - goerr113
   exclude:
     - "shadow: declaration of .err. shadows declaration"
     - "sloppyTestFuncName: function cleanUpInitFiles should be of form"

airflow/container.go

@@ -46,8 +46,8 @@ type RegistryHandler interface {
 // ImageHandler defines methods require to handle all operations on/for container images
 type ImageHandler interface {
 	Build(dockerfile, buildSecretString string, config types.ImageBuildConfig) error
-	Push(registry, username, token, remoteImage string) error
-	Pull(registry, username, token, remoteImage string) error
+	Push(remoteImage, username, token string) error
+	Pull(remoteImage, username, token string) error
 	GetLabel(altImageName, labelName string) (string, error)
 	DoesImageExist(image string) error
 	ListLabels() (map[string]string, error)

airflow/docker.go

@@ -625,27 +625,15 @@ func (d *DockerCompose) pullImageFromDeployment(deploymentID string, platformCor
 	if err != nil {
 		return err
 	}
-	domain := c.Domain
-	if domain == "" {
-		return errors.New("no domain set, re-authenticate")
-	}
 	ws := c.Workspace
-	registry := GetRegistryURL(domain)
-	repository := registry + "/" + c.Organization + "/" + deploymentID
 	currentDeployment, err := deployment.GetDeployment(ws, deploymentID, "", true, platformCoreClient, nil)
 	if err != nil {
 		return err
 	}
-	currentImageTag := currentDeployment.ImageTag
-	deploymentImage := fmt.Sprintf("%s:%s", repository, currentImageTag)
+	deploymentImage := fmt.Sprintf("%s:%s", currentDeployment.ImageRepository, currentDeployment.ImageTag)
 	token := c.Token
-	// Splitting out the Bearer part from the token
-	splittedToken := strings.Split(token, " ")
-	if len(splittedToken) > 1 {
-		token = strings.Split(token, " ")[1]
-	}
 	fmt.Printf("\nPulling image from Astro Deployment %s\n\n", currentDeployment.Name)
-	err = d.imageHandler.Pull(registry, registryUsername, token, deploymentImage)
+	err = d.imageHandler.Pull(deploymentImage, registryUsername, token)
 	if err != nil {
 		return err
 	}
@@ -771,16 +759,6 @@ func (d *DockerCompose) dagTest(testHomeDirectory, newAirflowVersion, newDockerF
 	return nil
 }
 
-func GetRegistryURL(domain string) string {
-	var registry string
-	if domain == "localhost" {
-		registry = config.CFG.LocalRegistry.GetString()
-	} else {
-		registry = "images." + strings.Split(domain, ".")[0] + ".cloud"
-	}
-	return registry
-}
-
 func upgradeDockerfile(oldDockerfilePath, newDockerfilePath, newTag, newImage string) error { //nolint:gocognit
 	// Read the content of the old Dockerfile
 	content, err := os.ReadFile(oldDockerfilePath)

airflow/docker_image.go

@@ -36,6 +36,10 @@ const (
 
 var errGetImageLabel = errors.New("error getting image label")
 
+var getDockerClient = func() (client.APIClient, error) {
+	return client.NewClientWithOpts(client.FromEnv)
+}
+
 type DockerImage struct {
 	imageName string
 }
@@ -299,13 +303,18 @@ func (d *DockerImage) CreatePipFreeze(altImageName, pipFreezeFile string) error
 	return nil
 }
 
-func (d *DockerImage) Push(registry, username, token, remoteImage string) error {
+func (d *DockerImage) Push(remoteImage, username, token string) error {
 	dockerCommand := config.CFG.DockerCommand.GetString()
 	err := cmdExec(dockerCommand, nil, nil, "tag", d.imageName, remoteImage)
 	if err != nil {
 		return fmt.Errorf("command '%s tag %s %s' failed: %w", dockerCommand, d.imageName, remoteImage, err)
 	}
 
+	registry, err := d.getRegistryToAuth(remoteImage)
+	if err != nil {
+		return err
+	}
+
 	// Push image to registry
 	fmt.Println(pushingImagePrompt)
 
@@ -334,13 +343,31 @@ func (d *DockerImage) Push(registry, username, token, remoteImage string) error
 
 	log.Debugf("Exec Push %s creds %v \n", dockerCommand, authConfig)
 
+	err = d.pushWithClient(&authConfig, remoteImage)
+
+	if err != nil {
+		// if it does not work with the go library use bash to run docker commands. Support for (old?) versions of Colima
+		err = pushWithBash(&authConfig, remoteImage)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Delete the image tags we just generated
+	err = cmdExec(dockerCommand, nil, nil, "rmi", remoteImage)
+	if err != nil {
+		return fmt.Errorf("command '%s rmi %s' failed: %w", dockerCommand, remoteImage, err)
+	}
+	return nil
+}
+
+func (d *DockerImage) pushWithClient(authConfig *cliTypes.AuthConfig, remoteImage string) error {
 	ctx := context.Background()
 
-	cli, err := client.NewClientWithOpts(client.FromEnv)
+	cli, err := getDockerClient()
 	if err != nil {
 		log.Debugf("Error setting up new Client ops %v", err)
-		// if NewClientWithOpt does not work use bash to run docker commands
-		return useBash(&authConfig, remoteImage)
+		return err
 	}
 	cli.NegotiateAPIVersion(ctx)
 	buf, err := json.Marshal(authConfig)
@@ -353,27 +380,40 @@ func (d *DockerImage) Push(registry, username, token, remoteImage string) error
 	if err != nil {
 		log.Debugf("Error pushing image to docker: %v", err)
 		// if NewClientWithOpt does not work use bash to run docker commands
-		return useBash(&authConfig, remoteImage)
+		return err
 	}
 	defer responseBody.Close()
-	err = displayJSONMessagesToStream(responseBody, nil)
+	return displayJSONMessagesToStream(responseBody, nil)
+}
+
+// Get the registry name to authenticate against
+func (d *DockerImage) getRegistryToAuth(imageName string) (string, error) {
+	domain, err := config.GetCurrentDomain()
 	if err != nil {
-		return useBash(&authConfig, remoteImage)
+		return "", err
 	}
-	// Delete the image tags we just generated
-	err = cmdExec(dockerCommand, nil, nil, "rmi", remoteImage)
-	if err != nil {
-		return fmt.Errorf("command '%s rmi %s' failed: %w", dockerCommand, remoteImage, err)
+
+	if domain == "localhost" {
+		return config.CFG.LocalRegistry.GetString(), nil
 	}
-	return nil
+	parts := strings.SplitN(imageName, "/", 2)
+	if len(parts) != 2 || !strings.Contains(parts[1], "/") {
+		// This _should_ be impossible for users to hit
+		return "", fmt.Errorf("internal logic error: unsure how to get registry from image name %q", imageName) //nolint:goerr113
+	}
+	return parts[0], nil
 }
 
-func (d *DockerImage) Pull(registry, username, token, remoteImage string) error {
+func (d *DockerImage) Pull(remoteImage, username, token string) error {
 	// Pulling image to registry
 	fmt.Println(pullingImagePrompt)
 	dockerCommand := config.CFG.DockerCommand.GetString()
 	var err error
 	if username != "" { // Case for cloud image push where we have both registry user & pass, for software login happens during `astro login` itself
+		var registry string
+		if registry, err = d.getRegistryToAuth(remoteImage); err != nil {
+			return err
+		}
 		pass := token
 		pass = strings.TrimPrefix(pass, prefix)
 		cmd := "echo \"" + pass + "\"" + " | " + dockerCommand + " login " + registry + " -u " + username + " --password-stdin"
@@ -583,7 +623,7 @@ var cmdExec = func(cmd string, stdout, stderr io.Writer, args ...string) error {
 }
 
 // When login and push do not work use bash to run docker commands, this function is for users using colima
-func useBash(authConfig *cliTypes.AuthConfig, image string) error {
+func pushWithBash(authConfig *cliTypes.AuthConfig, image string) error {
 	dockerCommand := config.CFG.DockerCommand.GetString()
 
 	var err error
@@ -592,19 +632,10 @@ func useBash(authConfig *cliTypes.AuthConfig, image string) error {
 		pass = strings.TrimPrefix(pass, prefix)
 		cmd := "echo \"" + pass + "\"" + " | " + dockerCommand + " login " + authConfig.ServerAddress + " -u " + authConfig.Username + " --password-stdin"
 		err = cmdExec("bash", os.Stdout, os.Stderr, "-c", cmd) // This command will only work on machines that have bash. If users have issues we will revist
-	}
-	if err != nil {
-		return err
+		if err != nil {
+			return err
+		}
 	}
 	// docker push <image>
-	err = cmdExec(dockerCommand, os.Stdout, os.Stderr, "push", image)
-	if err != nil {
-		return err
-	}
-	// Delete the image tags we just generated
-	err = cmdExec(dockerCommand, nil, nil, "rmi", image)
-	if err != nil {
-		return fmt.Errorf("command '%s rmi %s' failed: %w", dockerCommand, image, err)
-	}
-	return nil
+	return cmdExec(dockerCommand, os.Stdout, os.Stderr, "push", image)
 }