// loadKeyspaceAndBlacklistRules does what the name suggests: // 1. load and build keyrange query rules // 2. load and build blacklist query rules func (agent *ActionAgent) loadKeyspaceAndBlacklistRules(tablet *pbt.Tablet, blacklistedTables []string) (err error) { // Keyrange rules keyrangeRules := tabletserver.NewQueryRules() if key.KeyRangeIsPartial(tablet.KeyRange) { log.Infof("Restricting to keyrange: %v", tablet.KeyRange) dmlPlans := []struct { planID planbuilder.PlanType onAbsent bool }{ {planbuilder.PlanInsertPK, true}, {planbuilder.PlanInsertSubquery, true}, {planbuilder.PlanPassDML, false}, {planbuilder.PlanDMLPK, false}, {planbuilder.PlanDMLSubquery, false}, {planbuilder.PlanUpsertPK, false}, } for _, plan := range dmlPlans { qr := tabletserver.NewQueryRule( fmt.Sprintf("enforce keyspace_id range for %v", plan.planID), fmt.Sprintf("keyspace_id_not_in_range_%v", plan.planID), tabletserver.QRFail, ) qr.AddPlanCond(plan.planID) err := qr.AddBindVarCond("keyspace_id", plan.onAbsent, true, tabletserver.QRNotIn, tablet.KeyRange) if err != nil { return fmt.Errorf("Unable to add keyspace rule: %v", err) } keyrangeRules.Add(qr) } } // Blacklisted tables blacklistRules := tabletserver.NewQueryRules() if len(blacklistedTables) > 0 { // tables, first resolve wildcards tables, err := mysqlctl.ResolveTables(agent.MysqlDaemon, topoproto.TabletDbName(tablet), blacklistedTables) if err != nil { return err } log.Infof("Blacklisting tables %v", strings.Join(tables, ", ")) qr := tabletserver.NewQueryRule("enforce blacklisted tables", "blacklisted_table", tabletserver.QRFailRetry) for _, t := range tables { qr.AddTableCond(t) } blacklistRules.Add(qr) } // Push all three sets of QueryRules to TabletServerRpcService loadRuleErr := agent.QueryServiceControl.SetQueryRules(keyrangeQueryRules, keyrangeRules) if loadRuleErr != nil { log.Warningf("Fail to load query rule set %s: %s", keyrangeQueryRules, loadRuleErr) } loadRuleErr = agent.QueryServiceControl.SetQueryRules(blacklistQueryRules, blacklistRules) if loadRuleErr != nil { log.Warningf("Fail to load query rule set %s: %s", blacklistQueryRules, loadRuleErr) } return nil }
// NewZkCustomRule Creates new ZkCustomRule structure func NewZkCustomRule(zkconn zk.Conn) *ZkCustomRule { return &ZkCustomRule{ zconn: zkconn, currentRuleSet: tabletserver.NewQueryRules(), currentRuleSetVersion: InvalidQueryRulesVersion, finish: make(chan int, 1)} }
// loadBlacklistRules loads and builds the blacklist query rules func (agent *ActionAgent) loadBlacklistRules(tablet *topodatapb.Tablet, blacklistedTables []string) (err error) { blacklistRules := tabletserver.NewQueryRules() if len(blacklistedTables) > 0 { // tables, first resolve wildcards tables, err := mysqlctl.ResolveTables(agent.MysqlDaemon, topoproto.TabletDbName(tablet), blacklistedTables) if err != nil { return err } // Verify that at least one table matches the wildcards, so // that we don't add a rule to blacklist all tables if len(tables) > 0 { log.Infof("Blacklisting tables %v", strings.Join(tables, ", ")) qr := tabletserver.NewQueryRule("enforce blacklisted tables", "blacklisted_table", tabletserver.QRFailRetry) for _, t := range tables { qr.AddTableCond(t) } blacklistRules.Add(qr) } } loadRuleErr := agent.QueryServiceControl.SetQueryRules(blacklistQueryRules, blacklistRules) if loadRuleErr != nil { log.Warningf("Fail to load query rule set %s: %s", blacklistQueryRules, loadRuleErr) } return nil }
// Open try to build query rules from local file and push the rules to vttablet func (fcr *FileCustomRule) Open(qsc tabletserver.QueryServiceControl, rulePath string) error { fcr.path = rulePath if fcr.path == "" { // Don't go further if path is empty return nil } data, err := ioutil.ReadFile(fcr.path) if err != nil { log.Warningf("Error reading file %v: %v", fcr.path, err) // Don't update any internal cache, just return error return err } qrs := tabletserver.NewQueryRules() err = qrs.UnmarshalJSON(data) if err != nil { log.Warningf("Error unmarshaling query rules %v", err) return err } fcr.currentRuleSetTimestamp = time.Now().Unix() fcr.currentRuleSet = qrs.Copy() // Push query rules to vttablet qsc.SetQueryRules(FileCustomRuleSource, qrs.Copy()) log.Infof("Custom rule loaded from file: %s", fcr.path) return nil }
func loadCustomRules(customrules string) *ts.QueryRules { if customrules == "" { return ts.NewQueryRules() } data, err := ioutil.ReadFile(customrules) if err != nil { relog.Fatal("Error reading file %v: %v", customrules, err) } qrs := ts.NewQueryRules() err = qrs.UnmarshalJSON(data) if err != nil { relog.Fatal("Error unmarshaling query rules %v", err) } return qrs }
// initializeKeyRangeRule will create and set the key range rules func (agent *ActionAgent) initializeKeyRangeRule(ctx context.Context, keyspace string, keyRange *topodatapb.KeyRange) error { // check we have a partial key range if !key.KeyRangeIsPartial(keyRange) { log.Infof("Tablet covers the full KeyRange, not adding KeyRange query rule") return nil } // read the keyspace to get the sharding column name keyspaceInfo, err := agent.TopoServer.GetKeyspace(ctx, keyspace) if err != nil { return fmt.Errorf("cannot read keyspace %v to get sharding key: %v", keyspace, err) } if keyspaceInfo.ShardingColumnName == "" { log.Infof("Keyspace %v has an empty ShardingColumnName, not adding KeyRange query rule", keyspace) return nil } // create the rules log.Infof("Restricting to keyrange: %v", key.KeyRangeString(keyRange)) keyrangeRules := tabletserver.NewQueryRules() dmlPlans := []struct { planID planbuilder.PlanType onAbsent bool }{ {planbuilder.PlanInsertPK, true}, {planbuilder.PlanInsertSubquery, true}, {planbuilder.PlanPassDML, false}, {planbuilder.PlanDMLPK, false}, {planbuilder.PlanDMLSubquery, false}, {planbuilder.PlanUpsertPK, false}, } for _, plan := range dmlPlans { qr := tabletserver.NewQueryRule( fmt.Sprintf("enforce %v range for %v", keyspaceInfo.ShardingColumnName, plan.planID), fmt.Sprintf("%v_not_in_range_%v", keyspaceInfo.ShardingColumnName, plan.planID), tabletserver.QRFail, ) qr.AddPlanCond(plan.planID) err := qr.AddBindVarCond(keyspaceInfo.ShardingColumnName, plan.onAbsent, true, tabletserver.QRNotIn, keyRange) if err != nil { return fmt.Errorf("Unable to add key range rule: %v", err) } keyrangeRules.Add(qr) } // and load them agent.QueryServiceControl.RegisterQueryRuleSource(keyrangeQueryRules) if err := agent.QueryServiceControl.SetQueryRules(keyrangeQueryRules, keyrangeRules); err != nil { return fmt.Errorf("failed to load query rule set %s: %s", keyrangeQueryRules, err) } return nil }
func TestQueryRules(t *testing.T) { rules := tabletserver.NewQueryRules() err := rules.UnmarshalJSON(rulesJSON) if err != nil { t.Error(err) return } err = framework.Server.SetQueryRules("endtoend", rules) want := "Rule source identifier endtoend is not valid" if err == nil || err.Error() != want { t.Errorf("Error: %v, want %s", err, want) } framework.Server.RegisterQueryRuleSource("endtoend") defer framework.Server.UnRegisterQueryRuleSource("endtoend") err = framework.Server.SetQueryRules("endtoend", rules) if err != nil { t.Error(err) return } client := framework.NewClient() query := "select * from vitess_test where intval=:asdfg" bv := map[string]interface{}{"asdfg": 1} _, err = client.Execute(query, bv) want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'" if err == nil || err.Error() != want { t.Errorf("Error: %v, want %s", err, want) } _, err = client.StreamExecute(query, bv) want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'" if err == nil || err.Error() != want { t.Errorf("Error: %v, want %s", err, want) } err = framework.Server.SetQueryRules("endtoend", nil) if err != nil { t.Error(err) return } _, err = client.Execute(query, bv) if err != nil { t.Error(err) return } }
// refreshData gets query rules from Zookeeper and refresh internal QueryRules cache // this function will also call TabletServer.SetQueryRules to propagate rule changes to query service func (zkcr *ZkCustomRule) refreshData(qsc tabletserver.QueryServiceControl, nodeRemoval bool) error { data, stat, err := zkcr.zconn.Get(zkcr.path) zkcr.mu.Lock() defer zkcr.mu.Unlock() if err == nil { qrs := tabletserver.NewQueryRules() if !nodeRemoval { err = qrs.UnmarshalJSON([]byte(data)) if err != nil { log.Warningf("Error unmarshaling query rules %v, original data '%s'", err, data) return nil } } zkcr.currentRuleSetVersion = stat.Mzxid() if !reflect.DeepEqual(zkcr.currentRuleSet, qrs) { zkcr.currentRuleSet = qrs.Copy() qsc.SetQueryRules(ZkCustomRuleSource, qrs.Copy()) log.Infof("Custom rule version %v fetched from Zookeeper and applied to vttablet", zkcr.currentRuleSetVersion) } return nil } log.Warningf("Error encountered when trying to get data and watch from Zk: %v", err) return err }
func TestZkCustomRule(t *testing.T) { tqsc := tabletservermock.NewController() setUpFakeZk(t) zkcr := NewZkCustomRule(conn) err := zkcr.Open(tqsc, "/zk/fake/customrules/testrules") if err != nil { t.Fatalf("Cannot open zookeeper custom rule service, err=%v", err) } var qrs *tabletserver.QueryRules // Test if we can successfully fetch the original rule (test GetRules) qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } qr := qrs.Find("r1") if qr == nil { t.Fatalf("Expect custom rule r1 to be found, but got nothing, qrs=%v", qrs) } // Test updating rules conn.Set("/zk/fake/customrules/testrules", customRule2, -1) <-time.After(time.Second) //Wait for the polling thread to respond qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } qr = qrs.Find("r2") if qr == nil { t.Fatalf("Expect custom rule r2 to be found, but got nothing, qrs=%v", qrs) } qr = qrs.Find("r1") if qr != nil { t.Fatalf("Custom rule r1 should not be found after r2 is set") } // Test rule path removal conn.Delete("/zk/fake/customrules/testrules", -1) <-time.After(time.Second) qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } if reflect.DeepEqual(qrs, tabletserver.NewQueryRules()) { t.Fatalf("Expect empty rule at this point") } // Test rule path revival conn.Create("/zk/fake/customrules/testrules", "customrule2", 0, zookeeper.WorldACL(zookeeper.PERM_ALL)) conn.Set("/zk/fake/customrules/testrules", customRule2, -1) <-time.After(time.Second) //Wait for the polling thread to respond qrs, _, err = zkcr.GetRules() if err != nil { t.Fatalf("GetRules of ZkCustomRule should always return nil error, but we receive %v", err) } qr = qrs.Find("r2") if qr == nil { t.Fatalf("Expect custom rule r2 to be found, but got nothing, qrs=%v", qrs) } zkcr.Close() }
// NewFileCustomRule returns pointer to new FileCustomRule structure func NewFileCustomRule() (fcr *FileCustomRule) { fcr = new(FileCustomRule) fcr.path = "" fcr.currentRuleSet = tabletserver.NewQueryRules() return fcr }
func TestQueryRules(t *testing.T) { rules := tabletserver.NewQueryRules() err := rules.UnmarshalJSON(rulesJSON) if err != nil { t.Error(err) return } err = framework.Server.SetQueryRules("endtoend", rules) want := "Rule source identifier endtoend is not valid" if err == nil || err.Error() != want { t.Errorf("Error: %v, want %s", err, want) } framework.Server.RegisterQueryRuleSource("endtoend") defer framework.Server.UnRegisterQueryRuleSource("endtoend") err = framework.Server.SetQueryRules("endtoend", rules) if err != nil { t.Error(err) return } rulesJSON := compacted(framework.FetchURL("/debug/query_rules")) want = compacted(`{ "endtoend":[{ "Description": "disallow bindvar 'asdfg'", "Name": "r1", "BindVarConds":[{ "Name": "asdfg", "OnAbsent": false, "Operator": "" }], "Action": "FAIL" }] }`) if rulesJSON != want { t.Errorf("/debug/query_rules:\n%v, want\n%s", rulesJSON, want) } client := framework.NewClient() query := "select * from vitess_test where intval=:asdfg" bv := map[string]interface{}{"asdfg": 1} _, err = client.Execute(query, bv) want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'" if err == nil || err.Error() != want { t.Errorf("Error: %v, want %s", err, want) } _, err = client.StreamExecute(query, bv) want = "error: Query disallowed due to rule: disallow bindvar 'asdfg'" if err == nil || err.Error() != want { t.Errorf("Error: %v, want %s", err, want) } err = framework.Server.SetQueryRules("endtoend", nil) if err != nil { t.Error(err) return } _, err = client.Execute(query, bv) if err != nil { t.Error(err) return } }