MOD: support chroot for a connection

Author: git-hulk

zk/conn.go

@@ -16,6 +16,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"reflect"
 	"strconv"
 	"strings"
 	"sync"
@@ -73,6 +74,7 @@ type Conn struct {
 	xid              uint32
 	sessionTimeoutMs int32 // session timeout in milliseconds
 	passwd           []byte
+	chroot           string
 
 	dialer         Dialer
 	hostProvider   HostProvider
@@ -196,6 +198,7 @@ func Connect(servers []string, sessionTimeout time.Duration, options ...connOpti
 		requests:       make(map[int32]*request),
 		watchers:       make(map[watchPathType][]chan Event),
 		passwd:         emptyPassword,
+		chroot:         "",
 		logger:         DefaultLogger,
 		buf:            make([]byte, bufferSize),
 
@@ -611,6 +614,19 @@ func (c *Conn) sendData(req *request) error {
 		return nil
 	}
 
+	if req != nil && req.pkt != nil && c.chroot != "" {
+		v := reflect.ValueOf(req.pkt)
+		for v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface {
+			v = v.Elem()
+		}
+		if v.Kind() == reflect.Struct {
+			field := v.FieldByName("Path")
+			if field.Kind() == reflect.String {
+				field.SetString(c.chroot + field.String())
+			}
+		}
+	}
+
 	n2, err := encodePacket(c.buf[4+n:], req.pkt)
 	if err != nil {
 		req.recvChan <- response{-1, err}
@@ -708,6 +724,9 @@ func (c *Conn) recvLoop(conn net.Conn) error {
 			if err != nil {
 				return err
 			}
+			if c.chroot != "" {
+				res.Path = strings.TrimPrefix(res.Path, c.chroot)
+			}
 			ev := Event{
 				Type:  res.Type,
 				State: res.State,
@@ -760,6 +779,19 @@ func (c *Conn) recvLoop(conn net.Conn) error {
 				} else {
 					_, err = decodePacket(buf[16:blen], req.recvStruct)
 				}
+
+				if req != nil && req.recvStruct != nil && c.chroot != "" {
+					v := reflect.ValueOf(req.recvStruct)
+					for v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface {
+						v = v.Elem()
+					}
+					if v.Kind() == reflect.Struct {
+						field := v.FieldByName("Path")
+						if field.Kind() == reflect.String {
+							field.SetString(strings.TrimPrefix(field.String(), c.chroot))
+						}
+					}
+				}
 				if req.recvFunc != nil {
 					req.recvFunc(req, &res, err)
 				}
@@ -804,6 +836,16 @@ func (c *Conn) request(opcode int32, req interface{}, res interface{}, recvFunc
 	return r.zxid, r.err
 }
 
+func (c *Conn) Chroot(path string) error {
+	res := &existsResponse{}
+	_, err := c.request(opExists, &existsRequest{Path: path, Watch: false}, res, nil)
+	if err != nil {
+		return err
+	}
+	c.chroot = path
+	return nil
+}
+
 func (c *Conn) AddAuth(scheme string, auth []byte) error {
 	_, err := c.request(opSetAuth, &setAuthRequest{Type: 0, Scheme: scheme, Auth: auth}, &setAuthResponse{}, nil)
 

zk/zk_test.go

@@ -575,6 +575,49 @@ func TestExpiringWatch(t *testing.T) {
 	}
 }
 
+func TestChroot(t *testing.T) {
+	ts, err := StartTestCluster(1, nil, logWriter{t: t, p: "[ZKERR] "})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ts.Stop()
+	zk, _, err := ts.ConnectAll()
+	if err != nil {
+		t.Fatalf("Connect returned error: %+v", err)
+	}
+	defer zk.Close()
+
+	path := "/gozk-test-chroot"
+	err = zk.Chroot(path)
+	if err == nil {
+		t.Fatal("Chroot expect error when path is not exist")
+	}
+	_, err = zk.Create(path, []byte{1, 2, 3, 4}, 0, WorldACL(PermAll))
+	if err != nil {
+		t.Fatalf("Chroot create path error, %s", err.Error())
+	}
+	err = zk.Chroot(path)
+	if err != nil {
+		t.Fatalf("Chroot error, %s", err.Error())
+	}
+	subPath := "/abc"
+	actualPath, err := zk.Create(subPath, []byte{1, 2, 3, 4}, 0, WorldACL(PermAll))
+	if err != nil {
+		t.Fatalf("Chroot create path error, %s", err.Error())
+	}
+	if actualPath != subPath {
+		t.Fatalf("path expect %s, got %s", path+subPath, actualPath)
+	}
+	err = zk.Chroot("")
+	if err != nil {
+		t.Fatalf("Chroot error, %s", err.Error())
+	}
+	exists, _, err := zk.Exists(path + subPath)
+	if !exists || err != nil {
+		t.Fatal("Chroot error as node exists")
+	}
+}
+
 func TestRequestFail(t *testing.T) {
 	// If connecting fails to all servers in the list then pending requests
 	// should be errored out so they don't hang forever.