From d32407f43c123baaeba4beb2d8e04ee54e24ee86 Mon Sep 17 00:00:00 2001
From: Knut Ahlers <knut@ahlers.me>
Date: Mon, 15 Aug 2016 16:28:08 +0200
Subject: [PATCH] Add support for multiple Vault instances

---
 main.go | 112 +++++++++++++++++++++++++++++++++-----------------------
 1 file changed, 67 insertions(+), 45 deletions(-)

diff --git a/main.go b/main.go
index d8f7fd6..b184cff 100644
--- a/main.go
+++ b/main.go
@@ -3,19 +3,24 @@ package main // import "github.com/Jimdo/vault-unseal"
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"log"
 	"net/http"
 	"os"
+	"sync"
 	"time"
 
+	"golang.org/x/net/context"
+	"golang.org/x/net/context/ctxhttp"
+
 	"github.com/Luzifer/rconfig"
 )
 
 var config = struct {
-	OneShot       bool     `flag:"oneshot,1" default:"false" description:"Only try once and exit after"`
-	SealTokens    []string `flag:"tokens" default:"" description:"Tokens to try for unsealing the vault instance"`
-	VaultInstance string   `flag:"instance" env:"VAULT_ADDR" default:"http://127.0.0.1:8200" description:"Vault instance to unlock"`
-	Sleep         int      `flag:"sleep" default:"30" description:"How long to wait between sealed-state checks"`
+	OneShot        bool     `flag:"oneshot,1" default:"false" description:"Only try once and exit after"`
+	SealTokens     []string `flag:"tokens" default:"" description:"Tokens to try for unsealing the vault instance"`
+	VaultInstances []string `flag:"instance" env:"VAULT_ADDR" default:"http://127.0.0.1:8200" description:"Vault instance to unlock"`
+	Sleep          int      `flag:"sleep" default:"30" description:"How long to wait between sealed-state checks"`
 }{}
 
 func init() {
@@ -34,52 +39,21 @@ func init() {
 }
 
 func main() {
-	timedClient := http.DefaultClient
-	timedClient.Timeout = 10 * time.Second
+	var wg sync.WaitGroup
 
 	for {
-		s := sealStatus{}
-		r, err := timedClient.Get(config.VaultInstance + "/v1/sys/seal-status")
-		if err != nil {
-			log.Printf("An error ocurred while reading seal-status: %s\n", err)
-			os.Exit(1)
-		}
-		defer r.Body.Close()
+		for i := range config.VaultInstances {
+			wg.Add(1)
+			go func(i int) {
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 
-		if err := json.NewDecoder(r.Body).Decode(&s); err != nil {
-			log.Printf("Unable to decode seal-status: %s\n", err)
-			os.Exit(1)
-		}
+				defer wg.Done()
+				defer cancel()
 
-		if s.Sealed {
-			for _, token := range config.SealTokens {
-				log.Printf("Vault instance is sealed (missing %d tokens), trying to unlock...\n", s.T-s.Progress)
-				body := bytes.NewBuffer([]byte{})
-				json.NewEncoder(body).Encode(map[string]interface{}{
-					"key": token,
-				})
-				r, _ := http.NewRequest("PUT", config.VaultInstance+"/v1/sys/unseal", body)
-				resp, err := timedClient.Do(r)
-				if err != nil {
-					log.Printf("An error ocurred while doing unseal: %s\n", err)
-					os.Exit(1)
+				if err := unsealInstance(ctx, config.VaultInstances[i]); err != nil {
+					log.Printf("[ERR] %s", err)
 				}
-				defer resp.Body.Close()
-
-				if err := json.NewDecoder(resp.Body).Decode(&s); err != nil {
-					log.Printf("Unable to decode seal-status: %s\n", err)
-					os.Exit(1)
-				}
-
-				if !s.Sealed {
-					log.Printf("Unseal successfully finished.\n")
-					break
-				}
-			}
-
-			if s.Sealed {
-				log.Printf("Vault instance is still sealed (missing %d tokens), I don't have any more tokens.\n", s.T-s.Progress)
-			}
+			}(i)
 		}
 
 		if config.OneShot {
@@ -88,4 +62,52 @@ func main() {
 			<-time.After(time.Duration(config.Sleep) * time.Second)
 		}
 	}
+
+	wg.Wait()
+}
+
+func unsealInstance(ctx context.Context, instance string) error {
+	s := sealStatus{}
+	r, err := ctxhttp.Get(ctx, http.DefaultClient, instance+"/v1/sys/seal-status")
+	if err != nil {
+		return fmt.Errorf("[%s] An error ocurred while reading seal-status: %s", instance, err)
+	}
+	defer r.Body.Close()
+
+	if err := json.NewDecoder(r.Body).Decode(&s); err != nil {
+		return fmt.Errorf("[%s] Unable to decode seal-status: %s", instance, err)
+	}
+
+	if s.Sealed {
+		for _, token := range config.SealTokens {
+			log.Printf("[%s] Vault instance is sealed (missing %d tokens), trying to unlock...", instance, s.T-s.Progress)
+			body := bytes.NewBuffer([]byte{})
+			json.NewEncoder(body).Encode(map[string]interface{}{
+				"key": token,
+			})
+			r, _ := http.NewRequest("PUT", instance+"/v1/sys/unseal", body)
+			resp, err := ctxhttp.Do(ctx, http.DefaultClient, r)
+			if err != nil {
+				return fmt.Errorf("[%s] An error ocurred while doing unseal: %s", instance, err)
+			}
+			defer resp.Body.Close()
+
+			if err := json.NewDecoder(resp.Body).Decode(&s); err != nil {
+				return fmt.Errorf("[%s] Unable to decode seal-status: %s", instance, err)
+			}
+
+			if !s.Sealed {
+				log.Printf("[%s] Unseal successfully finished.", instance)
+				break
+			}
+		}
+
+		if s.Sealed {
+			log.Printf("[%s] Vault instance is still sealed (missing %d tokens), I don't have any more tokens.", instance, s.T-s.Progress)
+		}
+	} else {
+		log.Printf("[%s] Vault instance is already unsealed.", instance)
+	}
+
+	return nil
 }