Skip to content

Commit 8f050a4

Browse files
Mzack9999Mzack9999
authored
Merge pull request #249 from projectdiscovery/dwisiswant0/refactor/adds-optionalReadLimit-params
refactor: adds optionalReadLimit params
2 parents e471587 + c41e3e9 commit 8f050a4

File tree

2 files changed

+110
-50
lines changed

2 files changed

+110
-50
lines changed

‎dsl.go

Lines changed: 90 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -282,21 +282,39 @@ func init() {
282282

283283
return buffer.String(), nil
284284
}))
285-
MustAddFunction(NewWithPositionalArgs("gzip_decode", 1, true, func(args ...interface{}) (interface{}, error) {
286-
reader, err := gzip.NewReader(strings.NewReader(args[0].(string)))
287-
if err != nil {
288-
return "", err
289-
}
290-
limitReader := io.LimitReader(reader, DefaultMaxDecompressionSize)
285+
MustAddFunction(NewWithSingleSignature("gzip_decode",
286+
"(data string, optionalReadLimit int) string",
287+
true,
288+
func(args ...interface{}) (interface{}, error) {
289+
if len(args) == 0 {
290+
return nil, ErrInvalidDslFunction
291+
}
291292

292-
data, err := io.ReadAll(limitReader)
293-
if err != nil {
293+
argData := toString(args[0])
294+
readLimit := DefaultMaxDecompressionSize
295+
296+
if len(args) > 1 {
297+
if limit, ok := args[1].(float64); ok {
298+
readLimit = int64(limit)
299+
}
300+
}
301+
302+
reader, err := gzip.NewReader(strings.NewReader(argData))
303+
if err != nil {
304+
return "", err
305+
}
306+
limitReader := io.LimitReader(reader, readLimit)
307+
308+
data, err := io.ReadAll(limitReader)
309+
if err != nil && err != io.EOF {
310+
_ = reader.Close()
311+
312+
return "", err
313+
}
294314
_ = reader.Close()
295-
return "", err
296-
}
297-
_ = reader.Close()
298-
return string(data), nil
299-
}))
315+
316+
return string(data), nil
317+
}))
300318
MustAddFunction(NewWithPositionalArgs("zlib", 1, true, func(args ...interface{}) (interface{}, error) {
301319
buffer := &bytes.Buffer{}
302320
writer := zlib.NewWriter(buffer)
@@ -308,21 +326,39 @@ func init() {
308326

309327
return buffer.String(), nil
310328
}))
311-
MustAddFunction(NewWithPositionalArgs("zlib_decode", 1, true, func(args ...interface{}) (interface{}, error) {
312-
reader, err := zlib.NewReader(strings.NewReader(args[0].(string)))
313-
if err != nil {
314-
return "", err
315-
}
316-
limitReader := io.LimitReader(reader, DefaultMaxDecompressionSize)
329+
MustAddFunction(NewWithSingleSignature("zlib_decode",
330+
"(data string, optionalReadLimit int) string",
331+
true,
332+
func(args ...interface{}) (interface{}, error) {
333+
if len(args) == 0 {
334+
return nil, ErrInvalidDslFunction
335+
}
317336

318-
data, err := io.ReadAll(limitReader)
319-
if err != nil {
337+
argData := toString(args[0])
338+
readLimit := DefaultMaxDecompressionSize
339+
340+
if len(args) > 1 {
341+
if limit, ok := args[1].(float64); ok {
342+
readLimit = int64(limit)
343+
}
344+
}
345+
346+
reader, err := zlib.NewReader(strings.NewReader(argData))
347+
if err != nil {
348+
return "", err
349+
}
350+
limitReader := io.LimitReader(reader, readLimit)
351+
352+
data, err := io.ReadAll(limitReader)
353+
if err != nil && err != io.EOF {
354+
_ = reader.Close()
355+
356+
return "", err
357+
}
320358
_ = reader.Close()
321-
return "", err
322-
}
323-
_ = reader.Close()
324-
return string(data), nil
325-
}))
359+
360+
return string(data), nil
361+
}))
326362

327363
MustAddFunction(NewWithPositionalArgs("deflate", 1, true, func(args ...interface{}) (interface{}, error) {
328364
buffer := &bytes.Buffer{}
@@ -338,18 +374,36 @@ func init() {
338374

339375
return buffer.String(), nil
340376
}))
341-
MustAddFunction(NewWithPositionalArgs("inflate", 1, true, func(args ...interface{}) (interface{}, error) {
342-
reader := flate.NewReader(strings.NewReader(args[0].(string)))
343-
limitReader := io.LimitReader(reader, DefaultMaxDecompressionSize)
377+
MustAddFunction(NewWithSingleSignature("inflate",
378+
"(data string, optionalReadLimit int) string",
379+
true,
380+
func(args ...interface{}) (interface{}, error) {
381+
if len(args) == 0 {
382+
return nil, ErrInvalidDslFunction
383+
}
344384

345-
data, err := io.ReadAll(limitReader)
346-
if err != nil {
385+
argData := toString(args[0])
386+
readLimit := DefaultMaxDecompressionSize
387+
388+
if len(args) > 1 {
389+
if limit, ok := args[1].(float64); ok {
390+
readLimit = int64(limit)
391+
}
392+
}
393+
394+
reader := flate.NewReader(strings.NewReader(argData))
395+
limitReader := io.LimitReader(reader, readLimit)
396+
397+
data, err := io.ReadAll(limitReader)
398+
if err != nil && err != io.EOF {
399+
_ = reader.Close()
400+
401+
return "", err
402+
}
347403
_ = reader.Close()
348-
return "", err
349-
}
350-
_ = reader.Close()
351-
return string(data), nil
352-
}))
404+
405+
return string(data), nil
406+
}))
353407

354408
MustAddFunction(NewWithSingleSignature("date_time",
355409
"(dateTimeFormat string, optionalUnixTime interface{}) string",

‎dsl_test.go

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -256,15 +256,15 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
256256
generate_java_gadget(arg1, arg2, arg3 interface{}) interface{}
257257
generate_jwt(jsonString, algorithm, optionalSignature string, optionalMaxAgeUnix interface{}) string
258258
gzip(arg1 interface{}) interface{}
259-
gzip_decode(arg1 interface{}) interface{}
259+
gzip_decode(data string, optionalReadLimit int) string
260260
hex_decode(arg1 interface{}) interface{}
261261
hex_encode(arg1 interface{}) interface{}
262262
hex_to_dec(arg1 interface{}) interface{}
263263
hmac(arg1, arg2, arg3 interface{}) interface{}
264264
html_escape(arg1 interface{}) interface{}
265265
html_unescape(arg1 interface{}) interface{}
266266
index(arg1, arg2 interface{}) interface{}
267-
inflate(arg1 interface{}) interface{}
267+
inflate(data string, optionalReadLimit int) string
268268
ip_format(arg1, arg2 interface{}) interface{}
269269
jarm(arg1 interface{}) interface{}
270270
join(separator string, elements ...interface{}) string
@@ -329,7 +329,7 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
329329
xor(args ...interface{}) interface{}
330330
zip(file_entry string, content string, ... ) []byte
331331
zlib(arg1 interface{}) interface{}
332-
zlib_decode(arg1 interface{}) interface{}
332+
zlib_decode(data string, optionalReadLimit int) string
333333
`
334334

335335
signatures := GetPrintableDslFunctionSignatures(true)
@@ -435,17 +435,23 @@ func TestDslExpressions(t *testing.T) {
435435
`uniq("ab", "cd", "12", "34", "12", "cd")`: []string{"ab", "cd", "12", "34"},
436436
`join(" ", uniq("ab", "cd", "12", "34", "12", "cd"))`: "ab cd 12 34",
437437
`join(", ", split(hex_encode("abcdefg"), 2))`: "61, 62, 63, 64, 65, 66, 67",
438-
`json_minify("{ \"name\": \"John Doe\", \"foo\": \"bar\" }")`: "{\"foo\":\"bar\",\"name\":\"John Doe\"}",
439-
`json_prettify("{\"foo\":\"bar\",\"name\":\"John Doe\"}")`: "{\n \"foo\": \"bar\",\n \"name\": \"John Doe\"\n}",
440-
`ip_format('127.0.0.1', '1')`: "127.0.0.1",
441-
`ip_format('127.0.0.1', '3')`: "0177.0.0.01",
442-
`ip_format('127.0.0.1', '5')`: "2130706433",
443-
`ip_format('127.0.1.0', '11')`: "127.0.256",
444-
"unpack('>I', '\xac\xd7\t\xd0')": -272646673,
445-
"xor('\x01\x02', '\x02\x01')": []uint8([]byte{0x3, 0x3}),
446-
`count("projectdiscovery", "e")`: 2,
447-
`concat(to_title("pRoJeCt"), to_title("diScOvErY"))`: "ProjectDiscovery",
448-
`concat(to_title("welcome "), "to", to_title(" watch"), to_title("mojo"))`: "Welcome to WatchMojo",
438+
`json_minify("{ \"name\": \"John Doe\", \"foo\": \"bar\" }")`: "{\"foo\":\"bar\",\"name\":\"John Doe\"}",
439+
`json_prettify("{\"foo\":\"bar\",\"name\":\"John Doe\"}")`: "{\n \"foo\": \"bar\",\n \"name\": \"John Doe\"\n}",
440+
`ip_format('127.0.0.1', '1')`: "127.0.0.1",
441+
`ip_format('127.0.0.1', '3')`: "0177.0.0.01",
442+
`ip_format('127.0.0.1', '5')`: "2130706433",
443+
`ip_format('127.0.1.0', '11')`: "127.0.256",
444+
"unpack('>I', '\xac\xd7\t\xd0')": -272646673,
445+
"xor('\x01\x02', '\x02\x01')": []uint8([]byte{0x3, 0x3}),
446+
`count("projectdiscovery", "e")`: 2,
447+
`concat(to_title("pRoJeCt"), to_title("diScOvErY"))`: "ProjectDiscovery",
448+
`concat(to_title("welcome "), "to", to_title(" watch"), to_title("mojo"))`: "Welcome to WatchMojo",
449+
`zlib_decode(hex_decode("789cf248cdc9c907040000ffff058c01f5"), 4)`: "Hell",
450+
`gzip_decode(hex_decode("1f8b08000000000000fff248cdc9c907040000ffff8289d1f705000000"), 4)`: "Hell",
451+
`inflate(hex_decode("f248cdc9c907040000ffff"), 4)`: "Hell",
452+
`zlib_decode(hex_decode("789cf248cdc9c907040000ffff058c01f5"), 100)`: "Hello",
453+
`gzip_decode(hex_decode("1f8b08000000000000fff248cdc9c907040000ffff8289d1f705000000"), 100)`: "Hello",
454+
`inflate(hex_decode("f248cdc9c907040000ffff"), 100)`: "Hello",
449455
}
450456

451457
testDslExpressions(t, dslExpressions)

0 commit comments

Comments
 (0)