enh: support whitelist

common/const.go

@@ -142,4 +142,7 @@ const (
 
 const ReqIDKey = "taos_req_id"
 
-const TAOS_NOTIFY_PASSVER = 0
+const (
+	TAOS_NOTIFY_PASSVER       = 0
+	TAOS_NOTIFY_WHITELIST_VER = 1
+)

wrapper/notify_test.go

@@ -49,4 +49,25 @@ func TestNotify(t *testing.T) {
 	case <-timeout.Done():
 		t.Error("wait for notify callback timeout")
 	}
+	{
+		notify := make(chan int64, 1)
+		handler := cgo.NewHandle(notify)
+		errCode := TaosSetNotifyCB(conn2, handler, common.TAOS_NOTIFY_WHITELIST_VER)
+		if errCode != 0 {
+			errStr := TaosErrorStr(nil)
+			t.Error(errCode, errStr)
+		}
+		err = exec(conn, "ALTER USER t_notify ADD HOST '192.168.1.98/0','192.168.1.98/32'")
+		assert.NoError(t, err)
+		timeout, cancel = context.WithTimeout(context.Background(), time.Second*5)
+		defer cancel()
+		now := time.Now()
+		select {
+		case version := <-notify:
+			fmt.Println(time.Now().Sub(now))
+			t.Log(version)
+		case <-timeout.Done():
+			t.Error("wait for notify callback timeout")
+		}
+	}
 }

wrapper/notifycb.go

@@ -17,11 +17,16 @@ import (
 //export NotifyCallback
 func NotifyCallback(p unsafe.Pointer, ext unsafe.Pointer, notifyType C.int) {
 	defer func() {
+		// channel may be closed
 		recover()
 	}()
 	if int(notifyType) == common.TAOS_NOTIFY_PASSVER {
 		version := int32(*(*C.int32_t)(ext))
 		c := (*(*cgo.Handle)(p)).Value().(chan int32)
 		c <- version
+	} else if int(notifyType) == common.TAOS_NOTIFY_WHITELIST_VER {
+		version := int64(*(*C.int64_t)(ext))
+		c := (*(*cgo.Handle)(p)).Value().(chan int64)
+		c <- version
 	}
 }

wrapper/tmq.go

@@ -327,3 +327,8 @@ func TMQCommitOffsetAsync(consumer unsafe.Pointer, topic string, vGroupID int32,
 	defer C.free(unsafe.Pointer(topicName))
 	C.tmq_commit_offset_async((*C.tmq_t)(consumer), topicName, (C.int32_t)(vGroupID), (C.int64_t)(offset), (*C.tmq_commit_cb)(C.TMQCommitOffsetCB), h.Pointer())
 }
+
+// TMQGetConnect TAOS *tmq_get_connect(tmq_t *tmq)
+func TMQGetConnect(consumer unsafe.Pointer) unsafe.Pointer {
+	return unsafe.Pointer(C.tmq_get_connect((*C.tmq_t)(consumer)))
+}

wrapper/whitelist.go

@@ -0,0 +1,31 @@
+package wrapper
+
+/*
+#cgo CFLAGS: -IC:/TDengine/include -I/usr/include
+#cgo linux LDFLAGS: -L/usr/lib -ltaos
+#cgo windows LDFLAGS: -LC:/TDengine/driver -ltaos
+#cgo darwin LDFLAGS: -L/usr/local/lib -ltaos
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <taos.h>
+extern void WhitelistCallback(void *param, int code, TAOS *taos, int numOfWhiteLists, uint64_t* pWhiteLists);
+void taos_fetch_whitelist_a_wrapper(TAOS *taos, void *param){
+	return taos_fetch_whitelist_a(taos, WhitelistCallback, param);
+};
+*/
+import "C"
+import (
+	"unsafe"
+
+	"github.com/taosdata/driver-go/v3/wrapper/cgo"
+)
+
+const TAOS_NOTIFY_WHITELIST_VER = 1
+
+// typedef void (*__taos_async_whitelist_fn_t)(void *param, int code, TAOS *taos, int numOfWhiteLists, uint64_t* pWhiteLists);
+
+// TaosFetchWhitelistA DLL_EXPORT void taos_fetch_whitelist_a(TAOS *taos, __taos_async_whitelist_fn_t fp, void *param);
+func TaosFetchWhitelistA(taosConnect unsafe.Pointer, caller cgo.Handle) {
+	C.taos_fetch_whitelist_a_wrapper(taosConnect, caller.Pointer())
+}

wrapper/whitelist_test.go

@@ -0,0 +1,21 @@
+package wrapper
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/taosdata/driver-go/v3/wrapper/cgo"
+)
+
+func TestGetWhiteList(t *testing.T) {
+	conn, err := TaosConnect("", "root", "taosdata", "", 0)
+	assert.NoError(t, err)
+	defer TaosClose(conn)
+	c := make(chan *WhitelistResult, 1)
+	handler := cgo.NewHandle(c)
+	TaosFetchWhitelistA(conn, handler)
+	data := <-c
+	assert.Equal(t, int32(0), data.ErrCode)
+	assert.Equal(t, 1, len(data.IPNets))
+	assert.Equal(t, "0.0.0.0/0", data.IPNets[0].String())
+}

wrapper/whitelistcb.go

@@ -0,0 +1,35 @@
+package wrapper
+
+import "C"
+import (
+	"net"
+	"unsafe"
+
+	"github.com/taosdata/driver-go/v3/wrapper/cgo"
+)
+
+type WhitelistResult struct {
+	ErrCode int32
+	IPNets  []*net.IPNet
+}
+
+//export WhitelistCallback
+func WhitelistCallback(param unsafe.Pointer, code int, taosConnect unsafe.Pointer, numOfWhiteLists int, pWhiteLists unsafe.Pointer) {
+	c := (*(*cgo.Handle)(param)).Value().(chan *WhitelistResult)
+	if code != 0 {
+		c <- &WhitelistResult{ErrCode: int32(code)}
+		return
+	}
+	ips := make([]*net.IPNet, 0, numOfWhiteLists)
+	for i := 0; i < numOfWhiteLists; i++ {
+		ipNet := make([]byte, 8)
+		for j := 0; j < 8; j++ {
+			ipNet[j] = *(*byte)(unsafe.Pointer(uintptr(pWhiteLists) + uintptr(i*8) + uintptr(j)))
+		}
+		ip := net.IP{ipNet[0], ipNet[1], ipNet[2], ipNet[3]}
+		ones := int(ipNet[4])
+		ipMask := net.CIDRMask(ones, 32)
+		ips = append(ips, &net.IPNet{IP: ip, Mask: ipMask})
+	}
+	c <- &WhitelistResult{IPNets: ips}
+}