Skip to content

Commit

Permalink
Refactor the sidekick. Move etcd login into separate package to allow…
Browse files Browse the repository at this point in the history
… different backend.
  • Loading branch information
monder committed Nov 25, 2016
1 parent 790e2e8 commit 55e28c4
Show file tree
Hide file tree
Showing 6 changed files with 211 additions and 134 deletions.
2 changes: 1 addition & 1 deletion build-aci.sh
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
export CGO_ENABLED=0
export GOOS=linux
export GOARCH=amd64
export VERSION=v0.0.4
export VERSION=v0.1.0

go build -ldflags '-extldflags "-static"'
acbuild begin
Expand Down
26 changes: 18 additions & 8 deletions glide.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion glide.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,4 @@ import:
- package: golang.org/x/net
subpackages:
- context
- package: github.com/spf13/pflag
- package: github.com/mitchellh/cli
27 changes: 27 additions & 0 deletions lib/getIPAddress.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
package lib

import (
"fmt"
"net"
)

func GetIPAddress(cidr string) (string, error) {
_, requirednet, err := net.ParseCIDR(cidr)
if err != nil {
return "", err
}
addrs, err := net.InterfaceAddrs()
if err != nil {
return "", err
}

for _, addr := range addrs {
if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
if requirednet.Contains(ipnet.IP) {
return ipnet.IP.String(), nil
}
}

}
return "", fmt.Errorf("No IP address matching %s found.", cidr)
}
134 changes: 10 additions & 124 deletions main.go
Original file line number Diff line number Diff line change
@@ -1,139 +1,25 @@
package main

import (
"fmt"
"github.com/coreos/etcd/client"
"github.com/spf13/pflag"
"golang.org/x/net/context"
"log"
"net"
"os"
"os/signal"
"strings"
"syscall"
"time"
)

var flags struct {
etcdAddress string
cidr string
format string
expireDir string
interval time.Duration
}
"github.com/mitchellh/cli"

func init() {
pflag.StringVarP(&flags.etcdAddress, "etcd-endpoint", "e", "http://172.16.28.1:2379", "an etcd address in the cluster")
pflag.StringVar(&flags.cidr, "cidr", "0.0.0.0/0", "cidr to match the ip")
pflag.StringVarP(&flags.format, "format", "f", "$ip", "format of the etcd key value. '$ip' will be replace by container's ip address")
pflag.StringVar(&flags.expireDir, "expireDir", "", "set expiration TTLs for all items under that directory, not only the leaf node")
pflag.DurationVarP(&flags.interval, "interval", "i", time.Minute, "refresh interval")
pflag.Usage = func() {
fmt.Fprintf(os.Stderr, "Usage:\n %s /key/in/etcd\n\nFlags:\n", os.Args[0])
pflag.PrintDefaults()
}
}
"github.com/monder/rkt-sidekick/modes"
)

func main() {
pflag.Parse()
if pflag.Arg(0) == "" {
pflag.Usage()
os.Exit(2)
}

sigc := make(chan os.Signal, 1)
signal.Notify(sigc, syscall.SIGTERM, syscall.SIGINT)

etcd, err := getEtcd(strings.Split(flags.etcdAddress, ","))
if err != nil {
log.Fatal(err)
}
ip, err := getIPAddress(flags.cidr)
if err != nil {
log.Fatal(err)
}

value := strings.Replace(flags.format, "$ip", ip, -1)

_, err = etcd.Set(context.Background(), pflag.Arg(0), value, &client.SetOptions{})
if err != nil {
log.Fatal(err)
c := cli.NewCLI("rkt-sidekick", "0.1.0")
c.Args = os.Args[1:]
c.Commands = map[string]cli.CommandFactory{
"etcd": modes.EtcdCommand,
}

ticker := time.NewTicker(flags.interval)

for {
path := pflag.Arg(0)
keepRoot := flags.expireDir
if keepRoot == "" {
keepRoot = path[0:strings.LastIndex(path, "/")]
}
if !strings.HasSuffix(keepRoot, "/") {
keepRoot += "/"
}
currentPath := "/"
if strings.HasPrefix(path, keepRoot) {
path = path[len(keepRoot):]
currentPath = keepRoot
} else {
keepRoot = currentPath
path = path[len(currentPath):]
}

for _, s := range strings.SplitAfter(path, "/") {
currentPath += s
fmt.Printf("Setting ttl for %s (dir: %s)\n", currentPath, strings.HasSuffix(currentPath, "/"))
_, err = etcd.Set(context.Background(), currentPath, "", &client.SetOptions{
Refresh: true,
PrevExist: client.PrevExist,
TTL: 2 * flags.interval,
Dir: strings.HasSuffix(currentPath, "/"),
})
if err != nil {
log.Fatal(err)
}
}

select {
case <-sigc:
etcd.Delete(context.Background(), pflag.Arg(0), nil)
//TODO clear empty dirs
os.Exit(0)
case <-ticker.C:
}
}
}

func getEtcd(endpoints []string) (client.KeysAPI, error) {
cfg := client.Config{
Endpoints: endpoints,
Transport: client.DefaultTransport,
}
c, err := client.New(cfg)
exitStatus, err := c.Run()
if err != nil {
return nil, err
log.Println(err)
}
kapi := client.NewKeysAPI(c)
return kapi, nil
}

func getIPAddress(cidr string) (string, error) {
_, requirednet, err := net.ParseCIDR(cidr)
if err != nil {
return "", err
}
addrs, err := net.InterfaceAddrs()
if err != nil {
return "", err
}

for _, addr := range addrs {
if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
if requirednet.Contains(ipnet.IP) {
return ipnet.IP.String(), nil
}
}

}
return "", fmt.Errorf("No IP address matching %s found.", cidr)
os.Exit(exitStatus)
}
154 changes: 154 additions & 0 deletions modes/etcd.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
package modes

import (
"flag"
"fmt"
"github.com/coreos/etcd/client"
"github.com/mitchellh/cli"
"github.com/monder/rkt-sidekick/lib"
"golang.org/x/net/context"
"log"
"os"
"os/signal"
"strings"
"syscall"
"time"
)

type command struct {
etcdAddress string
cidr string
format string
expireDir string
interval time.Duration
}

var EtcdCommand = func() (cli.Command, error) {
return &command{}, nil
}

func (c command) Help() string {
return `
Usage: rkt-sidekick [options] [KEY_IN_ETCD]
Options:
-cidr cidr to match the ip (default: "0.0.0.0/0")
-etcd-endpoint an etcd address in the cluster (default: "http://172.16.28.1:2379")
-format format of the etcd key value. '$ip' will be replace by container's ip address (default: "$ip")
-expire-dir set expiration TTLs for all items under that directory, not only the leaf node
-interval refresh interval (default: "1m")
`
}

func (c command) Run(args []string) int {
cmdFlags := flag.NewFlagSet("etcd", flag.ContinueOnError)
cmdFlags.Usage = func() { fmt.Fprintf(os.Stderr, "%s", c.Help()) }

cmdFlags.StringVar(&c.etcdAddress, "etcd-endpoint", "http://172.16.28.1:2379", "")
cmdFlags.StringVar(&c.cidr, "cidr", "0.0.0.0/0", "")
cmdFlags.StringVar(&c.format, "format", "$ip", "")
cmdFlags.StringVar(&c.expireDir, "expire-dir", "", "")
cmdFlags.DurationVar(&c.interval, "interval", time.Minute, "")

if err := cmdFlags.Parse(args); err != nil {
return 1
}

key := ""

// Check for arg validation
args = cmdFlags.Args()
switch len(args) {
case 0:
key = ""
case 1:
key = args[0]
default:
fmt.Fprintf(os.Stderr, "Too many arguments (expected 1, got %d)\n", len(args))
return 1
}
if key == "" {
fmt.Fprintf(os.Stderr, "Error! Missing KEY argument\n")
return 1
}

sigc := make(chan os.Signal, 1)
signal.Notify(sigc, syscall.SIGTERM, syscall.SIGINT)

etcd, err := getEtcd(strings.Split(c.etcdAddress, ","))
if err != nil {
log.Fatal(err)
}
ip, err := lib.GetIPAddress(c.cidr)
if err != nil {
log.Fatal(err)
}

value := strings.Replace(c.format, "$ip", ip, -1)

_, err = etcd.Set(context.Background(), key, value, &client.SetOptions{})
if err != nil {
log.Fatal(err)
}

ticker := time.NewTicker(c.interval)

for {
path := key
keepRoot := c.expireDir
if keepRoot == "" {
keepRoot = path[0:strings.LastIndex(path, "/")]
}
if !strings.HasSuffix(keepRoot, "/") {
keepRoot += "/"
}
currentPath := "/"
if strings.HasPrefix(path, keepRoot) {
path = path[len(keepRoot):]
currentPath = keepRoot
} else {
keepRoot = currentPath
path = path[len(currentPath):]
}

for _, s := range strings.SplitAfter(path, "/") {
currentPath += s
fmt.Printf("Setting ttl for %s (dir: %s)\n", currentPath, strings.HasSuffix(currentPath, "/"))
_, err = etcd.Set(context.Background(), currentPath, "", &client.SetOptions{
Refresh: true,
PrevExist: client.PrevExist,
TTL: 2 * c.interval,
Dir: strings.HasSuffix(currentPath, "/"),
})
if err != nil {
log.Fatal(err)
}
}

select {
case <-sigc:
etcd.Delete(context.Background(), key, nil)
//TODO clear empty dirs
os.Exit(0)
case <-ticker.C:
}
}
return 0
}

func getEtcd(endpoints []string) (client.KeysAPI, error) {
cfg := client.Config{
Endpoints: endpoints,
Transport: client.DefaultTransport,
}
c, err := client.New(cfg)
if err != nil {
return nil, err
}
kapi := client.NewKeysAPI(c)
return kapi, nil
}

func (c command) Synopsis() string {
return "Updates a key in etcd with container IP address"
}

0 comments on commit 55e28c4

Please sign in to comment.