Golang callinfo.NewContext函数代码示例

func TestLogStatsRemoteAddrUsername(t *testing.T) {
	logStats := newLogStats("test", context.Background())
	addr, user := logStats.RemoteAddrUsername()
	if addr != "" {
		t.Fatalf("remote addr should be empty")
	}
	if user != "" {
		t.Fatalf("username should be empty")
	}

	remoteAddr := "1.2.3.4"
	username := "vt"
	callInfo := &fakeCallInfo{
		remoteAddr: remoteAddr,
		username:   username,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	logStats = newLogStats("test", ctx)
	addr, user = logStats.RemoteAddrUsername()
	if addr != remoteAddr {
		t.Fatalf("expected to get remote addr: %s, but got: %s", remoteAddr, addr)
	}
	if user != username {
		t.Fatalf("expected to get username: %s, but got: %s", username, user)
	}
}

func TestQueryExecutorBlacklistQRFail(t *testing.T) {
	db := setUpQueryExecutorTest()
	query := "select * from test_table where name = 1 limit 1000"
	expandedQuery := "select pk from test_table use index (`index`) where name = 1 limit 1000"
	expected := &mproto.QueryResult{
		Fields: getTestTableFields(),
	}
	db.AddQuery(query, expected)
	db.AddQuery(expandedQuery, expected)

	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	bannedAddr := "127.0.0.1"
	bannedUser := "u2"

	alterRule := NewQueryRule("disable update", "disable update", QR_FAIL)
	alterRule.SetIPCond(bannedAddr)
	alterRule.SetUserCond(bannedUser)
	alterRule.SetQueryCond("select.*")
	alterRule.AddPlanCond(planbuilder.PLAN_SELECT_SUBQUERY)
	alterRule.AddTableCond("test_table")

	rulesName := "blacklistedRulesQRFail"
	rules := NewQueryRules()
	rules.Add(alterRule)

	callInfo := &fakeCallInfo{
		remoteAddr: bannedAddr,
		username:   bannedUser,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	tsv := newTestTabletServer(ctx, enableRowCache|enableStrict, db)
	tsv.qe.schemaInfo.queryRuleSources.UnRegisterQueryRuleSource(rulesName)
	tsv.qe.schemaInfo.queryRuleSources.RegisterQueryRuleSource(rulesName)
	defer tsv.qe.schemaInfo.queryRuleSources.UnRegisterQueryRuleSource(rulesName)

	if err := tsv.qe.schemaInfo.queryRuleSources.SetRules(rulesName, rules); err != nil {
		t.Fatalf("failed to set rule, error: %v", err)
	}

	qre := newTestQueryExecutor(ctx, tsv, query, 0)
	defer tsv.StopService()

	checkPlanID(t, planbuilder.PLAN_SELECT_SUBQUERY, qre.plan.PlanId)
	// execute should fail because query has been blacklisted
	_, err := qre.Execute()
	if err == nil {
		t.Fatal("got: nil, want: error")
	}
	got, ok := err.(*TabletError)
	if !ok {
		t.Fatalf("got: %v, want: *TabletError", err)
	}
	if got.ErrorType != ErrFail {
		t.Fatalf("got: %s, want: ErrFail", getTabletErrorString(got.ErrorType))
	}
}

func TestQueryExecutorTableAclDryRun(t *testing.T) {
	aclName := fmt.Sprintf("simpleacl-test-%d", rand.Int63())
	tableacl.Register(aclName, &simpleacl.Factory{})
	tableacl.SetDefaultACL(aclName)
	db := setUpQueryExecutorTest()
	query := "select * from test_table limit 1000"
	want := &mproto.QueryResult{
		Fields:       getTestTableFields(),
		RowsAffected: 0,
		Rows:         [][]sqltypes.Value{},
	}
	db.AddQuery(query, want)
	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	username := "u2"
	callInfo := &fakeCallInfo{
		remoteAddr: "1.2.3.4",
		username:   username,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)

	config := &tableaclpb.Config{
		TableGroups: []*tableaclpb.TableGroupSpec{{
			Name:                 "group02",
			TableNamesOrPrefixes: []string{"test_table"},
			Readers:              []string{"u1"},
		}},
	}

	if err := tableacl.InitFromProto(config); err != nil {
		t.Fatalf("unable to load tableacl config, error: %v", err)
	}

	tableACLStatsKey := strings.Join([]string{
		"test_table",
		username,
		planbuilder.PLAN_PASS_SELECT.String(),
		username,
	}, ".")
	// enable Config.StrictTableAcl
	sqlQuery := newTestSQLQuery(ctx, enableRowCache|enableSchemaOverrides|enableStrict|enableStrictTableAcl)
	sqlQuery.qe.enableTableAclDryRun = true
	qre := newTestQueryExecutor(ctx, sqlQuery, query, 0)
	defer sqlQuery.disallowQueries()
	checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
	beforeCount := sqlQuery.qe.tableaclPseudoDenied.Counters.Counts()[tableACLStatsKey]
	// query should fail because current user do not have read permissions
	_, err := qre.Execute()
	if err != nil {
		t.Fatalf("qre.Execute() = %v, want: nil", err)
	}
	afterCount := sqlQuery.qe.tableaclPseudoDenied.Counters.Counts()[tableACLStatsKey]
	if afterCount-beforeCount != 1 {
		t.Fatalf("table acl pseudo denied count should increase by one. got: %d, want: %d", afterCount, beforeCount+1)
	}
}

func TestQueryExecutorBlacklistQRRetry(t *testing.T) {
	db := setUpQueryExecutorTest()
	query := "select * from test_table where name = 1 limit 1000"
	expandedQuery := "select pk from test_table use index (`index`) where name = 1 limit 1000"
	expected := &sqltypes.Result{
		Fields: getTestTableFields(),
	}
	db.AddQuery(query, expected)
	db.AddQuery(expandedQuery, expected)

	db.AddQuery("select * from test_table where 1 != 1", &sqltypes.Result{
		Fields: getTestTableFields(),
	})

	bannedAddr := "127.0.0.1"
	bannedUser := "x"

	alterRule := NewQueryRule("disable update", "disable update", QRFailRetry)
	alterRule.SetIPCond(bannedAddr)
	alterRule.SetUserCond(bannedUser)
	alterRule.SetQueryCond("select.*")
	alterRule.AddPlanCond(planbuilder.PlanPassSelect)
	alterRule.AddTableCond("test_table")

	rulesName := "blacklistedRulesQRRetry"
	rules := NewQueryRules()
	rules.Add(alterRule)

	callInfo := &fakeCallInfo{
		remoteAddr: bannedAddr,
		username:   bannedUser,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	tsv := newTestTabletServer(ctx, enableStrict, db)
	tsv.qe.schemaInfo.queryRuleSources.UnRegisterQueryRuleSource(rulesName)
	tsv.qe.schemaInfo.queryRuleSources.RegisterQueryRuleSource(rulesName)
	defer tsv.qe.schemaInfo.queryRuleSources.UnRegisterQueryRuleSource(rulesName)

	if err := tsv.qe.schemaInfo.queryRuleSources.SetRules(rulesName, rules); err != nil {
		t.Fatalf("failed to set rule, error: %v", err)
	}

	qre := newTestQueryExecutor(ctx, tsv, query, 0)
	defer tsv.StopService()

	checkPlanID(t, planbuilder.PlanPassSelect, qre.plan.PlanID)
	_, err := qre.Execute()
	if err == nil {
		t.Fatal("got: nil, want: error")
	}
	got, ok := err.(*TabletError)
	if !ok {
		t.Fatalf("got: %v, want: *TabletError", err)
	}
	if got.ErrorCode != vtrpcpb.ErrorCode_QUERY_NOT_SERVED {
		t.Fatalf("got: %s, want: QUERY_NOT_SERVED", got.ErrorCode)
	}
}

func TestQueryExecutorBlacklistQRRetry(t *testing.T) {
	db := setUpQueryExecutorTest()
	query := "select * from test_table where name = 1 limit 1000"
	expandedQuery := "select pk from test_table use index (INDEX) where name = 1 limit 1000"
	expected := &mproto.QueryResult{
		Fields: getTestTableFields(),
	}
	db.AddQuery(query, expected)
	db.AddQuery(expandedQuery, expected)

	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	bannedAddr := "127.0.0.1"
	bannedUser := "x"

	alterRule := NewQueryRule("disable update", "disable update", QR_FAIL_RETRY)
	alterRule.SetIPCond(bannedAddr)
	alterRule.SetUserCond(bannedUser)
	alterRule.SetQueryCond("select.*")
	alterRule.AddPlanCond(planbuilder.PLAN_SELECT_SUBQUERY)
	alterRule.AddTableCond("test_table")

	rulesName := "blacklistedRulesQRRetry"
	rules := NewQueryRules()
	rules.Add(alterRule)

	QueryRuleSources.UnRegisterQueryRuleSource(rulesName)
	QueryRuleSources.RegisterQueryRuleSource(rulesName)
	defer QueryRuleSources.UnRegisterQueryRuleSource(rulesName)

	if err := QueryRuleSources.SetRules(rulesName, rules); err != nil {
		t.Fatalf("failed to set rule, error: %v", err)
	}

	callInfo := &fakeCallInfo{
		remoteAddr: bannedAddr,
		username:   bannedUser,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	sqlQuery := newTestSQLQuery(ctx, enableRowCache|enableStrict)
	qre := newTestQueryExecutor(ctx, sqlQuery, query, 0)
	defer sqlQuery.disallowQueries()
	checkPlanID(t, planbuilder.PLAN_SELECT_SUBQUERY, qre.plan.PlanId)
	_, err := qre.Execute()
	if err == nil {
		t.Fatal("got: nil, want: error")
	}
	got, ok := err.(*TabletError)
	if !ok {
		t.Fatalf("got: %v, want: *TabletError", err)
	}
	if got.ErrorType != ErrRetry {
		t.Fatalf("got: %s, want: ErrRetry", getTabletErrorString(got.ErrorType))
	}
}

func TestLogStatsContextHTML(t *testing.T) {
	html := "HtmlContext"
	callInfo := &fakeCallInfo{
		html: html,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	logStats := newLogStats("test", ctx)
	if string(logStats.ContextHTML()) != html {
		t.Fatalf("expect to get html: %s, but got: %s", html, string(logStats.ContextHTML()))
	}
}

func TestQueryExecutorTableAcl(t *testing.T) {
	testUtils := &testUtils{}
	aclName := fmt.Sprintf("simpleacl-test-%d", rand.Int63())
	tableacl.Register(aclName, &simpleacl.Factory{})
	tableacl.SetDefaultACL(aclName)

	db := setUpQueryExecutorTest()
	query := "select * from test_table limit 1000"
	expected := &mproto.QueryResult{
		Fields:       getTestTableFields(),
		RowsAffected: 0,
		Rows:         [][]sqltypes.Value{},
	}
	db.AddQuery(query, expected)
	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	username := "u2"
	callInfo := &fakeCallInfo{
		remoteAddr: "1.2.3.4",
		username:   username,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	if err := tableacl.InitFromBytes(
		[]byte(fmt.Sprintf(`{"test_table":{"READER":"%s"}}`, username))); err != nil {
		t.Fatalf("unable to load tableacl config, error: %v", err)
	}

	qre, sqlQuery := newTestQueryExecutor(
		query, ctx, enableRowCache|enableSchemaOverrides|enableStrict)
	checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
	testUtils.checkEqual(t, expected, qre.Execute())
	sqlQuery.disallowQueries()

	if err := tableacl.InitFromBytes([]byte(`{"test_table":{"READER":"superuser"}}`)); err != nil {
		t.Fatalf("unable to load tableacl config, error: %v", err)
	}
	// without enabling Config.StrictTableAcl
	qre, sqlQuery = newTestQueryExecutor(
		query, ctx, enableRowCache|enableSchemaOverrides|enableStrict)
	checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
	qre.Execute()
	sqlQuery.disallowQueries()
	// enable Config.StrictTableAcl
	qre, sqlQuery = newTestQueryExecutor(
		query, ctx, enableRowCache|enableSchemaOverrides|enableStrict|enableStrictTableAcl)
	defer sqlQuery.disallowQueries()
	checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
	defer handleAndVerifyTabletError(t, "query should fail because current user do not have read permissions", ErrFail)
	qre.Execute()
}

func TestQueryExecutorBlacklistQRRetry(t *testing.T) {
	db := setUpQueryExecutorTest()
	query := "select * from test_table where name = 1 limit 1000"
	expandedQuery := "select pk from test_table use index (INDEX) where name = 1 limit 1000"
	expected := &mproto.QueryResult{
		Fields: getTestTableFields(),
	}
	db.AddQuery(query, expected)
	db.AddQuery(expandedQuery, expected)

	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	bannedAddr := "127.0.0.1"
	bannedUser := "x"

	alterRule := NewQueryRule("disable update", "disable update", QR_FAIL_RETRY)
	alterRule.SetIPCond(bannedAddr)
	alterRule.SetUserCond(bannedUser)
	alterRule.SetQueryCond("select.*")
	alterRule.AddPlanCond(planbuilder.PLAN_SELECT_SUBQUERY)
	alterRule.AddTableCond("test_table")

	rulesName := "blacklistedRulesQRRetry"
	rules := NewQueryRules()
	rules.Add(alterRule)

	QueryRuleSources.UnRegisterQueryRuleSource(rulesName)
	QueryRuleSources.RegisterQueryRuleSource(rulesName)
	defer QueryRuleSources.UnRegisterQueryRuleSource(rulesName)

	if err := QueryRuleSources.SetRules(rulesName, rules); err != nil {
		t.Fatalf("failed to set rule, error: %v", err)
	}

	callInfo := &fakeCallInfo{
		remoteAddr: bannedAddr,
		username:   bannedUser,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	qre, sqlQuery := newTestQueryExecutor(query, ctx, enableRowCache|enableStrict)
	defer sqlQuery.disallowQueries()
	checkPlanID(t, planbuilder.PLAN_SELECT_SUBQUERY, qre.plan.PlanId)
	defer handleAndVerifyTabletError(t, "execute should fail because query has been blacklisted", ErrRetry)
	qre.Execute()
}

func TestQueryExecutorTableAcl(t *testing.T) {
	aclName := fmt.Sprintf("simpleacl-test-%d", rand.Int63())
	tableacl.Register(aclName, &simpleacl.Factory{})
	tableacl.SetDefaultACL(aclName)
	db := setUpQueryExecutorTest()
	query := "select * from test_table limit 1000"
	want := &mproto.QueryResult{
		Fields:       getTestTableFields(),
		RowsAffected: 0,
		Rows:         [][]sqltypes.Value{},
	}
	db.AddQuery(query, want)
	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	username := "u2"
	callInfo := &fakeCallInfo{
		remoteAddr: "1.2.3.4",
		username:   username,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)
	config := &tableaclpb.Config{
		TableGroups: []*tableaclpb.TableGroupSpec{{
			Name:                 "group01",
			TableNamesOrPrefixes: []string{"test_table"},
			Readers:              []string{username},
		}},
	}
	if err := tableacl.InitFromProto(config); err != nil {
		t.Fatalf("unable to load tableacl config, error: %v", err)
	}

	sqlQuery := newTestSQLQuery(ctx, enableRowCache|enableSchemaOverrides|enableStrict)
	qre := newTestQueryExecutor(ctx, sqlQuery, query, 0)
	defer sqlQuery.disallowQueries()
	checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
	got, err := qre.Execute()
	if err != nil {
		t.Fatalf("got: %v, want nil", err)
	}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("qre.Execute() = %v, want: %v", got, want)
	}
}

func TestQueryExecutorTableAclExemptACL(t *testing.T) {
	aclName := fmt.Sprintf("simpleacl-test-%d", rand.Int63())
	tableacl.Register(aclName, &simpleacl.Factory{})
	tableacl.SetDefaultACL(aclName)
	db := setUpQueryExecutorTest()
	query := "select * from test_table limit 1000"
	want := &mproto.QueryResult{
		Fields:       getTestTableFields(),
		RowsAffected: 0,
		Rows:         [][]sqltypes.Value{},
	}
	db.AddQuery(query, want)
	db.AddQuery("select * from test_table where 1 != 1", &mproto.QueryResult{
		Fields: getTestTableFields(),
	})

	username := "u2"
	callInfo := &fakeCallInfo{
		remoteAddr: "1.2.3.4",
		username:   username,
	}
	ctx := callinfo.NewContext(context.Background(), callInfo)

	config := &tableaclpb.Config{
		TableGroups: []*tableaclpb.TableGroupSpec{{
			Name:                 "group02",
			TableNamesOrPrefixes: []string{"test_table"},
			Readers:              []string{"u1"},
		}},
	}

	if err := tableacl.InitFromProto(config); err != nil {
		t.Fatalf("unable to load tableacl config, error: %v", err)
	}

	// enable Config.StrictTableAcl
	sqlQuery := newTestSQLQuery(ctx, enableRowCache|enableSchemaOverrides|enableStrict|enableStrictTableAcl)
	qre := newTestQueryExecutor(ctx, sqlQuery, query, 0)
	defer sqlQuery.disallowQueries()
	checkPlanID(t, planbuilder.PLAN_PASS_SELECT, qre.plan.PlanId)
	// query should fail because current user do not have read permissions
	_, err := qre.Execute()
	if err == nil {
		t.Fatal("got: nil, want: error")
	}
	tabletError, ok := err.(*TabletError)
	if !ok {
		t.Fatalf("got: %v, want: *TabletError", err)
	}
	if tabletError.ErrorType != ErrFail {
		t.Fatalf("got: %s, want: ErrFail", getTabletErrorString(tabletError.ErrorType))
	}
	if !strings.Contains(tabletError.Error(), "table acl error") {
		t.Fatalf("got %s, want tablet errorL table acl error", tabletError.Error())
	}

	// table acl should be ignored since this is an exempt user.
	username = "exempt-acl"
	sqlQuery.qe.exemptACL = username
	callInfo = &fakeCallInfo{
		remoteAddr: "1.2.3.4",
		username:   username,
	}
	ctx = callinfo.NewContext(context.Background(), callInfo)
	qre = newTestQueryExecutor(ctx, sqlQuery, query, 0)
	_, err = qre.Execute()
	if err != nil {
		t.Fatal("qre.Execute: nil, want: error")
	}
}