will be merged: fix kms address parsing, add shuffle

Author: urykhy

kms.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"crypto/aes"
 	"io/ioutil"
+	"math/rand"
 
 	"encoding/base64"
 	"encoding/json"
@@ -51,38 +52,74 @@ func (c *Client) kmsAuth(url string) error {
 	return nil
 }
 
-func kmsSplit(keyProvider string) []string {
-	return strings.Split(keyProvider, ",")
-}
+// parse uri like kms://[email protected];kms02.example.com:9600/kms
+func kmsParseProviderUri(uri string) ([]string, error) {
+	original_uri := uri
 
-// kmsUrl parse KeyProviderUri to list of URL's
-func (c *Client) kmsUrl(einfo *hdfs.FileEncryptionInfoProto) ([]string, error) {
-	defaults, err := c.fetchDefaults()
-	if err != nil {
-		return nil, err
-	}
-
-	uri := defaults.GetKeyProviderUri()
 	if uri == "" {
-		return nil, errors.New("KeyProviderUri not configured on server")
+		return nil, errors.New("KeyProviderUri empty. not configured on server ?")
 	}
 
 	var urls []string
 	var proto string
 	if strings.HasPrefix(uri, kmsSchemeHTTPS) {
 		proto = "https://"
-		urls = kmsSplit(uri[len(kmsSchemeHTTPS):])
+		uri = uri[len(kmsSchemeHTTPS):]
 	}
 	if proto == "" && strings.HasPrefix(uri, kmsSchemeHTTP) {
 		proto = "http://"
-		urls = kmsSplit(uri[len(kmsSchemeHTTP):])
+		uri = uri[len(kmsSchemeHTTP):]
 	}
 	if proto == "" {
-		return nil, fmt.Errorf("not supported scheme %v", uri)
+		return nil, fmt.Errorf("not supported uri %v", original_uri)
+	}
+
+	port := ":9600" // default kms port
+	path := ""      // default path
+
+	parts := strings.Split(uri, ";")
+	for i, s := range parts {
+		path_index := strings.Index(s, "/")
+		if path_index > -1 {
+			path = s[path_index:]
+			s = s[:path_index]
+		}
+		port_index := strings.Index(s, ":")
+		if port_index > -1 {
+			port = s[port_index:]
+			s = s[:port_index]
+		}
+		if (path_index > -1 || port_index > -1) && i+1 != len(parts) {
+			return nil, fmt.Errorf("bad uri: %v", original_uri)
+		}
+		urls = append(urls, proto+s)
 	}
 
 	for i := range urls {
-		urls[i] = proto + urls[i] + "/v1/keyversion/" + url.QueryEscape(*einfo.EzKeyVersionName) + "/_eek?eek_op=decrypt"
+		urls[i] += port
+		urls[i] += path
+	}
+
+	return urls, nil
+}
+
+// kmsUrl parse KeyProviderUri to list of URL's
+func (c *Client) kmsUrl(einfo *hdfs.FileEncryptionInfoProto) ([]string, error) {
+	defaults, err := c.fetchDefaults()
+	if err != nil {
+		return nil, err
+	}
+
+	urls, err := kmsParseProviderUri(defaults.GetKeyProviderUri())
+	if err != nil {
+		return nil, err
+	}
+
+	// Reorder urls. Simple method to round robin calls across em.
+	rand.Shuffle(len(urls), func(i, j int) { urls[i], urls[j] = urls[j], urls[i] })
+
+	for i := range urls {
+		urls[i] = urls[i] + "/v1/keyversion/" + url.QueryEscape(*einfo.EzKeyVersionName) + "/_eek?eek_op=decrypt"
 	}
 
 	return urls, nil
@@ -156,7 +193,7 @@ func (c *Client) kmsGetKey(einfo *hdfs.FileEncryptionInfoProto) (*TransparentEnc
 
 	urls, err := c.kmsUrl(einfo)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "fail to get KMS address")
 	}
 
 	requestBody, err := json.Marshal(map[string]string{

kms_test.go

@@ -0,0 +1,53 @@
+package hdfs
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestKmsParseProviderUri(t *testing.T) {
+	assert.Equal(t, nil, nil)
+
+	urls, err := kmsParseProviderUri("")
+	assert.Error(t, err)
+
+	urls, err = kmsParseProviderUri("http")
+	assert.Error(t, err)
+
+	urls, err = kmsParseProviderUri("kms://https@localhost:9600/kms")
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(urls))
+	assert.Equal(t, "https://localhost:9600/kms", urls[0])
+
+	urls, err = kmsParseProviderUri("kms://[email protected]:9600;kms02.example.com")
+	assert.Error(t, err)
+
+	urls, err = kmsParseProviderUri("kms://[email protected]/kms;kms02.example.com")
+	assert.Error(t, err)
+
+	urls, err = kmsParseProviderUri("kms://[email protected];kms02.example.com:9600/kms")
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(urls))
+	assert.Equal(t, "http://kms01.example.com:9600/kms", urls[0])
+	assert.Equal(t, "http://kms02.example.com:9600/kms", urls[1])
+
+	urls, err = kmsParseProviderUri("kms://[email protected];kms02.example.com/kms")
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(urls))
+	assert.Equal(t, "http://kms01.example.com:9600/kms", urls[0])
+	assert.Equal(t, "http://kms02.example.com:9600/kms", urls[1])
+
+	urls, err = kmsParseProviderUri("kms://[email protected];kms02.example.com:9600")
+	assert.NoError(t, err)
+	assert.Equal(t, 2, len(urls))
+	assert.Equal(t, "http://kms01.example.com:9600", urls[0])
+	assert.Equal(t, "http://kms02.example.com:9600", urls[1])
+
+	urls, err = kmsParseProviderUri("kms://[email protected];kms02.example.com;kms03.example.com")
+	assert.NoError(t, err)
+	assert.Equal(t, 3, len(urls))
+	assert.Equal(t, "http://kms01.example.com:9600", urls[0])
+	assert.Equal(t, "http://kms02.example.com:9600", urls[1])
+	assert.Equal(t, "http://kms03.example.com:9600", urls[2])
+}