Add support for custom partitions

.gitignore

@@ -8,4 +8,5 @@ bin
 *-initrd.img
 *-kernel
 
+*.img
 dist

pkg/api/build.go

@@ -18,6 +18,11 @@ import (
 func createBuild(name, format, output string, w http.ResponseWriter, r *http.Request) {
 	log.Debugf("create build, name: %s, format: %s, output: %s", name, format, output)
 	body, _ := ioutil.ReadAll(r.Body)
+	imgConfig, err := image.NewConfig(body)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
 
 	buildDir, err := ioutil.TempDir("", "linuxkit")
 	if err != nil {
@@ -27,22 +32,22 @@ func createBuild(name, format, output string, w http.ResponseWriter, r *http.Req
 
 	if err := linuxkit.Build(name, body, []string{format}, buildDir); err != nil {
 		if linuxkit.IsInvalidConfiguration(err) {
-			http.Error(w, err.Error(), 400)
+			http.Error(w, err.Error(), http.StatusBadRequest)
 		} else if linuxkit.IsBuildFailed(err) {
-			http.Error(w, err.Error(), 503)
+			http.Error(w, err.Error(), http.StatusServiceUnavailable)
 		} else {
-			http.Error(w, err.Error(), 500)
+			http.Error(w, err.Error(), http.StatusInternalServerError)
 		}
 		return
 	}
 
-	if err := writeResponse(buildDir, name, format, output, w); err != nil {
+	if err := writeResponse(buildDir, name, imgConfig, format, output, w); err != nil {
 		log.Debugf("Writing response caused error: %s", err)
 		http.Error(w, err.Error(), 500)
 	}
 }
 
-func writeResponse(buildDir, name, format, output string, w io.Writer) error {
+func writeResponse(buildDir, name string, config *image.Config, format, output string, w io.Writer) error {
 	log.Debugf("write response buildDir: %s, name: %s, format: %s, output: %s", buildDir, name, format, output)
 	switch format {
 	case "rpi3":
@@ -65,7 +70,7 @@ func writeResponse(buildDir, name, format, output string, w io.Writer) error {
 				return errors.Wrap(err, "Error while unpacking rpi3 package")
 			}
 
-			if err := image.Build(tempDir, w); err != nil {
+			if err := image.Build(tempDir, config.Image.Partitions, w); err != nil {
 				return errors.Wrap(err, "Failed to build img file")
 			}
 		default:

pkg/image/builder.go

@@ -1,27 +1,27 @@
 package image
 
 import (
-	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
 	"os/exec"
+	"strconv"
 	"strings"
 
 	"github.com/ernoaapa/linuxkit-server/pkg/utils"
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 )
 
-func Build(sourceDir string, w io.Writer) error {
+func Build(sourceDir string, partitions []Partition, w io.Writer) error {
 	tmpfile, err := ioutil.TempFile("", "img")
 	if err != nil {
 		return errors.Wrapf(err, "Failed to create temporary file")
 	}
 	defer tmpfile.Close()
 	defer os.Remove(tmpfile.Name())
 
-	if err := buildImage(sourceDir, tmpfile.Name()); err != nil {
+	if err := buildImage(sourceDir, partitions, tmpfile.Name()); err != nil {
 		return err
 	}
 
@@ -32,11 +32,20 @@ func Build(sourceDir string, w io.Writer) error {
 	return nil
 }
 
-func buildImage(sourceDir, filename string) error {
-	if err := createZeroFile(filename, 1024*1024*100); err != nil {
+func buildImage(sourceDir string, partitions []Partition, filename string) error {
+	bootSize, err := utils.GetDirSize(sourceDir)
+	if err != nil {
+		return errors.Wrapf(err, "Failed to resolve image source size")
+	}
+
+	if err := validatePartitionTable(partitions, uint64(bootSize)); err != nil {
 		return err
 	}
-	if err := createFat32Partition(filename); err != nil {
+
+	if err := createZeroFile(filename, getTotalSize(partitions)); err != nil {
+		return err
+	}
+	if err := createPartitions(filename, partitions); err != nil {
 		return err
 	}
 
@@ -50,28 +59,40 @@ func buildImage(sourceDir, filename string) error {
 		return err
 	}
 
-	var device string
-	switch l := len(devices); l {
-	case 0:
-		return fmt.Errorf("%s don't have any paritions. There must be single partition", filename)
-	case 1:
-		device = devices[0]
-	default:
-		return fmt.Errorf("%s contain multiple paritions, but we support only single partition img files", filename)
-	}
+	for i, device := range devices {
+		partition := partitions[i]
+
+		switch partition.FsType {
+		case "fat32":
+			if err := formatFat32(device); err != nil {
+				return errors.Wrapf(err, "Failed to format device %s as Fat32", device)
+			}
+
+		case "ext4":
+			if err := formatExt4(device); err != nil {
+				return errors.Wrapf(err, "Failed to format device %s as ext4", device)
+			}
+		}
 
-	if err := formatFat32(device); err != nil {
-		return errors.Wrapf(err, "Failed to format device %s as Fat32", device)
+		if partition.Boot {
+			if err := writeBootPartition(device, sourceDir); err != nil {
+				return err
+			}
+		}
 	}
 
+	return nil
+}
+
+func writeBootPartition(device, sourceDir string) error {
 	buildDir, err := ioutil.TempDir("/mnt", "")
 	if err != nil {
 		return errors.Wrapf(err, "Failed to create temporary build directory")
 	}
 	defer os.RemoveAll(buildDir)
 
 	if err := mountDevice(device, buildDir); err != nil {
-		return errors.Wrapf(err, "Failed to mount device %s to dir %s", device, buildDir)
+		return errors.Wrapf(err, "Failed to mount root partition %s to dir %s", device, buildDir)
 	}
 	defer unmountDevice(buildDir)
 
@@ -82,30 +103,31 @@ func buildImage(sourceDir, filename string) error {
 	return nil
 }
 
-func createZeroFile(path string, size int64) error {
+func createZeroFile(path string, size uint64) error {
 	log.Debugf("Create %d bytes empty zero file to %s", size, path)
 	f, err := os.Create(path)
 	if err != nil {
 		return errors.Wrapf(err, "Failed to create empty file to %s", path)
 	}
 	defer f.Close()
 
-	if err := f.Truncate(size); err != nil {
+	if err := f.Truncate(int64(size)); err != nil {
 		return errors.Wrapf(err, "Failed to fill file %s to size %d", path, size)
 	}
 	return nil
 }
 
-func createFat32Partition(path string) error {
-	log.Debugf("Create Fat32 partition to %s", path)
+func createPartitions(path string, table []Partition) error {
+	log.Debugf("Create partitions to %s", path)
 	if err := runParted(path, "mklabel", "msdos"); err != nil {
 		return errors.Wrapf(err, "Failed to execute 'parted mklabel' to %s", path)
 	}
-	if err := runParted("--script", "--align=opt", path, "mkpart", "primary", "fat32", "2048s", "100%"); err != nil {
-		return errors.Wrapf(err, "Failed to execute 'parted mkpart' to %s", path)
-	}
-	if err := runParted(path, "set", "1", "boot", "on"); err != nil {
-		return errors.Wrapf(err, "Failed to 'parted set boot on' to %s", path)
+	for i, partition := range table {
+		args := []string{"--script", "--align=opt", path, "mkpart", "primary", partition.FsType, strconv.FormatUint(partition.Start, 10) + "B", strconv.FormatUint(partition.Start+partition.Size, 10) + "B"}
+		if err := runParted(args...); err != nil {
+			log.Debugf("Failed to execute (parted %s): %s", strings.Join(args, " "), err)
+			return errors.Wrapf(err, "Failed to create %d/%d partition to %s", i, len(table), path)
+		}
 	}
 	return nil
 }
@@ -123,6 +145,13 @@ func formatFat32(device string) error {
 	return cmd.Run()
 }
 
+func formatExt4(device string) error {
+	log.Debugf("Format device %s", device)
+	cmd := exec.Command("mkfs.ext4", device)
+	cmd.Stdout = os.Stdout
+	return cmd.Run()
+}
+
 func mountDevice(device, path string) error {
 	log.Debugf("Mount device %s to path %s", device, path)
 	cmd := exec.Command("mount", device, path)

pkg/image/config.go

@@ -0,0 +1,42 @@
+package image
+
+import (
+	"github.com/c2h5oh/datasize"
+	yaml "gopkg.in/yaml.v2"
+)
+
+// Config includes extra fields what linuxkit-server supports in the Linuxkit config
+type Config struct {
+	Image Image
+}
+
+// Image output configuration
+type Image struct {
+	Partitions []Partition
+}
+
+// Partition contains information about disk partition
+type Partition struct {
+	Boot   bool   `yaml:"boot"`
+	Start  uint64 `yaml:"start"`
+	Size   uint64 `yaml:"size"`
+	FsType string `yaml:"type"`
+}
+
+// NewConfig parses a config file
+func NewConfig(raw []byte) (*Config, error) {
+	var config = &Config{}
+
+	err := yaml.Unmarshal(raw, config)
+	if err != nil {
+		return config, err
+	}
+
+	if len(config.Image.Partitions) == 0 {
+		config.Image.Partitions = []Partition{
+			{Start: startOffset, Size: uint64(datasize.MB * 100), Boot: true, FsType: "fat32"},
+		}
+	}
+
+	return config, nil
+}

pkg/image/config_test.go

@@ -0,0 +1,35 @@
+package image
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewConfigWithNormalLinuxkitConfig(t *testing.T) {
+	_, err := NewConfig([]byte(`
+kernel:
+  image: linuxkit/kernel:4.9.69
+  cmdline: "console=tty0 console=ttyS0 console=ttyAMA0"
+trust:
+  org:
+    - linuxkit`))
+	assert.NoError(t, err)
+}
+
+func TestNewConfigWithExtraFields(t *testing.T) {
+	config, err := NewConfig([]byte(`
+image:
+  partitions:
+    - { start: 0, size: 104857600, boot: true, type: "fat32" }
+    - { start: 104857600, size: 209715200, type: "ext4" }
+kernel:
+  image: linuxkit/kernel:4.9.69
+  cmdline: "console=tty0 console=ttyS0 console=ttyAMA0"
+trust:
+  org:
+    - linuxkit`))
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(config.Image.Partitions))
+	assert.Equal(t, "fat32", config.Image.Partitions[0].FsType)
+}

pkg/image/partition.go

@@ -0,0 +1,73 @@
+package image
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+)
+
+// The default partition sector size
+var sectorSize = uint64(512)
+var startOffset = sectorSize * 2048
+var endOffset = sectorSize
+
+var supportedTypes = []string{"fat32", "ext4"}
+
+func validatePartitionTable(table []Partition, bootSize uint64) error {
+	if len(table) > 4 {
+		return errors.New("Only max 4 partitions are supported")
+	}
+
+	for _, partition := range table {
+		if !contains(supportedTypes, partition.FsType) {
+			return fmt.Errorf("%s is unsupported partition file system type. Supported types are %s", partition.FsType, strings.Join(supportedTypes, ","))
+		}
+	}
+
+	bootPartitions := getBootPartitions(table)
+	if len(bootPartitions) == 0 {
+		return errors.New("Partition table must have at least one boot partition")
+	} else if len(bootPartitions) > 1 {
+		return errors.New("Partition table can contain only one partition with boot=true")
+	}
+
+	boot := bootPartitions[0]
+
+	if boot.FsType != "fat32" {
+		return errors.New("Boot partition file system type must be fat32")
+	}
+
+	if boot.Size < bootSize {
+		return errors.New("Boot partition is not large enough for the boot files")
+	}
+
+	return nil
+}
+
+func contains(s []string, e string) bool {
+	for _, a := range s {
+		if a == e {
+			return true
+		}
+	}
+	return false
+}
+
+func getBootPartitions(table []Partition) (result []Partition) {
+	for _, p := range table {
+		if p.Boot {
+			result = append(result, p)
+		}
+	}
+	return result
+}
+
+func getTotalSize(table []Partition) uint64 {
+	var last Partition
+	for _, p := range table {
+		if last.Start < p.Start {
+			last = p
+		}
+	}
+	return last.Start + last.Size + sectorSize
+}