Add support for build secrets

docs/yaml.md

@@ -86,6 +86,20 @@ build:
 
 Your code is _not_ available to commands in `run`. This is so we can build your image efficiently when running locally.
 
+Each command in `run` can be either a string or a dictionary in the following format:
+
+````yaml
+build:
+  run:
+    - command: pip install
+      mounts:
+        - type: secret
+          id: pip
+          target: /etc/pip.conf
+```
+
+You can use secret mounts to securely pass credentials to setup commands, without baking them into the image. For more information, see [Dockerfile reference](https://docs.docker.com/engine/reference/builder/#run---mounttypesecret).
+
 ### `system_packages`
 
 A list of Ubuntu APT packages to install. For example:
@@ -95,7 +109,7 @@ build:
   system_packages:
     - "ffmpeg"
     - "libavcodec-dev"
-```
+````
 
 ## `image`
 

pkg/config/config.go

@@ -19,16 +19,25 @@ import (
 // TODO(andreas): custom cpu/gpu installs
 // TODO(andreas): suggest valid torchvision versions (e.g. if the user wants to use 0.8.0, suggest 0.8.1)
 
+type RunItem struct {
+	Command string `json:"command,omitempty" yaml:"command"`
+	Mounts  []struct {
+		Type   string `json:"type,omitempty" yaml:"type"`
+		ID     string `json:"id,omitempty" yaml:"id"`
+		Target string `json:"target,omitempty" yaml:"target"`
+	} `json:"mounts,omitempty" yaml:"mounts"`
+}
+
 type Build struct {
-	GPU                bool     `json:"gpu,omitempty" yaml:"gpu"`
-	PythonVersion      string   `json:"python_version,omitempty" yaml:"python_version"`
-	PythonRequirements string   `json:"python_requirements,omitempty" yaml:"python_requirements"`
-	PythonPackages     []string `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility
-	Run                []string `json:"run,omitempty" yaml:"run"`
-	SystemPackages     []string `json:"system_packages,omitempty" yaml:"system_packages"`
-	PreInstall         []string `json:"pre_install,omitempty" yaml:"pre_install"` // Deprecated, but included for backwards compatibility
-	CUDA               string   `json:"cuda,omitempty" yaml:"cuda"`
-	CuDNN              string   `json:"cudnn,omitempty" yaml:"cudnn"`
+	GPU                bool      `json:"gpu,omitempty" yaml:"gpu"`
+	PythonVersion      string    `json:"python_version,omitempty" yaml:"python_version"`
+	PythonRequirements string    `json:"python_requirements,omitempty" yaml:"python_requirements"`
+	PythonPackages     []string  `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility
+	Run                []RunItem `json:"run,omitempty" yaml:"run"`
+	SystemPackages     []string  `json:"system_packages,omitempty" yaml:"system_packages"`
+	PreInstall         []string  `json:"pre_install,omitempty" yaml:"pre_install"` // Deprecated, but included for backwards compatibility
+	CUDA               string    `json:"cuda,omitempty" yaml:"cuda"`
+	CuDNN              string    `json:"cudnn,omitempty" yaml:"cudnn"`
 
 	pythonRequirementsContent []string
 }
@@ -54,6 +63,44 @@ func DefaultConfig() *Config {
 	}
 }
 
+func (r *RunItem) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var commandOrMap interface{}
+	if err := unmarshal(&commandOrMap); err != nil {
+		return err
+	}
+
+	switch v := commandOrMap.(type) {
+	case string:
+		r.Command = v
+	case map[interface{}]interface{}:
+		var data []byte
+		var err error
+
+		if data, err = yaml.Marshal(v); err != nil {
+			return err
+		}
+
+		aux := struct {
+			Command string `yaml:"command"`
+			Mounts  []struct {
+				Type   string `yaml:"type"`
+				ID     string `yaml:"id"`
+				Target string `yaml:"target"`
+			} `yaml:"mounts,omitempty"`
+		}{}
+
+		if err := yaml.Unmarshal(data, &aux); err != nil {
+			return err
+		}
+
+		*r = RunItem(aux)
+	default:
+		return fmt.Errorf("unexpected type %T for RunItem", v)
+	}
+
+	return nil
+}
+
 func FromYAML(contents []byte) (*Config, error) {
 	config := DefaultConfig()
 	if err := yaml.Unmarshal(contents, config); err != nil {

pkg/config/data/config_schema_v1.0.json

@@ -85,6 +85,35 @@
               {
                 "$id": "#/properties/build/properties/run/items/anyOf/0",
                 "type": "string"
+              },
+              {
+                "$id": "#/properties/build/properties/run/items/anyOf/1",
+                "type": "object",
+                "properties": {
+                  "command": {
+                    "type": "string"
+                  },
+                  "mounts": {
+                    "type": "array",
+                    "items": {
+                      "type": "object",
+                      "properties": {
+                        "type": {
+                          "type": "string",
+                          "enum": ["secret"]
+                        },
+                        "id": {
+                          "type": "string"
+                        },
+                        "target": {
+                          "type": "string"
+                        }
+                      },
+                      "required": ["type", "id", "target"]
+                    }
+                  }
+                },
+                "required": ["command"]
               }
             ]
           }

pkg/dockerfile/generator.go

@@ -228,17 +228,31 @@ func (g *Generator) run() (string, error) {
 	runCommands := g.Config.Build.Run
 
 	// For backwards compatibility
-	runCommands = append(runCommands, g.Config.Build.PreInstall...)
+	for _, command := range g.Config.Build.PreInstall {
+		runCommands = append(runCommands, config.RunItem{Command: command})
+	}
 
 	lines := []string{}
 	for _, run := range runCommands {
-		run = strings.TrimSpace(run)
-		if strings.Contains(run, "\n") {
+		command := strings.TrimSpace(run.Command)
+		if strings.Contains(command, "\n") {
 			return "", fmt.Errorf(`One of the commands in 'run' contains a new line, which won't work. You need to create a new list item in YAML prefixed with '-' for each command.
 
-This is the offending line: %s`, run)
+This is the offending line: %s`, command)
+		}
+
+		if len(run.Mounts) > 0 {
+			mounts := []string{}
+			for _, mount := range run.Mounts {
+				if mount.Type == "secret" {
+					secretMount := fmt.Sprintf("--mount=type=secret,id=%s,target=%s", mount.ID, mount.Target)
+					mounts = append(mounts, secretMount)
+				}
+			}
+			lines = append(lines, fmt.Sprintf("RUN %s %s", strings.Join(mounts, " "), command))
+		} else {
+			lines = append(lines, "RUN "+command)
 		}
-		lines = append(lines, "RUN "+run)
 	}
 	return strings.Join(lines, "\n"), nil
 }

test-integration/test_integration/test_run.py

@@ -17,3 +17,32 @@ def test_run(tmpdir_factory):
         capture_output=True,
     )
     assert b"hello world" in result.stdout
+
+
+def test_run_with_secret(tmpdir_factory):
+    tmpdir = tmpdir_factory.mktemp("project")
+    with open(tmpdir / "cog.yaml", "w") as f:
+        cog_yaml = """
+build:
+  python_version: "3.8"
+  run:
+    - echo hello world
+    - command: >-
+        echo shh
+      mounts:
+        - type: secret
+          id: foo
+          target: secret.txt
+        """
+        f.write(cog_yaml)
+    with open(tmpdir / "secret.txt", "w") as f:
+        f.write("🤫")
+
+    result = subprocess.run(
+        ["cog", "debug"],
+        cwd=tmpdir,
+        check=True,
+        capture_output=True,
+    )
+    assert b"RUN echo hello world" in result.stdout
+    assert b"RUN --mount=type=secret,id=foo,target=secret.txt echo shh" in result.stdout