First version

Author: tr4cks

cmd/root.go

@@ -1,26 +1,42 @@
 package cmd
 
 import (
+"bytes"
 "context"
+yamlConfig "docker-netns/config"
 "fmt"
 "github.com/docker/docker/api/types"
 "github.com/docker/docker/api/types/filters"
 "github.com/docker/docker/client"
 "github.com/kardianos/service"
+"github.com/kballard/go-shellquote"
 "github.com/spf13/cobra"
+"github.com/vishvananda/netns"
 "log"
 "os"
+"os/exec"
+"path"
+"runtime"
+"syscall"
 )
 
+func init() {
+rootCmd.Flags().StringVar(&configPath, "config", path.Join("/opt", appName, "config.yaml"), "YAML configuration file")
+}
+
+const appName = "docker-netns"
+
 var (
-logger service.Logger
-rootCmd = &cobra.Command{
-Use: "docker-netns",
-Short: "Docker network namespace manager",
+configPath string
+logger service.Logger
+rootCmd = &cobra.Command{
+Use: appName,
+Short: "Docker network namespace manager",
 Version: "1.0.0",
 Run: func(cmd *cobra.Command, args []string) {
-prg := &program{}
+prg := NewProgram()
 s, err := service.New(prg, svcConfig)
+prg.service = s
 if err != nil {
 fmt.Fprintf(os.Stderr, "[service.New] %v\n", err)
 os.Exit(1)
@@ -33,13 +49,14 @@ var (
 err = s.Run()
 if err != nil {
 logger.Error(err)
+os.Exit(1)
 }
 },
 }
 svcConfig = &service.Config{
-Name: "docker-netns",
+Name:  "docker-netns",
 DisplayName: "Docker network namespace service",
-UserName: "root",
+UserName:  "root",
 Dependencies: []string{
 "After=network.target syslog.target docker.service",
 },
@@ -49,37 +66,140 @@ var (
 }
 )
 
-type program struct {}
+func nsContext(containerID string, callback func() (interface{}, error)) (interface{}, error) {
+runtime.LockOSThread()
+defer runtime.UnlockOSThread()
+originNamespace, err := netns.Get()
+if err != nil {
+return nil, err
+}
+defer originNamespace.Close()
+cli, err := client.NewClientWithOpts(client.WithVersion("1.40"))
+if err != nil {
+return nil, err
+}
+container, err := cli.ContainerInspect(context.Background(), containerID)
+if err != nil {
+return nil, err
+}
+ContainerNamespace, err := netns.GetFromPid(container.State.Pid)
+if err != nil {
+return nil, err
+}
+defer ContainerNamespace.Close()
+err = netns.Set(ContainerNamespace)
+if err != nil {
+return nil, err
+}
+res, err := callback()
+if err != nil {
+return res, err
+}
+err = netns.Set(originNamespace)
+if err != nil {
+return nil, err
+}
+return res, nil
+}
+
+func execCommands(containerID string, commands []string) error {
+_, err := nsContext(containerID, func() (interface{}, error) {
+for _, command := range commands {
+command, err := shellquote.Split(command)
+if err != nil {
+return nil, err
+}
+execCmd := exec.Command(command[0], command[1:]...)
+var stderr bytes.Buffer
+execCmd.Stderr = &stderr
+err = execCmd.Run()
+if err != nil {
+return nil, fmt.Errorf("%v: %v", err, stderr.String())
+}
+}
+return nil, nil
+})
+return err
+}
+
+type program struct {
+service service.Service
+ctx context.Context
+cancel context.CancelFunc
+exited chan error
+}
+
+func NewProgram() *program {
+ctx, cancel := context.WithCancel(context.Background())
+exited := make(chan error)
+return &program{ctx: ctx, cancel: cancel, exited: exited}
+}
 
 func (p *program) Start(s service.Service) error {
 // Start should not block. Do the actual work async.
+config, err := yamlConfig.NewConfig(configPath)
+if err != nil {
+return err
+}
 cli, err := client.NewClientWithOpts(client.WithVersion("1.40"))
 if err != nil {
 return err
 }
-go p.run(cli)
+go p.run(config, cli)
 return nil
 }
 
-func (p *program) run(cli *client.Client) {
-ctx := context.Background()
+func (p *program) run(config *yamlConfig.Config, cli *client.Client) {
+defer close(p.exited)
+process, _ := os.FindProcess(os.Getpid())
+containers, err := cli.ContainerList(context.Background(), types.ContainerListOptions{})
+if err != nil {
+process.Signal(syscall.SIGTERM)
+p.exited <- err
+return
+}
+for containerID := range *config {
+for _, container := range containers {
+if containerID == container.ID[:len(containerID)] {
+err := execCommands(containerID, (*config)[containerID])
+if err != nil {
+process.Signal(syscall.SIGTERM)
+p.exited <- err
+return
+}
+break
+}
+}
+}
+
 f := filters.NewArgs()
 f.Add("event", "start")
-msgs, errs := cli.Events(ctx, types.EventsOptions{Filters: f})
+msgs, errs := cli.Events(p.ctx, types.EventsOptions{Filters: f})
 for {
 select {
-case msg := <-msgs:
-_ = msg
+case <-p.ctx.Done():
+return
+case <-msgs:
+for containerID := range *config {
+err := execCommands(containerID, (*config)[containerID])
+if err != nil {
+process.Signal(syscall.SIGTERM)
+p.exited <- err
+return
+}
+}
 case err := <-errs:
-_ = err
+process.Signal(syscall.SIGTERM)
+p.exited <- err
 return
 }
 }
 }
 
 func (p *program) Stop(s service.Service) error {
 // Stop should not block. Return with a few seconds.
-return nil
+p.cancel()
+return <-p.exited
 }
 
 func Execute() {

cmd/service.go

@@ -4,30 +4,88 @@ import (
 "fmt"
 "github.com/kardianos/service"
 "github.com/spf13/cobra"
+"io/ioutil"
 "os"
-"strings"
+"path"
 )
 
 func init() {
-serviceCmd.Flags().StringP("action", "a", "", strings.Join(service.ControlAction[:], "|"))
-serviceCmd.MarkFlagRequired("action")
+installCmd.Flags().String("config", "./config.yaml", "YAML configuration file")
+serviceCmd.AddCommand(startCmd, stopCmd, restartCmd, installCmd, uninstallCmd)
 rootCmd.AddCommand(serviceCmd)
 }
 
-var serviceCmd = &cobra.Command{
-Use: "service",
-Short: "Service management",
-Run: func(cmd *cobra.Command, args []string) {
-prg := &program{}
+func serviceControl(action string) func(*cobra.Command, []string) {
+return func(cmd *cobra.Command, args []string) {
+prg := NewProgram()
 s, err := service.New(prg, svcConfig)
+prg.service = s
 if err != nil {
 fmt.Fprintf(os.Stderr, "[service.New] %v\n", err)
 os.Exit(1)
 }
-err = service.Control(s, cmd.Flag("action").Value.String())
+err = service.Control(s, action)
 if err != nil {
 fmt.Fprintf(os.Stderr, "[service.Control] %v\n", err)
 os.Exit(1)
 }
-},
+}
 }
+
+var (
+serviceCmd = &cobra.Command{
+Use: "service",
+Short: "Service management",
+}
+startCmd = &cobra.Command{
+Use: "start",
+Short: "Start service",
+Run: serviceControl("start"),
+}
+stopCmd = &cobra.Command{
+Use: "stop",
+Short: "Stop service",
+Run: serviceControl("stop"),
+}
+restartCmd = &cobra.Command{
+Use: "restart",
+Short: "Restart service",
+Run: serviceControl("restart"),
+}
+installCmd = &cobra.Command{
+Use: "install",
+Short: "Install service",
+Run: func(cmd *cobra.Command, args []string) {
+// TODO: transaction
+err := os.Mkdir(path.Join("/opt", appName), 0755)
+if err != nil {
+fmt.Fprintf(os.Stderr, "[os.Mkdir] %v\n", err)
+os.Exit(1)
+}
+bytes, err := ioutil.ReadFile(cmd.Flag("config").Value.String())
+if err != nil {
+fmt.Fprintf(os.Stderr, "[ioutil.ReadFile] %v\n", err)
+os.Exit(1)
+}
+err = ioutil.WriteFile(path.Join("/opt", appName, "config.yaml"), bytes, 0644)
+if err != nil {
+fmt.Fprintf(os.Stderr, "[ioutil.WriteFile] %v\n", err)
+os.Exit(1)
+}
+serviceControl("install")(cmd, args)
+},
+}
+uninstallCmd = &cobra.Command{
+Use: "uninstall",
+Short: "Uninstall service",
+Run: func(cmd *cobra.Command, args []string) {
+// TODO: transaction
+err := os.RemoveAll(path.Join("/opt", appName))
+if err != nil {
+fmt.Fprintf(os.Stderr, "[ioutil.RemoveAll] %v\n", err)
+os.Exit(1)
+}
+serviceControl("uninstall")(cmd, args)
+},
+}
+)

cmd/shell.go

@@ -1,14 +1,10 @@
 package cmd
 
 import (
-"context"
 "fmt"
-"github.com/docker/docker/client"
 "github.com/spf13/cobra"
-"github.com/vishvananda/netns"
 "os"
 "os/exec"
-"runtime"
 )
 
 func init() {
@@ -21,49 +17,20 @@ var shellCmd = &cobra.Command{
 Use: "shell",
 Short: "Start a shell in a specific network namespace",
 Run: func(cmd *cobra.Command, args []string) {
-runtime.LockOSThread()
-defer runtime.UnlockOSThread()
-originNamespace, err := netns.Get()
-if err != nil {
-fmt.Fprintf(os.Stderr, "[netns.Get] %v\n", err)
-os.Exit(1)
-}
-defer originNamespace.Close()
-cli, err := client.NewClientWithOpts(client.WithVersion("1.40"))
-if err != nil {
-fmt.Fprintf(os.Stderr, "[client.NewClientWithOpts] %v\n", err)
-os.Exit(1)
-}
-container, err := cli.ContainerInspect(context.Background(), cmd.Flag("container").Value.String())
-if err != nil {
-fmt.Fprintf(os.Stderr, "[cli.ContainerInspect] %v\n", err)
-os.Exit(1)
-}
-println(container.State.Pid)
-ContainerNamespace, err := netns.GetFromPid(container.State.Pid)
-if err != nil {
-fmt.Fprintf(os.Stderr, "[netns.GetFromPid] %v\n", err)
-os.Exit(1)
-}
-defer ContainerNamespace.Close()
-err = netns.Set(ContainerNamespace)
-if err != nil {
-fmt.Fprintf(os.Stderr, "[netns.Set] %v\n", err)
-os.Exit(1)
-}
-shell := exec.Command(os.Getenv("SHELL"))
-shell.Stdin = os.Stdin
-shell.Stdout = os.Stdout
-shell.Stderr = os.Stderr
-err = shell.Run()
+_, err := nsContext(cmd.Flag("container").Value.String(), func() (interface{}, error) {
+shell := exec.Command(os.Getenv("SHELL"))
+shell.Stdin = os.Stdin
+shell.Stdout = os.Stdout
+shell.Stderr = os.Stderr
+err := shell.Run()
+if err != nil {
+return nil, err
+}
+return nil, nil
+})
 if err != nil {
 fmt.Fprintf(os.Stderr, "[shell.Run] %v\n", err)
 os.Exit(1)
 }
-err = netns.Set(originNamespace)
-if err != nil {
-fmt.Fprintf(os.Stderr, "[netns.Set] %v\n", err)
-os.Exit(1)
-}
 },
 }

config.yaml

@@ -1,3 +1,2 @@
-d8b71122deaa:
- - ip route add default via 192.168.1.254 dev eth0
-
+360f9c13f2ab9:
+ - iptables -P OUTPUT DROP