Skip to content

Commit 229a68f

Browse files
piquark6046piquark6046
committed
feat: implement budgeted depth check for WeakMap and add regex test guard
1 parent b0bda37 commit 229a68f

4 files changed

Lines changed: 306 additions & 41 deletions

File tree

‎userscript/source/as-weakmap.ts‎

Lines changed: 121 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -8,31 +8,132 @@
88
* - See Git history at https://github.com/FilteringDev/tinyShield for detailed authorship information.
99
*/
1010

11-
import * as Utils from './utils.js'
1211
import { OriginalRegExpTest } from './index.js'
1312

14-
export function CheckDepthInASWeakMap(Args: [object, unknown]) {
15-
const FirstArg = Args[0] as Record<string, unknown>
16-
if (Utils.CountCommonStrings(['device', 'id', 'imp', 'regs', 'site', 'source'], Object.keys(FirstArg)) < 5) {
17-
return false
13+
type CheckDepthResult = { Status: 'matched' } | { Status: 'not-matched' } | { Status: 'too-expensive' } | { Status: 'unsafe-object'; Reason: unknown }
14+
15+
type CheckBudget = {
16+
MaxTopLevelKeys: number;
17+
MaxArrayItems: number;
18+
MaxInnerKeysPerObject: number;
19+
MaxOperations: number;
20+
};
21+
22+
const DefaultBudget: CheckBudget = {
23+
MaxTopLevelKeys: 300,
24+
MaxArrayItems: 1_000,
25+
MaxInnerKeysPerObject: 100,
26+
MaxOperations: 10_000,
27+
}
28+
29+
const ImportantKeys = ['device', 'id', 'imp', 'regs', 'site', 'source'] as const
30+
const PropertyIsEnumerable = Object.prototype.propertyIsEnumerable
31+
32+
function HasOwnEnumerableStringKey(Obj: object, Key: string): boolean {
33+
return PropertyIsEnumerable.call(Obj, Key)
34+
}
35+
36+
function CountCommonKnownKeys(Obj: object): number {
37+
let Count = 0
38+
39+
for (const Key of ImportantKeys) {
40+
if (HasOwnEnumerableStringKey(Obj, Key)) {
41+
Count++
42+
}
1843
}
1944

20-
const ASBannerFrameIdRegExp = /^[0-9]+\/[a-zA-Z0-9]+\/[a-zA-Z0-9]+\/[a-z0-9-\(\)]+\/[a-zA-Z0-9_]+_slot[0-9]+_+/
21-
const ASBannerFrameKey = Object.keys(FirstArg).find(Arg => {
22-
const Candidate = FirstArg[Arg]
23-
if (!Array.isArray(Candidate)) {
24-
return false
45+
return Count
46+
}
47+
48+
export function CheckDepthInASWeakMapBudgeted(
49+
Args: [object, unknown],
50+
Budget: CheckBudget = DefaultBudget,
51+
): CheckDepthResult {
52+
let Operations = 0
53+
54+
const Tick = (Amount = 1): boolean => {
55+
Operations += Amount
56+
return Operations <= Budget.MaxOperations
57+
}
58+
59+
try {
60+
const FirstArg = Args[0] as Record<string, unknown>
61+
62+
if (!Tick(ImportantKeys.length)) {
63+
return { Status: 'too-expensive' }
2564
}
26-
return Candidate.filter(SubArg => {
27-
if (typeof SubArg !== 'object' || SubArg === null) {
28-
return false
65+
66+
if (CountCommonKnownKeys(FirstArg) < 5) {
67+
return { Status: 'not-matched' }
68+
}
69+
70+
const AsBannerFrameIdRegExp =
71+
/^[0-9]+\/[a-zA-Z0-9]+\/[a-zA-Z0-9]+\/[a-z0-9-\(\)]+\/[a-zA-Z0-9_]+_slot[0-9]+_+/
72+
73+
let TopLevelKeyCount = 0
74+
75+
for (const Arg in FirstArg) {
76+
if (!HasOwnEnumerableStringKey(FirstArg, Arg)) {
77+
continue
2978
}
30-
return Object.keys(SubArg).some(InnerArg => OriginalRegExpTest.call(ASBannerFrameIdRegExp, InnerArg) as boolean)
31-
}).length >= 1
32-
})
33-
if (typeof ASBannerFrameKey === 'undefined') {
34-
return false
35-
}
3679

37-
return true
80+
TopLevelKeyCount++
81+
82+
if (
83+
TopLevelKeyCount > Budget.MaxTopLevelKeys ||
84+
!Tick()
85+
) {
86+
return { Status: 'too-expensive' }
87+
}
88+
89+
const Candidate = FirstArg[Arg]
90+
91+
if (!Array.isArray(Candidate)) {
92+
continue
93+
}
94+
95+
const MaxArrayIndex = Math.min(Candidate.length, Budget.MaxArrayItems)
96+
97+
for (let I = 0; I < MaxArrayIndex; I++) {
98+
if (!Tick()) {
99+
return { Status: 'too-expensive' }
100+
}
101+
102+
const SubArg = Candidate[I]
103+
104+
if (typeof SubArg !== 'object' || SubArg === null) {
105+
continue
106+
}
107+
108+
let InnerKeyCount = 0
109+
110+
for (const InnerArg in SubArg as Record<string, unknown>) {
111+
if (!HasOwnEnumerableStringKey(SubArg, InnerArg)) {
112+
continue
113+
}
114+
115+
InnerKeyCount++
116+
117+
if (
118+
InnerKeyCount > Budget.MaxInnerKeysPerObject ||
119+
!Tick()
120+
) {
121+
return { Status: 'too-expensive' }
122+
}
123+
124+
if (OriginalRegExpTest.call(AsBannerFrameIdRegExp, InnerArg) as boolean) {
125+
return { Status: 'matched' }
126+
}
127+
}
128+
}
129+
130+
if (Candidate.length > Budget.MaxArrayItems) {
131+
return { Status: 'too-expensive' }
132+
}
133+
}
134+
135+
return { Status: 'not-matched' }
136+
} catch (error) {
137+
return { Status: 'unsafe-object', Reason: error }
138+
}
38139
}

‎userscript/source/index.ts‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ declare const unsafeWindow: unsafeWindow
1515
const BrowserWindow = typeof unsafeWindow !== 'undefined' ? unsafeWindow : window
1616
const UserscriptName = 'tinyShield'
1717

18-
import { CheckDepthInASWeakMap } from './as-weakmap.js'
18+
import { CheckDepthInASWeakMapBudgeted } from './as-weakmap.js'
19+
import { ShouldSkipRegExpTest } from './regexp-cheap-guard.js'
1920
import { SafeArrayToString } from './safe-ArrayToString.js'
2021

2122
export const OriginalRegExpTest = BrowserWindow.RegExp.prototype.test
@@ -49,7 +50,8 @@ BrowserWindow.Map.prototype.get = new Proxy(BrowserWindow.Map.prototype.get, {
4950
}
5051

5152
let ArgText = SafeArrayToString(Args, { OriginalArrayMap, OriginalString, OriginalArrayJoin, OriginalObjectGetPrototypeOf })
52-
if (ASInitPositiveRegExps.filter(ASInitPositiveRegExp => ASInitPositiveRegExp.filter(Index => OriginalRegExpTest.call(Index, ArgText) as boolean).length >= 2).length === 1) {
53+
54+
if (!ShouldSkipRegExpTest(ArgText) && ASInitPositiveRegExps.filter(ASInitPositiveRegExp => ASInitPositiveRegExp.filter(Index => OriginalRegExpTest.call(Index, ArgText) as boolean).length >= 2).length === 1) {
5355
console.debug(`[${UserscriptName}]: Map.prototype.get:`, ThisArg, Args)
5456
throw new Error()
5557
}
@@ -71,7 +73,8 @@ BrowserWindow.Map.prototype.set = new Proxy(BrowserWindow.Map.prototype.set, {
7173
apply(Target: (key: unknown, value: unknown) => Map<unknown, unknown>, ThisArg: Map<unknown, unknown>, Args: [unknown, unknown]) {
7274
let ArgText = ''
7375
ArgText = SafeArrayToString(Args, { OriginalArrayMap, OriginalString, OriginalArrayJoin, OriginalObjectGetPrototypeOf })
74-
if (ASReinsertedAdvInvenPositiveRegExps.filter(ASReinsertedAdvInvenPositiveRegExp => ASReinsertedAdvInvenPositiveRegExp.filter(Index => OriginalRegExpTest.call(Index, ArgText) as boolean).length >= 3).length === 1) {
76+
77+
if (!ShouldSkipRegExpTest(ArgText) && ASReinsertedAdvInvenPositiveRegExps.filter(ASReinsertedAdvInvenPositiveRegExp => ASReinsertedAdvInvenPositiveRegExp.filter(Index => OriginalRegExpTest.call(Index, ArgText) as boolean).length >= 3).length === 1) {
7578
console.debug(`[${UserscriptName}]: Map.prototype.set:`, ThisArg, Args)
7679
throw new Error()
7780
}
@@ -81,9 +84,19 @@ BrowserWindow.Map.prototype.set = new Proxy(BrowserWindow.Map.prototype.set, {
8184

8285
BrowserWindow.WeakMap.prototype.set = new Proxy(BrowserWindow.WeakMap.prototype.set, {
8386
apply(Target: (key: object, value: unknown) => WeakMap<object, unknown>, ThisArg: WeakMap<object, unknown>, Args: [object, unknown]) {
84-
if (CheckDepthInASWeakMap(Args)) {
85-
console.debug(`[${UserscriptName}]: WeakMap.prototype.set:`, ThisArg, Args)
86-
throw new Error()
87+
let CheckResult = CheckDepthInASWeakMapBudgeted(Args)
88+
switch (CheckResult.Status) {
89+
case 'matched':
90+
console.debug(`[${UserscriptName}]: WeakMap.prototype.set:`, ThisArg, Args)
91+
throw new Error()
92+
case 'not-matched':
93+
break
94+
case 'too-expensive':
95+
console.warn(`[${UserscriptName}]: WeakMap.prototype.set: Check too expensive:`, ThisArg, Args)
96+
break
97+
case 'unsafe-object':
98+
console.warn(`[${UserscriptName}]: WeakMap.prototype.set: Unsafe object:`, ThisArg, Args, CheckResult.Reason)
99+
break
87100
}
88101

89102
return Reflect.apply(Target, ThisArg, Args)

‎userscript/source/regexp-cheap-guard.ts‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
const MinimumInputLength = 512
2+
const MaximumSampleLength = 4096
3+
4+
const MinimumObjectTokenCount = 10
5+
const MinimumCommaCount = 50
6+
const MinimumSplitPartCount = 50
7+
8+
const UselessPartRatioThreshold = 0.65
9+
10+
const ObjectTokenMarker = '[object '
11+
12+
const MaskedDumpMarkers = [
13+
'[object ',
14+
'[native code]',
15+
'function ',
16+
'=>',
17+
'HTMLElement',
18+
'HTML',
19+
'ShadowRoot',
20+
'Comment',
21+
'Text',
22+
'Object',
23+
] as const
24+
25+
const DomObjectSuffixMarkers = [
26+
'HTMLElement]',
27+
'ShadowRoot]',
28+
'Comment]',
29+
'Text]',
30+
] as const
31+
32+
const LiteralUselessPartSet: ReadonlySet<string> = new Set([
33+
'',
34+
'true',
35+
'false',
36+
])
37+
38+
export function SafeRegExpTest(RegExpObject: RegExp, Input: string): boolean {
39+
if (ShouldSkipRegExpTest(Input)) {
40+
return false
41+
}
42+
43+
return RegExpObject.test(Input)
44+
}
45+
46+
export function ShouldSkipRegExpTest(Input: string): boolean {
47+
if (Input.length < MinimumInputLength) {
48+
return false
49+
}
50+
51+
const Sample = Input.length > MaximumSampleLength
52+
? Input.slice(0, MaximumSampleLength)
53+
: Input
54+
55+
if (!Sample.includes(ObjectTokenMarker)) {
56+
return false
57+
}
58+
59+
const ObjectTokenCount = CountOccurrences(Sample, ObjectTokenMarker)
60+
61+
if (ObjectTokenCount < MinimumObjectTokenCount) {
62+
return false
63+
}
64+
65+
const CommaCount = CountOccurrences(Sample, ',')
66+
67+
if (CommaCount < MinimumCommaCount) {
68+
return false
69+
}
70+
71+
const HasMaskedDumpMarker = MaskedDumpMarkers.some((Marker) =>
72+
Sample.includes(Marker),
73+
)
74+
75+
if (!HasMaskedDumpMarker) {
76+
return false
77+
}
78+
79+
const Parts = Sample.split(',')
80+
81+
if (Parts.length < MinimumSplitPartCount) {
82+
return false
83+
}
84+
85+
let UselessPartCount = 0
86+
87+
for (const RawPart of Parts) {
88+
const Part = RawPart.trim()
89+
90+
if (IsLikelyUselessMaskedPart(Part)) {
91+
UselessPartCount += 1
92+
}
93+
}
94+
95+
const UselessPartRatio = UselessPartCount / Parts.length
96+
97+
return UselessPartRatio >= UselessPartRatioThreshold
98+
}
99+
100+
function IsLikelyUselessMaskedPart(Part: string): boolean {
101+
if (LiteralUselessPartSet.has(Part)) {
102+
return true
103+
}
104+
105+
if (IsIntegerLikeString(Part)) {
106+
return true
107+
}
108+
109+
if (Part.startsWith(ObjectTokenMarker)) {
110+
return true
111+
}
112+
113+
if (DomObjectSuffixMarkers.some((Marker) => Part.includes(Marker))) {
114+
return true
115+
}
116+
117+
return (
118+
Part.includes('[native code]') ||
119+
Part.startsWith('function ') ||
120+
Part.startsWith('()=>')
121+
)
122+
}
123+
124+
function IsIntegerLikeString(Value: string): boolean {
125+
if (Value.length === 0) {
126+
return false
127+
}
128+
129+
let StartIndex = 0
130+
131+
if (Value[0] === '-') {
132+
if (Value.length === 1) {
133+
return false
134+
}
135+
136+
StartIndex = 1
137+
}
138+
139+
for (let Index = StartIndex; Index < Value.length; Index += 1) {
140+
const CharCode = Value.charCodeAt(Index)
141+
142+
if (CharCode < 48 || CharCode > 57) {
143+
return false
144+
}
145+
}
146+
147+
return true
148+
}
149+
150+
function CountOccurrences(Haystack: string, Needle: string): number {
151+
let Count = 0
152+
let Index = 0
153+
154+
while (true) {
155+
const FoundIndex = Haystack.indexOf(Needle, Index)
156+
157+
if (FoundIndex === -1) {
158+
break
159+
}
160+
161+
Count += 1
162+
Index = FoundIndex + Needle.length
163+
}
164+
165+
return Count
166+
}

‎userscript/source/utils.ts‎

Lines changed: 0 additions & 15 deletions
This file was deleted.

0 commit comments

Comments
 (0)