Преглед на файлове

Fix a panic that could happen when using the TokenAwareHostPolicy

This was due to that the fallbackHost could be nil but was still checked
accessed to check if the host was the same as the filtered.
Max Ekman преди 9 години
родител
ревизия
6eee003f22
променени са 2 файла, в които са добавени 39 реда и са изтрити 1 реда
  1. 1 1
      policies.go
  2. 38 0
      policies_test.go

+ 1 - 1
policies.go

@@ -370,7 +370,7 @@ func (t *tokenAwareHostPolicy) Pick(qry *Query) NextHost {
 		fallbackHost := fallbackIter()
 
 		// filter the token aware selected hosts from the fallback hosts
-		if fallbackHost.Info() == host {
+		if fallbackHost != nil && fallbackHost.Info() == host {
 			fallbackHost = fallbackIter()
 		}
 

+ 38 - 0
policies_test.go

@@ -192,6 +192,44 @@ func TestRoundRobinNilHostInfo(t *testing.T) {
 	}
 }
 
+func TestTokenAwareNilHostInfo(t *testing.T) {
+	policy := TokenAwareHostPolicy(RoundRobinHostPolicy())
+
+	hosts := []*HostInfo{
+		{peer: "0", tokens: []string{"00"}},
+		{peer: "1", tokens: []string{"25"}},
+		{peer: "2", tokens: []string{"50"}},
+		{peer: "3", tokens: []string{"75"}},
+	}
+	policy.SetHosts(hosts)
+	policy.SetPartitioner("OrderedPartitioner")
+
+	query := &Query{}
+	query.RoutingKey([]byte("20"))
+
+	iter := policy.Pick(query)
+	next := iter()
+	if next == nil {
+		t.Fatal("got nil host")
+	} else if v := next.Info(); v == nil {
+		t.Fatal("got nil HostInfo")
+	} else if v.Peer() != "1" {
+		t.Fatalf("expected peer 1 got %v", v.Peer())
+	}
+
+	// Empty the hosts to trigger the panic when using the fallback.
+	hosts = []*HostInfo{}
+	policy.SetHosts(hosts)
+
+	next = iter()
+	if next != nil {
+		t.Errorf("expected to get nil host got %+v", next)
+		if next.Info() == nil {
+			t.Fatalf("HostInfo is nil")
+		}
+	}
+}
+
 func TestCOWList_Add(t *testing.T) {
 	var cow cowHostList