From 0e81c620316db95b752da8023e416ef52ae95858 Mon Sep 17 00:00:00 2001
From: bigsheeper <yihao.dai@zilliz.com>
Date: Fri, 10 Sep 2021 18:12:01 +0800
Subject: [PATCH] Check collection and plan before calling cgo functions
 (#7733)

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
---
 internal/querynode/plan.go        |  9 +++++++++
 internal/querynode/plan_test.go   | 12 ++++++++++++
 internal/querynode/reduce.go      |  4 ++++
 internal/querynode/reduce_test.go |  6 ++++++
 4 files changed, 31 insertions(+)

diff --git a/internal/querynode/plan.go b/internal/querynode/plan.go
index dcb8ef9fd6..9be6dba362 100644
--- a/internal/querynode/plan.go
+++ b/internal/querynode/plan.go
@@ -22,6 +22,7 @@ package querynode
 import "C"
 import (
 	"errors"
+	"fmt"
 	"unsafe"
 )
 
@@ -30,6 +31,10 @@ type SearchPlan struct {
 }
 
 func createSearchPlan(col *Collection, dsl string) (*SearchPlan, error) {
+	if col.collectionPtr == nil {
+		return nil, errors.New("nil collection ptr, collectionID = " + fmt.Sprintln(col.id))
+	}
+
 	cDsl := C.CString(dsl)
 	defer C.free(unsafe.Pointer(cDsl))
 	var cPlan C.CSearchPlan
@@ -45,6 +50,10 @@ func createSearchPlan(col *Collection, dsl string) (*SearchPlan, error) {
 }
 
 func createSearchPlanByExpr(col *Collection, expr []byte) (*SearchPlan, error) {
+	if col.collectionPtr == nil {
+		return nil, errors.New("nil collection ptr, collectionID = " + fmt.Sprintln(col.id))
+	}
+
 	var cPlan C.CSearchPlan
 	status := C.CreateSearchPlanByExpr(col.collectionPtr, (*C.char)(unsafe.Pointer(&expr[0])), (C.int64_t)(len(expr)), &cPlan)
 
diff --git a/internal/querynode/plan_test.go b/internal/querynode/plan_test.go
index d4cc1ea113..953c1c78a7 100644
--- a/internal/querynode/plan_test.go
+++ b/internal/querynode/plan_test.go
@@ -41,6 +41,18 @@ func TestPlan_Plan(t *testing.T) {
 	deleteCollection(collection)
 }
 
+func TestPlan_NilCollection(t *testing.T) {
+	collection := &Collection{
+		id: defaultCollectionID,
+	}
+
+	_, err := createSearchPlan(collection, "")
+	assert.Error(t, err)
+
+	_, err = createSearchPlanByExpr(collection, nil)
+	assert.Error(t, err)
+}
+
 func TestPlan_PlaceholderGroup(t *testing.T) {
 	collectionID := UniqueID(0)
 	collectionMeta := genTestCollectionMeta(collectionID, false)
diff --git a/internal/querynode/reduce.go b/internal/querynode/reduce.go
index 5dbba1dc26..0f248c2178 100644
--- a/internal/querynode/reduce.go
+++ b/internal/querynode/reduce.go
@@ -35,6 +35,10 @@ type MarshaledHits struct {
 }
 
 func reduceSearchResultsAndFillData(plan *SearchPlan, searchResults []*SearchResult, numSegments int64) error {
+	if plan.cSearchPlan == nil {
+		return errors.New("nil search plan")
+	}
+
 	cSearchResults := make([]C.CSearchResult, 0)
 	for _, res := range searchResults {
 		cSearchResults = append(cSearchResults, res.cSearchResult)
diff --git a/internal/querynode/reduce_test.go b/internal/querynode/reduce_test.go
index 1175e21a88..9777ba9396 100644
--- a/internal/querynode/reduce_test.go
+++ b/internal/querynode/reduce_test.go
@@ -106,3 +106,9 @@ func TestReduce_AllFunc(t *testing.T) {
 	deleteSegment(segment)
 	deleteCollection(collection)
 }
+
+func TestReduce_nilPlan(t *testing.T) {
+	plan := &SearchPlan{}
+	err := reduceSearchResultsAndFillData(plan, nil, 1)
+	assert.Error(t, err)
+}