Download runtime/go/parse.go from Dibachain/DibaGuard: direct link, hf CLI and curl.
- Browser
- Download file 13.7 kB
-
https://huggingface.co/Dibachain/DibaGuard/resolve/main/runtime/go/parse.go
- Command line
-
hf download hf://Dibachain/DibaGuard/runtime/go/parse.go
-
curl -L -o parse.go https://huggingface.co/Dibachain/DibaGuard/resolve/main/runtime/go/parse.go
13.7 kB
| package dibaguard | |
| import ( | |
| "bytes" | |
| "encoding/json" | |
| "math" | |
| "net/http" | |
| "net/url" | |
| "regexp" | |
| "sort" | |
| "strconv" | |
| "strings" | |
| "unicode" | |
| ) | |
| // The Python request parser (dibaguard/http/parse.py Request.params / model_text) in Go: query pairs, form, JSON | |
| // (every occurrence of a duplicate key, Python's number spelling), multipart (every field; a file's name and the | |
| // head of its content), cookies, the four standard headers and any other header whose value looks like a payload | |
| // (every value of a repeated header), all under the same judging budget - so Canonical is the Python model text. | |
| const ( | |
| parseMaxPart = 8192 // parse.MAX_PART | |
| parseChunkOverlap = 512 // parse.CHUNK_OVERLAP | |
| parseMaxText = 32768 // parse.MAX_TEXT | |
| parseExtraText = 32768 // parse.EXTRA_TEXT | |
| parseMaxFields = 2000 // parse.MAX_FIELDS | |
| parseMaxDepth = 512 // parse.MAX_DEPTH | |
| parseMaxFile = 8192 // parse.MAX_FILE | |
| parseRestMax = 1 << 20 // parse.REST_MAX | |
| ) | |
| var ( | |
| stdHeaders = []string{"user-agent", "referer", "x-forwarded-for", "content-type"} | |
| neverHeaders = map[string]bool{"cookie": true, "host": true, "content-length": true, "connection": true, "accept-encoding": true, | |
| "cache-control": true, "pragma": true, "upgrade-insecure-requests": true} | |
| hdrSuspicious = regexp.MustCompile("['\"<>`]|\\$\\{|\\{\\{|--|/\\*|\\.\\./|(?:^|[^a-z0-9_])(?:union|select|sleep|benchmark|jndi|script|javascript|vbscript|document|onerror|eval|exec" + | |
| "|cmd|wget|curl|waitfor|extractvalue|updatexml|information_schema|passwd|whoami)(?:$|[^a-z0-9_])") | |
| keepWords = map[string]bool{} | |
| reBound = regexp.MustCompile(`(?i)boundary=(?:"([^"]+)"|([^";,\s]+))`) | |
| reDispNm = regexp.MustCompile(`(?i)(?:^|;)\s*name\s*=\s*(?:"([^"]*)"|([^;\s]*))`) | |
| reDispFn = regexp.MustCompile(`(?i)(?:^|;)\s*filename\s*=\s*(?:"([^"]*)"|([^;\s]*))`) | |
| ) | |
| func init() { | |
| for _, w := range strings.Fields("or and xor not union select sleep benchmark waitfor having exec execute declare drop insert update delete truncate " + | |
| "shutdown grant information_schema pg_sleep extractvalue updatexml dual concat char chr rlike regexp like null jndi script javascript passwd " + | |
| "whoami cmd eval system") { | |
| keepWords[w] = true | |
| } | |
| } | |
| // keepPart is parse._keep: can this canonical text carry a payload? | |
| func keepPart(s string) bool { | |
| if strings.Contains(s, "--") || strings.Contains(s, "..") { | |
| return true | |
| } | |
| for _, r := range s { | |
| if !(unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' || r == ' ' || r == '.' || r == '-') { | |
| return true | |
| } | |
| } | |
| for i := 0; i < len(s); { | |
| c := s[i] | |
| if 'a' <= c && c <= 'z' || '0' <= c && c <= '9' || c == '_' { | |
| j := i | |
| for j < len(s) && ('a' <= s[j] && s[j] <= 'z' || '0' <= s[j] && s[j] <= '9' || s[j] == '_') { | |
| j++ | |
| } | |
| if keepWords[s[i:j]] { | |
| return true | |
| } | |
| i = j | |
| continue | |
| } | |
| i++ | |
| } | |
| return false | |
| } | |
| // chunks is parse._chunks (code points). | |
| func chunks(v string) []string { | |
| n := runeLen(v) | |
| if n <= parseMaxPart { | |
| return []string{v} | |
| } | |
| rs := []rune(v) | |
| step := parseMaxPart - parseChunkOverlap | |
| var out []string | |
| for i := 0; i < n-parseChunkOverlap; i += step { | |
| e := i + parseMaxPart | |
| if e > n { | |
| e = n | |
| } | |
| out = append(out, string(rs[i:e])) | |
| } | |
| return out | |
| } | |
| // textBuilder is parse._Text. | |
| type textBuilder struct { | |
| parts []string | |
| n, extra, over, rawLeft, d int | |
| } | |
| func (t *textBuilder) add(prefix, value string) { | |
| pl, vl := runeLen(prefix), runeLen(value) | |
| if t.n+pl+vl+3 <= parseMaxText { | |
| t.parts = append(t.parts, prefix+value) | |
| t.n += pl + vl + 3 | |
| return | |
| } | |
| for _, c := range chunks(value) { | |
| k := pl + runeLen(c) + 3 | |
| switch { | |
| case t.n+k <= parseMaxText: | |
| t.parts = append(t.parts, prefix+c) | |
| t.n += k | |
| case keepPart(c): | |
| if t.extra+k <= parseExtraText { | |
| t.parts = append(t.parts, prefix+c) | |
| t.extra += k | |
| } else { | |
| t.over += runeLen(c) | |
| } | |
| } | |
| } | |
| } | |
| func (t *textBuilder) addRaw(prefix, raw string) { | |
| n := runeLen(raw) | |
| if n > t.rawLeft { | |
| t.over += n - t.rawLeft | |
| raw = runeCut(raw, t.rawLeft) | |
| n = t.rawLeft | |
| } | |
| t.rawLeft -= n | |
| if n == 0 { | |
| return | |
| } | |
| rs := []rune(raw) | |
| step := canonMaxLen - parseChunkOverlap | |
| hi := n - parseChunkOverlap | |
| if hi < 1 { | |
| hi = 1 | |
| } | |
| for i := 0; i < hi; i += step { | |
| e := i + canonMaxLen | |
| if e > n { | |
| e = n | |
| } | |
| a, p, ok, r := canonText(string(rs[i:e]), false) | |
| if r > t.d { | |
| t.d = r | |
| } | |
| t.add(prefix, a) | |
| if ok { | |
| t.add(prefix, p) | |
| } | |
| } | |
| } | |
| func (t *textBuilder) text() string { | |
| if t.over > 0 { | |
| t.parts = append(t.parts, "x:oversized="+strconv.Itoa(t.over)) | |
| } | |
| return strings.Join(t.parts, " | ") | |
| } | |
| type param struct{ where, k, v string } | |
| // pairs is parse._pairs: split on '&' only, a pair at its first '='; nothing decoded yet. | |
| func pairs(s string, where string, out []param) []param { | |
| for _, p := range strings.Split(s, "&") { | |
| if p == "" { | |
| continue | |
| } | |
| k, v, _ := strings.Cut(p, "=") | |
| out = append(out, param{where, k, v}) | |
| } | |
| return out | |
| } | |
| // pyFloat is Python's repr(float). | |
| func pyFloat(f float64) string { | |
| switch { | |
| case math.IsInf(f, 1): | |
| return "inf" | |
| case math.IsInf(f, -1): | |
| return "-inf" | |
| } | |
| e := strconv.FormatFloat(f, 'e', -1, 64) // [-]d[.ddd]e±XX, the shortest digits that round-trip | |
| sign := "" | |
| if e[0] == '-' { | |
| sign, e = "-", e[1:] | |
| } | |
| mant, exps, _ := strings.Cut(e, "e") | |
| digits := strings.Replace(mant, ".", "", 1) | |
| x, _ := strconv.Atoi(exps) | |
| decpt := x + 1 | |
| if decpt <= -4 || decpt > 16 { | |
| m := digits[:1] | |
| if len(digits) > 1 { | |
| m += "." + digits[1:] | |
| } | |
| es := "+" | |
| if x < 0 { | |
| es, x = "-", -x | |
| } | |
| xs := strconv.Itoa(x) | |
| if len(xs) < 2 { | |
| xs = "0" + xs | |
| } | |
| return sign + m + "e" + es + xs | |
| } | |
| switch { | |
| case decpt <= 0: | |
| return sign + "0." + strings.Repeat("0", -decpt) + digits | |
| case decpt >= len(digits): | |
| return sign + digits + strings.Repeat("0", decpt-len(digits)) + ".0" | |
| } | |
| return sign + digits[:decpt] + "." + digits[decpt:] | |
| } | |
| type jsonFail struct{} | |
| // flattenJSON is json.loads(b, object_pairs_hook=..) + parse._flatten: every leaf as (path, Python's str of it), every | |
| // duplicate key kept, in order. ok false: not JSON for Python (judged as one raw value). | |
| func flattenJSON(b string) (out []param, ok bool) { | |
| defer func() { | |
| if r := recover(); r != nil { | |
| if _, is := r.(jsonFail); !is { | |
| panic(r) | |
| } | |
| out, ok = nil, false | |
| } | |
| }() | |
| dec := json.NewDecoder(strings.NewReader(b)) | |
| dec.UseNumber() | |
| next := func() json.Token { | |
| t, err := dec.Token() | |
| if err != nil { | |
| panic(jsonFail{}) | |
| } | |
| return t | |
| } | |
| var value func(t json.Token, prefix string, depth int) | |
| value = func(t json.Token, prefix string, depth int) { | |
| if depth > parseMaxDepth { | |
| panic(jsonFail{}) | |
| } | |
| switch x := t.(type) { | |
| case json.Delim: | |
| switch x { | |
| case '{': | |
| for { | |
| kt := next() | |
| if d, is := kt.(json.Delim); is && d == '}' { | |
| return | |
| } | |
| k, is := kt.(string) | |
| if !is { | |
| panic(jsonFail{}) | |
| } | |
| p := k | |
| if prefix != "" { | |
| p = prefix + "." + k | |
| } | |
| value(next(), p, depth+1) | |
| } | |
| case '[': | |
| for i := 0; ; i++ { | |
| vt := next() | |
| if d, is := vt.(json.Delim); is && d == ']' { | |
| return | |
| } | |
| value(vt, prefix+"["+strconv.Itoa(i)+"]", depth+1) | |
| } | |
| } | |
| panic(jsonFail{}) | |
| case string: | |
| out = append(out, param{"b", prefix, x}) | |
| case json.Number: | |
| s := string(x) | |
| if strings.ContainsAny(s, ".eE") { | |
| f, _ := strconv.ParseFloat(s, 64) // the decoder checked the syntax; out of range is ±Inf or 0, as float() | |
| out = append(out, param{"b", prefix, pyFloat(f)}) | |
| } else { | |
| digits := strings.TrimPrefix(s, "-") | |
| if len(digits) > 4300 { // Python's own int limit (parse._int) | |
| panic(jsonFail{}) | |
| } | |
| if strings.Trim(digits, "0") == "" { | |
| s = "0" | |
| } | |
| out = append(out, param{"b", prefix, s}) | |
| } | |
| case bool: | |
| if x { | |
| out = append(out, param{"b", prefix, "True"}) | |
| } else { | |
| out = append(out, param{"b", prefix, "False"}) | |
| } | |
| case nil: | |
| out = append(out, param{"b", prefix, ""}) | |
| default: | |
| panic(jsonFail{}) | |
| } | |
| } | |
| value(next(), "", 0) | |
| // json.loads: nothing but JSON whitespace may follow the value | |
| if strings.Trim(b[int(dec.InputOffset()):], " \t\n\r") != "" { | |
| return nil, false | |
| } | |
| return out, true | |
| } | |
| // multipart is parse._multipart. | |
| func multipart(body, ct string) ([]param, bool) { | |
| m := reBound.FindStringSubmatch(ct) | |
| if m == nil { | |
| return nil, false | |
| } | |
| bnd := m[1] | |
| if bnd == "" { | |
| bnd = m[2] | |
| } | |
| pieces := strings.Split(body, "--"+bnd) | |
| if len(pieces) < 2 { | |
| return nil, false | |
| } | |
| out := []param{} | |
| for _, piece := range pieces[1:] { | |
| if strings.HasPrefix(piece, "--") { | |
| break | |
| } | |
| if strings.HasPrefix(piece, "\r\n") { | |
| piece = piece[2:] | |
| } else if strings.HasPrefix(piece, "\n") { | |
| piece = piece[1:] | |
| } | |
| i, j := strings.Index(piece, "\r\n\r\n"), strings.Index(piece, "\n\n") | |
| if i < 0 && j < 0 { | |
| out = append(out, param{"b", "", piece}) | |
| continue | |
| } | |
| cut, skip := j, 2 | |
| if j < 0 || 0 <= i && i < j { | |
| cut, skip = i, 4 | |
| } | |
| head, content := piece[:cut], piece[cut+skip:] | |
| if strings.HasSuffix(content, "\r\n") { | |
| content = content[:len(content)-2] | |
| } else if strings.HasSuffix(content, "\n") { | |
| content = content[:len(content)-1] | |
| } | |
| disp := "" | |
| for _, ln := range strings.Split(head, "\n") { | |
| hk, hv, sep := strings.Cut(ln, ":") | |
| if sep && strings.ToLower(pyStrip(hk)) == "content-disposition" { | |
| disp = pyStrip(hv) | |
| } | |
| } | |
| name := "" | |
| if nm := reDispNm.FindStringSubmatchIndex(disp); nm != nil { | |
| if nm[2] >= 0 { | |
| name = disp[nm[2]:nm[3]] | |
| } else if nm[4] >= 0 { | |
| name = disp[nm[4]:nm[5]] | |
| } | |
| } | |
| if fn := reDispFn.FindStringSubmatchIndex(disp); fn != nil { | |
| f := "" | |
| if fn[2] >= 0 { | |
| f = disp[fn[2]:fn[3]] | |
| } else if fn[4] >= 0 { | |
| f = disp[fn[4]:fn[5]] | |
| } | |
| out = append(out, param{"b", name, f}, param{"b", name, runeCut(content, parseMaxFile)}) | |
| } else { | |
| out = append(out, param{"b", name, content}) | |
| } | |
| } | |
| return out, true | |
| } | |
| // pyRequest is the dibaguard.http.parse.Request a request becomes: lower-cased header names, every value of a | |
| // repeated header joined (cookies with "; ", others with ", "), the raw path and query. | |
| type pyRequest struct { | |
| method, path, query, body string | |
| headers map[string]string | |
| } | |
| func newPyRequest(method string, u *url.URL, hdr http.Header, body []byte) pyRequest { | |
| // the path as it came on the wire: RawPath whenever Go's own escaping would differ (EscapedPath re-escapes a raw | |
| // "{" that the client sent unescaped, which would add a decoding layer the Python edge never saw) | |
| path := u.RawPath | |
| if path == "" { | |
| path = u.EscapedPath() | |
| } | |
| r := pyRequest{method: strings.ToUpper(method), path: validUTF8([]byte(path)), query: validUTF8([]byte(u.RawQuery)), | |
| body: validUTF8(body), headers: map[string]string{}} | |
| if r.path == "" { | |
| r.path = "/" | |
| } | |
| for k, vs := range hdr { | |
| name := strings.ToLower(k) | |
| sep := ", " | |
| if name == "cookie" { | |
| sep = "; " | |
| } | |
| v := validUTF8([]byte(strings.Join(vs, sep))) | |
| if old, ok := r.headers[name]; ok { // two spellings of one name: joined like a repeat | |
| v = old + sep + v | |
| } | |
| r.headers[name] = v | |
| } | |
| return r | |
| } | |
| func (r pyRequest) params() []param { | |
| out := pairs(r.query, "q", nil) | |
| ctRaw := r.headers["content-type"] | |
| ct := strings.ToLower(ctRaw) | |
| b := pyStrip(r.body) | |
| if b != "" { | |
| switch { | |
| case strings.Contains(ct, "json") || b[0] == '[' || b[0] == '{': | |
| if js, ok := flattenJSON(b); ok { | |
| out = append(out, js...) | |
| } else { | |
| out = append(out, param{"b", "", b}) | |
| } | |
| case strings.Contains(ct, "multipart/"): | |
| if mp, ok := multipart(b, ctRaw); ok { | |
| out = append(out, mp...) | |
| } else { | |
| out = append(out, param{"b", "", b}) | |
| } | |
| case strings.Contains(b, "="): | |
| out = pairs(b, "b", out) | |
| default: | |
| out = append(out, param{"b", "", b}) | |
| } | |
| } | |
| if ck := r.headers["cookie"]; ck != "" { | |
| for _, p := range strings.Split(ck, ";") { | |
| if pyStrip(p) == "" { | |
| continue | |
| } | |
| k, v, _ := strings.Cut(p, "=") | |
| out = append(out, param{"c", pyStrip(k), v}) | |
| } | |
| } | |
| return out | |
| } | |
| // modelText is parse.Request.model_text. | |
| func (r pyRequest) modelText() (string, int) { | |
| t := &textBuilder{rawLeft: parseRestMax} | |
| p, pa, pok, d := canonText(r.path, false) | |
| t.d = d | |
| t.add(strings.ToLower(r.method)+" ", p) | |
| if pok { | |
| t.add("p:=", pa) | |
| } | |
| ps := r.params() | |
| for i, x := range ps { | |
| if i >= parseMaxFields { | |
| break | |
| } | |
| cv, va, vok, d := canonText(x.v, x.where != "c") | |
| ck, ka, kok, _ := canonText(x.k, false) | |
| if d > t.d { | |
| t.d = d | |
| } | |
| prefix := x.where + ":" + ck + "=" | |
| t.add(prefix, cv) | |
| if vok { | |
| t.add(prefix, va) | |
| } | |
| if kok { | |
| t.add(x.where+":"+ka+"=", cv) | |
| } | |
| if runeLen(x.v) > canonMaxLen { | |
| t.addRaw(prefix, runeFrom(x.v, canonMaxLen-parseChunkOverlap)) | |
| } | |
| } | |
| if len(ps) > parseMaxFields { | |
| var rest bytes.Buffer | |
| for i, x := range ps[parseMaxFields:] { | |
| if i > 0 { | |
| rest.WriteByte('&') | |
| } | |
| rest.WriteString(x.k + "=" + x.v) | |
| } | |
| t.addRaw("x:rest=", rest.String()) | |
| } | |
| header := func(h, v string) { | |
| ch, ha, hok, d := canonText(v, false) | |
| if d > t.d { | |
| t.d = d | |
| } | |
| t.add("h:"+h+"=", ch) | |
| if hok { | |
| t.add("h:"+h+"=", ha) | |
| } | |
| } | |
| for _, h := range stdHeaders { | |
| if v := r.headers[h]; v != "" { | |
| header(h, v) | |
| } | |
| } | |
| names := make([]string, 0, len(r.headers)) | |
| for h := range r.headers { | |
| names = append(names, h) | |
| } | |
| sort.Strings(names) | |
| for _, h := range names { | |
| v := r.headers[h] | |
| if v == "" || neverHeaders[h] || h == "user-agent" || h == "referer" || h == "x-forwarded-for" || h == "content-type" { | |
| continue | |
| } | |
| ch, ha, hok, _ := canonText(v, false) | |
| if hdrSuspicious.MatchString(ch) || hok && hdrSuspicious.MatchString(ha) { | |
| header(h, v) | |
| } | |
| } | |
| return t.text(), t.d | |
| } | |