Merge pull request #48 from zegl/riak-security

#47 - Added support for password based authentication

client.go

@@ -2,11 +2,18 @@ package goriak
 
 import (
 	riak "github.com/basho/riak-go-client"
+
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"io/ioutil"
+	"strconv"
+	"strings"
 )
 
 // Session holds the connection to Riak
 type Session struct {
-	riak *riak.Client
+	riak *riak.Cluster
 	opts ConnectOpts
 }
 
@@ -15,6 +22,16 @@ type ConnectOpts struct {
 	// Both Address and Addresses should be on the form HOST|IP[:PORT]
 	Address   string   // Address to a single Riak host. Will be used in case Addresses is empty
 	Addresses []string // Addresses to all Riak hosts.
+
+	// Username and password for connection to servers with secirity enabled
+	User     string
+	Password string
+
+	// Path to root CA certificate. Required if security is used
+	CARootCert string
+
+	// Option to override port. Is set to 8087 by default
+	Port uint32
 }
 
 // Connect creates a new Riak connection. See ConnectOpts for the available options.
@@ -37,10 +54,71 @@ func (c *Session) connect() error {
 		c.opts.Addresses = []string{c.opts.Address}
 	}
 
-	con, err := riak.NewClient(&riak.NewClientOptions{
-		RemoteAddresses: c.opts.Addresses,
+	var authOptions *riak.AuthOptions
+
+	// Build auth options
+	if c.opts.User != "" {
+		rootCertPemData, err := ioutil.ReadFile(c.opts.CARootCert)
+		if err != nil {
+			return errors.New("Opening CARootCert: " + err.Error())
+		}
+
+		rootCertPool := x509.NewCertPool()
+		if !rootCertPool.AppendCertsFromPEM(rootCertPemData) {
+			return errors.New("Invalid PEM certificate file")
+		}
+
+		tlsConf := &tls.Config{
+			InsecureSkipVerify: true,
+			RootCAs:            rootCertPool,
+		}
+
+		authOptions = &riak.AuthOptions{
+			User:      c.opts.User,
+			Password:  c.opts.Password,
+			TlsConfig: tlsConf,
+		}
+	}
+
+	var nodes []*riak.Node
+
+	// Set to default port if not provided
+	port := c.opts.Port
+	if port == 0 {
+		port = 8087
+	}
+
+	for _, address := range c.opts.Addresses {
+		if !strings.Contains(address, ":") {
+			// Add port if not set in the user config
+			address = address + ":" + strconv.FormatUint(uint64(port), 10)
+		}
+
+		// Set ServerName based on the address we're connecting to
+		if authOptions != nil {
+			addressWithoutPort := address[0:strings.Index(address, ":")]
+			authOptions.TlsConfig.ServerName = addressWithoutPort
+		}
+
+		node, err := riak.NewNode(&riak.NodeOptions{
+			RemoteAddress: address,
+			AuthOptions:   authOptions,
+		})
+		if err != nil {
+			return err
+		}
+
+		nodes = append(nodes, node)
+	}
+
+	con, err := riak.NewCluster(&riak.ClusterOptions{
+		Nodes: nodes,
 	})
+	if err != nil {
+		return err
+	}
 
+	err = con.Start()
 	if err != nil {
 		return err
 	}