foo.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
	"os"
	"time"
)

// Stolen from: https://github.com/digitalocean/digitalocean-cloud-controller-manager/blob/5a64f9a0729ece886c65c59dcc2a0ecbe7f3b6eb/do/cloud.go#L47
type cloud struct {
	client *godo.Client
}

// Stolen from: https://github.com/kokhang/rook/commit/a71ebe90949bdd653108d5918bfb3c9656f738a0#diff-da8e79a24ddf6dde0e271796cfbd850bR35
type Result struct {
	// Status of the callout. One of "Success", "Failure" or "Not supported".
	Status string `json:"status"`
	// Reason for success/failure.
	Message string `json:"message,omitempty"`
	// Path to the device attached. This field is valid only for attach calls.
	// ie: /dev/sdx
	DevicePath string `json:"device,omitempty"`
	// Cluster wide unique name of the volume.
	VolumeName string `json:"volumeName,omitempty"`
	// Represents volume is attached on the node
	Attached bool `json:"attached,omitempty"`
	// Returns capabilities of the driver.
	// By default we assume all the capabilities are supported.
	// If the plugin does not support a capability, it can return false for that capability.
	Capabilities map[string]bool
}

// Stripped down version of: https://github.com/kokhang/rook/commit/a71ebe90949bdd653108d5918bfb3c9656f738a0#diff-da8e79a24ddf6dde0e271796cfbd850bR24
type AttachOptions struct {
	RW         string `json:"kubernetes.io/readwrite"`
	FsType     string `json:"kubernetes.io/fsType"`
	VolumeName string `json:"kubernetes.io/pvOrVolumeName"`
}

type tokenSource struct {
	AccessToken string
}

func (t *tokenSource) Token() (*oauth2.Token, error) {
	token := &oauth2.Token{
		AccessToken: t.AccessToken,
	}
	return token, nil
}

func (c *cloud) findDropletID(nodeName string) (int, error) {
	opt := &godo.ListOptions{}

	// Stolen from: https://github.com/digitalocean/godo/blob/master/README.md
	for {
		droplets, resp, err := c.client.Droplets.List(context.TODO(), opt)
		if err != nil {
			return 0, err
		}

		for _, droplet := range droplets {
			if droplet.Name == nodeName {
				return droplet.ID, nil
			}
			for _, n := range droplet.Networks.V4 {
				if n.IPAddress == nodeName {
					return droplet.ID, nil
				}
			}
		}

		if resp.Links == nil || resp.Links.IsLastPage() {
			break
		}

		page, err := resp.Links.CurrentPage()
		if err != nil {
			return 0, err
		}
		opt.Page = page + 1
	}

	return 0, errors.New("No matching droplet found")
}

func (c *cloud) waitForAction(actionID int) {
	for {
		// TODO: Check err (?)
		a, _, _ := c.client.Actions.Get(context.TODO(), actionID)
		if a.Status == godo.ActionCompleted {
			break
		}
		time.Sleep(1 * time.Second)
	}
}

func (c *cloud) attach(args []string) (Result, error) {
	var a AttachOptions
	if err := json.Unmarshal([]byte(args[0]), a); err != nil {
		return Result{}, err
	}
	nodeName := args[1]

	dropletID, err := c.findDropletID(nodeName)
	if err != nil {
		return Result{}, err
	}

	// TODO: Detach if attached to another node
	// TODO: Check if already attached
	action, _, err := c.client.StorageActions.Attach(context.TODO(), a.VolumeName, dropletID)
	if err != nil {
		return Result{}, err
	}
	c.waitForAction(action.ID)

	r := Result{
		Status:     "Success",
		DevicePath: "/dev/disk/by-id/scsi-0DO_Volume_" + a.VolumeName,
	}
	return r, nil
}

func (c *cloud) detach(args []string) (Result, error) {
	volumeName := args[0]
	nodeName := args[1]

	dropletID, err := c.findDropletID(nodeName)
	if err != nil {
		return Result{}, nil
	}

	v, _, err := c.client.Storage.GetVolume(context.TODO(), volumeName)
	if err != nil {
		return Result{}, nil
	}
	if len(v.DropletIDs) == 1 && v.DropletIDs[0] == dropletID {
		action, _, err := c.client.StorageActions.DetachByDropletID(context.TODO(), v.ID, dropletID)
		if err != nil {
			return Result{}, nil
		}
		c.waitForAction(action.ID)
	}

	r := Result{
		Status: "Success",
	}
	return r, nil
}

func (c *cloud) waitForAttach(args []string) (Result, error) {
	mountDev := args[0]
	for {
		if _, err := os.Stat(mountDev); !os.IsNotExist(err) {
			break
		}
		time.Sleep(1 * time.Second)
	}

	r := Result{
		Status:     "Success",
		DevicePath: mountDev,
	}
	return r, nil
}

func (c *cloud) isAttached(args []string) (Result, error) {
	var a AttachOptions
	if err := json.Unmarshal([]byte(args[0]), a); err != nil {
		return Result{}, nil
	}
	nodeName := args[1]

	v, _, err := c.client.Storage.GetVolume(context.TODO(), a.VolumeName)
	if err != nil {
		return Result{}, nil
	}

	dropletID, err := c.findDropletID(nodeName)
	if err != nil {
		return Result{}, nil
	}

	r := Result{
		Status:   "Success",
		Attached: len(v.DropletIDs) == 1 && v.DropletIDs[0] == dropletID,
	}
	return r, nil
}

func main() {
	// TODO: Pull from env or file
	tokenSource := &tokenSource{
		AccessToken: "xxxx",
	}

	oauthClient := oauth2.NewClient(context.Background(), tokenSource)

	c := &cloud{
		client: godo.NewClient(oauthClient),
	}

	r := func() (Result, error) {
		// TODO: Check number of args
		switch args := os.Args; args[1] {
		case "init":
			return Result{
				Status: "Success",
				Capabilities: map[string]bool{
					"attach": true,
				},
			}, nil
		case "attach":
			return c.attach(args[2:])
		case "detach":
			return c.detach(args[2:])
		case "waitforattach":
			return c.waitForAttach(args[2:])
		case "isattached":
			return c.isAttached(args[2:])
		default:
			return Result{
				Status: "Not supported",
			}, nil
		}
		return Result{}, nil
	}
	result, err := r()
	if err != nil {
		result.Status = "Failure"
		result.Message = err.Error()
	}
	fmt.Println(json.Marshal(result))
}