From 662ab0edae0de81364a55617337b4dac204d2b95 Mon Sep 17 00:00:00 2001
From: zikaeroh <48577114+zikaeroh@users.noreply.github.com>
Date: Sat, 24 Dec 2022 13:34:58 -0800
Subject: [PATCH] Add support for devices

Signed-off-by: Zik Aeroh <48577114+zikaeroh@users.noreply.github.com>
---
 .../cmd/swarmctl/service/flagparser/device.go | 49 +++++++++++++++++++
 .../cmd/swarmctl/service/flagparser/flags.go  |  5 ++
 swarmd/dockerexec/container.go                |  9 ++++
 swarmd/dockerexec/container_test.go           | 32 ++++++++++++
 4 files changed, 95 insertions(+)
 create mode 100644 swarmd/cmd/swarmctl/service/flagparser/device.go

diff --git a/swarmd/cmd/swarmctl/service/flagparser/device.go b/swarmd/cmd/swarmctl/service/flagparser/device.go
new file mode 100644
index 0000000000..d82a44c048
--- /dev/null
+++ b/swarmd/cmd/swarmctl/service/flagparser/device.go
@@ -0,0 +1,49 @@
+package flagparser
+
+import (
+	"strings"
+
+	"github.com/moby/swarmkit/v2/api"
+	"github.com/pkg/errors"
+	"github.com/spf13/pflag"
+)
+
+func parseDevice(flags *pflag.FlagSet, spec *api.ServiceSpec) error {
+	if flags.Changed("device") {
+		container := spec.Task.GetContainer()
+		if container == nil {
+			return nil
+		}
+
+		devices, err := flags.GetStringSlice("device")
+		if err != nil {
+			return err
+		}
+
+		container.Devices = make([]*api.ContainerSpec_DeviceMapping, len(devices))
+
+		for i, device := range devices {
+			parts := strings.Split(device, ":")
+			if len(parts) < 1 || len(parts) > 3 {
+				return errors.Wrap(err, "failed to parse device")
+			}
+
+			mapping := &api.ContainerSpec_DeviceMapping{
+				PathOnHost:      parts[0],
+				PathInContainer: parts[0],
+			}
+
+			if len(parts) > 1 {
+				mapping.PathInContainer = parts[1]
+			}
+
+			if len(parts) == 3 {
+				mapping.CgroupPermissions = parts[2]
+			}
+
+			container.Devices[i] = mapping
+		}
+	}
+
+	return nil
+}
diff --git a/swarmd/cmd/swarmctl/service/flagparser/flags.go b/swarmd/cmd/swarmctl/service/flagparser/flags.go
index 7146519236..1da37d1c87 100644
--- a/swarmd/cmd/swarmctl/service/flagparser/flags.go
+++ b/swarmd/cmd/swarmctl/service/flagparser/flags.go
@@ -59,6 +59,7 @@ func AddServiceFlags(flags *pflag.FlagSet) {
 	flags.StringSlice("volume", nil, "define a volume mount")
 	flags.StringSlice("tmpfs", nil, "define a tmpfs mount")
 	flags.StringSlice("npipe", nil, "define a npipe mount")
+	flags.StringSlice("device", nil, "device options")
 
 	flags.String("log-driver", "", "specify a log driver")
 	flags.StringSlice("log-opt", nil, "log driver options, as key value pairs")
@@ -151,6 +152,10 @@ func Merge(cmd *cobra.Command, spec *api.ServiceSpec, c api.ControlClient) error
 		return err
 	}
 
+	if err := parseDevice(flags, spec); err != nil {
+		return err
+	}
+
 	driver, err := common.ParseLogDriverFlags(flags)
 	if err != nil {
 		return err
diff --git a/swarmd/dockerexec/container.go b/swarmd/dockerexec/container.go
index 2ab35f38ef..6f5ea80ac9 100644
--- a/swarmd/dockerexec/container.go
+++ b/swarmd/dockerexec/container.go
@@ -467,6 +467,15 @@ func (c *containerConfig) resources() enginecontainer.Resources {
 		}
 	}
 
+	resources.Devices = make([]enginecontainer.DeviceMapping, len(c.spec().Devices))
+	for i, device := range c.spec().Devices {
+		resources.Devices[i] = enginecontainer.DeviceMapping{
+			PathOnHost:        device.PathOnHost,
+			PathInContainer:   device.PathInContainer,
+			CgroupPermissions: device.CgroupPermissions,
+		}
+	}
+
 	// If no limits are specified let the engine use its defaults.
 	//
 	// TODO(aluzzardi): We might want to set some limits anyway otherwise
diff --git a/swarmd/dockerexec/container_test.go b/swarmd/dockerexec/container_test.go
index 0b9add675e..cba79c4505 100644
--- a/swarmd/dockerexec/container_test.go
+++ b/swarmd/dockerexec/container_test.go
@@ -332,3 +332,35 @@ func TestUlimits(t *testing.T) {
 		t.Fatalf("expected %v, got %v", expected, actual)
 	}
 }
+
+func TestDevices(t *testing.T) {
+	c := containerConfig{
+		task: &api.Task{
+			Spec: api.TaskSpec{
+				Runtime: &api.TaskSpec_Container{
+					Container: &api.ContainerSpec{
+						Devices: []*api.ContainerSpec_DeviceMapping{
+							{
+								PathOnHost:        "/dev/dri/card0",
+								PathInContainer:   "/dev/card0",
+								CgroupPermissions: "ro",
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	expected := []enginecontainer.DeviceMapping{
+		{
+			PathOnHost:        "/dev/dri/card0",
+			PathInContainer:   "/dev/card0",
+			CgroupPermissions: "ro",
+		},
+	}
+	actual := c.resources().Devices
+	if !reflect.DeepEqual(actual, expected) {
+		t.Fatalf("expected %v, got %v", expected, actual)
+	}
+}