Merge pull request #69 from yuce/tls-support

Added TLS support

README.md

@@ -4,7 +4,7 @@
 <a href="https://godoc.org/github.com/pilosa/go-pilosa"><img src="https://godoc.org/github.com/pilosa/go-pilosa?status.svg" alt="GoDoc"></a>
 <a href="https://travis-ci.com/pilosa/go-pilosa"><img src="https://api.travis-ci.com/pilosa/go-pilosa.svg?token=vqssvEWV3KAhu8oVFx9s&branch=master"></a>
 <a href="https://goreportcard.com/report/github.com/pilosa/go-pilosa"><img src="https://goreportcard.com/badge/github.com/pilosa/go-pilosa?updated=1"></a>
-<a href="https://coveralls.io/github/pilosa/go-pilosa"><img src="https://coveralls.io/repos/github/pilosa/go-pilosa/badge.svg"></a>
+<a href="https://coveralls.io/github/pilosa/go-pilosa"><img src="https://coveralls.io/repos/github/pilosa/go-pilosa/badge.svg?updated=1"></a>
 
 <img src="https://www.pilosa.com/img/speed_sloth.svg" style="float: right" align="right" height="301">
 
@@ -22,6 +22,7 @@ Go client for Pilosa high performance distributed bitmap index.
     * Added support for excluding bits or attributes from bitmap calls. In order to exclude bits, pass `ExcludeBits: true` in your `QueryOptions`. In order to exclude attributes, pass `ExcludeAttrs: true`.
     * Added range field operations.
     * Customizable CSV timestamp format.
+    * `HTTPS connections are supported.
     * **Deprecation** Row and column labels are deprecated, and will be removed in a future release of this library. Do not use `ColumnLabel` option for `IndexOptions` and `RowLabel` for `FrameOption` for new code. See: https://github.com/pilosa/pilosa/issues/752 for more info.
 
 * **v0.5.0** (2017-08-03):

client.go

@@ -34,6 +34,7 @@ package pilosa
 
 import (
 	"bytes"
+	"crypto/tls"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -396,6 +397,7 @@ func (c *Client) importBits(indexName string, frameName string, slice uint64, bi
 		if err != nil {
 			return err
 		}
+		uri.SetScheme(node.Scheme)
 		err = c.importNode(uri, bitsToImportRequest(indexName, frameName, slice, bits))
 		if err != nil {
 			return err
@@ -416,6 +418,7 @@ func (c *Client) importValues(indexName string, frameName string, slice uint64,
 		if err != nil {
 			return err
 		}
+		uri.SetScheme(node.Scheme)
 		err = c.importValueNode(uri, valsToImportRequest(indexName, frameName, slice, fieldName, vals))
 		if err != nil {
 			return err
@@ -468,7 +471,7 @@ func (c *Client) ExportFrame(frame *Frame, view string) (BitIterator, error) {
 	if err != nil {
 		return nil, err
 	}
-	sliceURIs := statusToNodeSlicesForIndex(status, frame.index.Name())
+	sliceURIs := c.statusToNodeSlicesForIndex(status, frame.index.Name())
 	return NewCSVBitIterator(newExportReader(c, sliceURIs, frame, view)), nil
 }
 
@@ -598,6 +601,31 @@ func (c *Client) doRequest(host *URI, method, path string, headers map[string]st
 	return c.client.Do(req)
 }
 
+// statusToNodeSlicesForIndex finds the hosts which contains slices for the given index
+func (c *Client) statusToNodeSlicesForIndex(status *Status, indexName string) map[uint64]*URI {
+	// /status endpoint doesn't return the node scheme yet, default to the scheme of the current URI
+	// TODO: remove the following when /status endpoint returns the scheme for nodes
+	scheme := c.cluster.hosts[0].Scheme()
+	result := make(map[uint64]*URI)
+	for _, node := range status.Nodes {
+		for _, index := range node.Indexes {
+			if index.Name != indexName {
+				continue
+			}
+			for _, slice := range index.Slices {
+				uri, err := NewURIFromAddress(node.Host)
+				// err will always be nil, but prevent a panic in the odd chance the server returns an invalid URI
+				if err == nil {
+					uri.SetScheme(scheme)
+					result[slice] = uri
+				}
+			}
+			break
+		}
+	}
+	return result
+}
+
 func makeRequest(host *URI, method, path string, headers map[string]string, reader io.Reader) (*http.Request, error) {
 	request, err := http.NewRequest(method, host.Normalize()+path, reader)
 	if err != nil {
@@ -616,6 +644,7 @@ func newHTTPClient(options *ClientOptions) *http.Client {
 		Dial: (&net.Dialer{
 			Timeout: options.ConnectTimeout,
 		}).Dial,
+		TLSClientConfig:     options.TLSConfig,
 		MaxIdleConnsPerHost: options.PoolSizePerRoute,
 		MaxIdleConns:        options.TotalPoolSize,
 	}
@@ -685,33 +714,13 @@ func valsToImportRequest(indexName string, frameName string, slice uint64, field
 	}
 }
 
-// statusToNodeSlicesForIndex finds the hosts which contains slices for the given index
-func statusToNodeSlicesForIndex(status *Status, indexName string) map[uint64]*URI {
-	result := make(map[uint64]*URI)
-	for _, node := range status.Nodes {
-		for _, index := range node.Indexes {
-			if index.Name != indexName {
-				continue
-			}
-			for _, slice := range index.Slices {
-				uri, err := NewURIFromAddress(node.Host)
-				// err will always be nil, but prevent a panic in the odd chance the server returns an invalid URI
-				if err == nil {
-					result[slice] = uri
-				}
-			}
-			break
-		}
-	}
-	return result
-}
-
 // ClientOptions control the properties of client connection to the server
 type ClientOptions struct {
 	SocketTimeout    time.Duration
 	ConnectTimeout   time.Duration
 	PoolSizePerRoute int
 	TotalPoolSize    int
+	TLSConfig        *tls.Config
 }
 
 func (options *ClientOptions) withDefaults() (updated *ClientOptions) {
@@ -731,6 +740,9 @@ func (options *ClientOptions) withDefaults() (updated *ClientOptions) {
 	if updated.TotalPoolSize <= 100 {
 		updated.TotalPoolSize = 100
 	}
+	if updated.TLSConfig == nil {
+		updated.TLSConfig = &tls.Config{}
+	}
 	return
 }
 
@@ -753,6 +765,7 @@ const (
 )
 
 type fragmentNode struct {
+	Scheme       string
 	Host         string
 	InternalHost string
 }
@@ -768,6 +781,7 @@ type Status struct {
 
 // StatusNode contains node information.
 type StatusNode struct {
+	Scheme  string
 	Host    string
 	Indexes []StatusIndex
 }

client_it_test.go

@@ -47,6 +47,7 @@ import (
 	"testing"
 	"time"
 
+	"crypto/tls"
 	"github.com/golang/protobuf/proto"
 	pbuf "github.com/pilosa/go-pilosa/gopilosa_pbuf"
 )
@@ -1332,10 +1333,16 @@ func TestStatusToNodeSlicesForIndex(t *testing.T) {
 	// it needs to access statusToNodeSlicesForIndex which is not
 	// available to client_test.go
 
+	uri, err := NewURIFromAddress("https://:10101")
+	if err != nil {
+		t.Fatal(err)
+	}
+	client := NewClientWithURI(uri)
 	status := &Status{
 		Nodes: []StatusNode{
 			{
-				Host: ":10101",
+				Scheme: "https",
+				Host:   ":10101",
 				Indexes: []StatusIndex{
 					{
 						Name:   "index1",
@@ -1349,12 +1356,12 @@ func TestStatusToNodeSlicesForIndex(t *testing.T) {
 			},
 		},
 	}
-	sliceMap := statusToNodeSlicesForIndex(status, "index2")
+	sliceMap := client.statusToNodeSlicesForIndex(status, "index2")
 	if len(sliceMap) != 1 {
 		t.Fatalf("slice map should have a single item")
 	}
 	if uri, ok := sliceMap[0]; ok {
-		target, _ := NewURIFromAddress(":10101")
+		target, _ := NewURIFromAddress("https://:10101")
 		if !uri.Equals(target) {
 			t.Fatalf("slice map should have the correct URI")
 		}
@@ -1364,11 +1371,12 @@ func TestStatusToNodeSlicesForIndex(t *testing.T) {
 }
 
 func getClient() *Client {
-	uri, err := NewURIFromAddress(":10101")
+	uri, err := NewURIFromAddress("http://:10101")
 	if err != nil {
 		panic(err)
 	}
-	return NewClientWithURI(uri)
+	cluster := NewClusterWithHost(uri)
+	return NewClientWithCluster(cluster, &ClientOptions{TLSConfig: &tls.Config{InsecureSkipVerify: true}})
 }
 
 func getMockServer(statusCode int, response []byte, contentLength int) *httptest.Server {