@@ -330,6 +330,7 @@ import (
330330// - noticetrace: runs the query and compares only the notices that
331331// appear. Cannot be combined with kvtrace.
332332// - nodeidx=N: runs the query on node N of the cluster.
333+ // - allowunsafe: allows access to unsafe internals during execution.
333334//
334335// The label is optional. If specified, the test runner stores a hash
335336// of the results of the query under the given label. If the label is
@@ -969,6 +970,9 @@ type logicQuery struct {
969970 // roundFloatsInStringsSigFigs specifies the number of significant figures
970971 // to round floats embedded in strings to where zero means do not round.
971972 roundFloatsInStringsSigFigs int
973+
974+ // allowUnsafe indicates whether unsafe operations are allowed during execution.
975+ allowUnsafe bool
972976}
973977
974978var allowedKVOpTypes = []string {
@@ -1119,6 +1123,10 @@ type logicTest struct {
11191123 // retryDuration is the maximum duration to retry a statement when using
11201124 // the retry directive.
11211125 retryDuration time.Duration
1126+
1127+ // allowUnsafe is a variable which controls whether the test can access
1128+ // unsafe internals.
1129+ allowUnsafe bool
11221130}
11231131
11241132func (t * logicTest ) t () * testing.T {
@@ -1493,6 +1501,7 @@ func (t *logicTest) newCluster(
14931501 DisableOptimizerRuleProbability : * disableOptRuleProbability ,
14941502 OptimizerCostPerturbation : * optimizerCostPerturbation ,
14951503 ForceProductionValues : serverArgs .ForceProductionValues ,
1504+ UnsafeOverride : func () * bool { return & t .allowUnsafe },
14961505 }
14971506 knobs .SQLExecutor = & sql.ExecutorTestingKnobs {
14981507 DeterministicExplain : true ,
@@ -2115,6 +2124,18 @@ func (c knobOptSynchronousEventLog) apply(args *base.TestingKnobs) {
21152124 args .EventLog .(* eventlog.EventLogTestingKnobs ).SyncWrites = true
21162125}
21172126
2127+ // knobOptAllowUnsafe always allows access to the unsafe internals.
2128+ type knobOptAllowUnsafe struct {}
2129+
2130+ var _ knobOpt = knobOptAllowUnsafe {}
2131+
2132+ // apply implements the clusterOpt interface.
2133+ func (c knobOptAllowUnsafe ) apply (args * base.TestingKnobs ) {
2134+ e := args .SQLEvalContext .(* eval.TestingKnobs )
2135+ v := true
2136+ e .UnsafeOverride = func () * bool { return & v }
2137+ }
2138+
21182139// clusterOptIgnoreStrictGCForTenants corresponds to the
21192140// ignore-tenant-strict-gc-enforcement directive.
21202141type clusterOptIgnoreStrictGCForTenants struct {}
@@ -2268,6 +2289,8 @@ func readKnobOptions(t *testing.T, path string) []knobOpt {
22682289 res = append (res , knobOptDisableCorpusGeneration {})
22692290 case "sync-event-log" :
22702291 res = append (res , knobOptSynchronousEventLog {})
2292+ case "allow-unsafe" :
2293+ res = append (res , knobOptAllowUnsafe {})
22712294 default :
22722295 t .Fatalf ("unrecognized knob option: %s" , opt )
22732296 }
@@ -2954,6 +2977,9 @@ func (t *logicTest) processSubtest(
29542977 case "async" :
29552978 query .expectAsync = true
29562979
2980+ case "allowunsafe" :
2981+ query .allowUnsafe = true
2982+
29572983 default :
29582984 if strings .HasPrefix (opt , "round-in-strings" ) {
29592985 significantFigures , err := floatcmp .ParseRoundInStringsDirective (opt )
@@ -3599,6 +3625,7 @@ func (t *logicTest) unexpectedError(sql string, pos string, err error) (bool, er
35993625var uniqueHashPattern = regexp .MustCompile (`UNIQUE.*USING\s+HASH` )
36003626
36013627func (t * logicTest ) execStatement (stmt logicStatement , disableCFMutator bool ) (bool , error ) {
3628+ defer t .setSafetyGate (stmt .sql , false )()
36023629 db := t .db
36033630 t .noticeBuffer = nil
36043631 if * showSQL {
@@ -3722,6 +3749,7 @@ func (t *logicTest) hashResults(results []string) (string, error) {
37223749}
37233750
37243751func (t * logicTest ) execQuery (query logicQuery ) error {
3752+ defer t .setSafetyGate (query .sql , query .allowUnsafe )()
37253753 if * showSQL {
37263754 t .outf ("%s;" , query .sql )
37273755 }
@@ -4609,6 +4637,7 @@ func RunLogicTest(
46094637 perErrorSummary : make (map [string ][]string ),
46104638 rng : rng ,
46114639 declarativeCorpusCollector : cc ,
4640+ allowUnsafe : true ,
46124641 }
46134642 if * printErrorSummary {
46144643 defer lt .printErrorSummary ()
@@ -4884,3 +4913,23 @@ func locateCockroachPredecessor(version string) (string, error) {
48844913 }
48854914 return path , nil
48864915}
4916+
4917+ // setSafetyGate is a utility function which controls whether access to the unsafe
4918+ // internals is allowed by the contained sql statement. The reasoning behind it is
4919+ // as follows. We want queries which explicitly access unsafe internals to have
4920+ // access to them, but we want to prevent indirect access wherever possible.
4921+ // Indirect access can be described as any query which doesn't reference an unsafe
4922+ // object, but still accesses it under the hood. We want to flush out these cases
4923+ // so that users can never execute statements which look safe, but block when executed.
4924+ func (t * logicTest ) setSafetyGate (sql string , skip bool ) func () {
4925+ sql = strings .ToLower (sql )
4926+ explicitlyUnsafe := strings .Contains (sql , "crdb_internal." ) || strings .Contains (sql , "system." )
4927+ if skip || explicitlyUnsafe {
4928+ return func () {}
4929+ }
4930+
4931+ t .allowUnsafe = false
4932+ return func () {
4933+ t .allowUnsafe = true
4934+ }
4935+ }
0 commit comments