From 9f7143cf445d21c4fa3d0e70f6a0ba16184c8cc7 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Tue, 14 Apr 2015 16:53:40 -0700
Subject: [PATCH] vault: expose the current leader

---
 vault/core.go      | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 vault/core_test.go | 36 ++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/vault/core.go b/vault/core.go
index 75f91b8fd..2745a4477 100644
--- a/vault/core.go
+++ b/vault/core.go
@@ -58,6 +58,10 @@ var (
 	// ErrInternalError is returned when we don't want to leak
 	// any information about an internal error
 	ErrInternalError = errors.New("internal error")
+
+	// ErrHANotEnabled is returned if the operation only makes sense
+	// in an HA setting
+	ErrHANotEnabled = errors.New("Vault is not configured for highly-available mode")
 )
 
 // SealConfig is used to describe the seal configuration
@@ -631,6 +635,54 @@ func (c *Core) Standby() (bool, error) {
 	return c.standby, nil
 }
 
+// Leader is used to get the current active leader
+func (c *Core) Leader() (bool, string, error) {
+	c.stateLock.RLock()
+	defer c.stateLock.RUnlock()
+	// Check if sealed
+	if c.sealed {
+		return false, "", ErrSealed
+	}
+
+	// Check if HA enabled
+	if c.ha == nil {
+		return false, "", ErrHANotEnabled
+	}
+
+	// Check if we are the leader
+	if !c.standby {
+		return true, c.advertiseAddr, nil
+	}
+
+	// Initialize a lock
+	lock, err := c.ha.LockWith(coreLockPath, "read")
+	if err != nil {
+		return false, "", err
+	}
+
+	// Read the value
+	held, value, err := lock.Value()
+	if err != nil {
+		return false, "", err
+	}
+	if !held {
+		return false, "", nil
+	}
+
+	// Value is the UUID of the leader, fetch the key
+	key := coreLeaderPrefix + value
+	entry, err := c.barrier.Get(key)
+	if err != nil {
+		return false, "", err
+	}
+	if entry == nil {
+		return false, "", nil
+	}
+
+	// Leader address is in the entry
+	return false, string(entry.Value), nil
+}
+
 // SealConfiguration is used to return information
 // about the configuration of the Vault and it's current
 // status.
diff --git a/vault/core_test.go b/vault/core_test.go
index f99ee75c7..8eeed9392 100644
--- a/vault/core_test.go
+++ b/vault/core_test.go
@@ -1035,6 +1035,18 @@ func TestCore_Standby(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
+	// Check the leader is local
+	isLeader, advertise, err := core.Leader()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if !isLeader {
+		t.Fatalf("should be leader")
+	}
+	if advertise != "foo" {
+		t.Fatalf("Bad advertise: %v", advertise)
+	}
+
 	// Create a second core, attached to same in-memory store
 	core2, err := NewCore(&CoreConfig{Physical: inm, AdvertiseAddr: "bar"})
 	if err != nil {
@@ -1068,6 +1080,18 @@ func TestCore_Standby(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
+	// Check the leader is not local
+	isLeader, advertise, err = core2.Leader()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if isLeader {
+		t.Fatalf("should not be leader")
+	}
+	if advertise != "foo" {
+		t.Fatalf("Bad advertise: %v", advertise)
+	}
+
 	// Seal the first core, should step down
 	err = core.Seal(root)
 	if err != nil {
@@ -1113,4 +1137,16 @@ func TestCore_Standby(t *testing.T) {
 	if resp.Data["foo"] != "bar" {
 		t.Fatalf("bad: %#v", resp)
 	}
+
+	// Check the leader is local
+	isLeader, advertise, err = core2.Leader()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if !isLeader {
+		t.Fatalf("should be leader")
+	}
+	if advertise != "bar" {
+		t.Fatalf("Bad advertise: %v", advertise)
+	}
 }