From afce2948d216ddb11ba7413b19a65d5ce3d34b96 Mon Sep 17 00:00:00 2001
From: Brian Waldon <bcwaldon@gmail.com>
Date: Thu, 18 Sep 2014 07:50:52 -0700
Subject: [PATCH] etcd: replace -proxy-mode with -proxy

---
 Procfile     |  2 +-
 main.go      | 42 ++++++++++++++++++++++++++++++++++++------
 main_test.go | 23 +++++++++++++++++++++++
 3 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/Procfile b/Procfile
index 1ceece178..40f312945 100644
--- a/Procfile
+++ b/Procfile
@@ -2,4 +2,4 @@
 etcd1: ./etcd -id 0x1 -bind-addr 127.0.0.1:4001 -peer-bind-addr :7001 -peers '0x1=localhost:7001&0x2=localhost:7002&0x3=localhost:7003'
 etcd2: ./etcd -id 0x2 -bind-addr 127.0.0.1:4002 -peer-bind-addr :7002 -peers '0x1=localhost:7001&0x2=localhost:7002&0x3=localhost:7003'
 etcd3: ./etcd -id 0x3 -bind-addr 127.0.0.1:4003 -peer-bind-addr :7003 -peers '0x1=localhost:7001&0x2=localhost:7002&0x3=localhost:7003'
-proxy: ./etcd -proxy-mode -bind-addr 127.0.0.1:8080 -peers '0x1=localhost:7001&0x2=localhost:7002&0x3=localhost:7003'
+proxy: ./etcd -proxy=on -bind-addr 127.0.0.1:8080 -peers '0x1=localhost:7001&0x2=localhost:7002&0x3=localhost:7003'
diff --git a/main.go b/main.go
index 9ef258413..7afe112cd 100644
--- a/main.go
+++ b/main.go
@@ -24,6 +24,9 @@ import (
 const (
 	// the owner can make/remove files inside the directory
 	privateDirMode = 0700
+
+	proxyFlagValueOff = "off"
+	proxyFlagValueOn  = "on"
 )
 
 var (
@@ -31,18 +34,25 @@ var (
 	timeout   = flag.Duration("timeout", 10*time.Second, "Request Timeout")
 	paddr     = flag.String("peer-bind-addr", ":7001", "Peer service address (e.g., ':7001')")
 	dir       = flag.String("data-dir", "", "Path to the data directory")
-	proxyMode = flag.Bool("proxy-mode", false, "Forward HTTP requests to peers, do not participate in raft.")
 	snapCount = flag.Int64("snapshot-count", etcdserver.DefaultSnapCount, "Number of committed transactions to trigger a snapshot")
 
-	peers = &etcdhttp.Peers{}
-	addrs = &Addrs{}
+	peers     = &etcdhttp.Peers{}
+	addrs     = &Addrs{}
+	proxyFlag = new(ProxyFlag)
+
+	proxyFlagValues = []string{
+		proxyFlagValueOff,
+		proxyFlagValueOn,
+	}
 )
 
 func init() {
 	flag.Var(peers, "peers", "your peers")
 	flag.Var(addrs, "bind-addr", "List of HTTP service addresses (e.g., '127.0.0.1:4001,10.0.0.1:8080')")
+	flag.Var(proxyFlag, "proxy", fmt.Sprintf("Valid values include %s", strings.Join(proxyFlagValues, ", ")))
 	peers.Set("0x1=localhost:8080")
 	addrs.Set("127.0.0.1:4001")
+	proxyFlag.Set(proxyFlagValueOff)
 }
 
 func main() {
@@ -50,10 +60,10 @@ func main() {
 
 	setFlagsFromEnv()
 
-	if *proxyMode {
-		startProxy()
-	} else {
+	if string(*proxyFlag) == proxyFlagValueOff {
 		startEtcd()
+	} else {
+		startProxy()
 	}
 
 	// Block indefinitely
@@ -201,6 +211,26 @@ func (as *Addrs) String() string {
 	return strings.Join(*as, ",")
 }
 
+// ProxyFlag implements the flag.Value interface.
+type ProxyFlag string
+
+// Set verifies the argument to be a valid member of proxyFlagValues
+// before setting the underlying flag value.
+func (pf *ProxyFlag) Set(s string) error {
+	for _, v := range proxyFlagValues {
+		if s == v {
+			*pf = ProxyFlag(s)
+			return nil
+		}
+	}
+
+	return errors.New("invalid value")
+}
+
+func (pf *ProxyFlag) String() string {
+	return string(*pf)
+}
+
 // setFlagsFromEnv parses all registered flags in the global flagset,
 // and if they are not already set it attempts to set their values from
 // environment variables. Environment variables take the name of the flag but
diff --git a/main_test.go b/main_test.go
index f3c976cd1..dffcb0b08 100644
--- a/main_test.go
+++ b/main_test.go
@@ -41,3 +41,26 @@ func TestSetFlagsFromEnv(t *testing.T) {
 		}
 	}
 }
+
+func TestProxyFlagSet(t *testing.T) {
+	tests := []struct {
+		val  string
+		pass bool
+	}{
+		// known values
+		{"on", true},
+		{"off", true},
+
+		// unrecognized values
+		{"foo", false},
+		{"", false},
+	}
+
+	for i, tt := range tests {
+		pf := new(ProxyFlag)
+		err := pf.Set(tt.val)
+		if tt.pass != (err == nil) {
+			t.Errorf("#%d: want pass=%t, but got err=%v", i, tt.pass, err)
+		}
+	}
+}