1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  //
  3  // Licensed under the Apache License, Version 2.0 (the "License");
  4  // you may not use this file except in compliance with the License.
  5  // You may obtain a copy of the License at
  6  //
  7  //     http://www.apache.org/licenses/LICENSE-2.0
  8  //
  9  // Unless required by applicable law or agreed to in writing, software
 10  // distributed under the License is distributed on an "AS IS" BASIS,
 11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  // See the License for the specific language governing permissions and
 13  // limitations under the License.
 14  //
 15  // Requirement: Any integration or derivative work must explicitly attribute
 16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  // documentation or user interface, as detailed in the NOTICE file.
 18  
 19  // Package httpx http相关
 20  package httpx
 21  
 22  import (
 23  	"bufio"
 24  	"bytes"
 25  	"fmt"
 26  	"io"
 27  	"io/ioutil"
 28  	"net/http"
 29  	"net/http/httputil"
 30  	"net/url"
 31  	"strings"
 32  
 33  	"github.com/projectdiscovery/rawhttp/client"
 34  
 35  	"github.com/projectdiscovery/retryablehttp-go"
 36  )
 37  
 38  const (
 39  	headerParts  = 2
 40  	requestParts = 3
 41  )
 42  const (
 43  	// HTTP defines the plain http scheme
 44  	HTTP = "http"
 45  	// HTTPS defines the secure http scheme
 46  	HTTPS = "https"
 47  	// HTTPorHTTPS defines the both http and https scheme
 48  	HTTPorHTTPS = "http|https"
 49  )
 50  
 51  // DumpRequest to string
 52  func DumpRequest(req *retryablehttp.Request) (string, error) {
 53  	dump, err := httputil.DumpRequestOut(req.Request, true)
 54  
 55  	return string(dump), err
 56  }
 57  
 58  // DumpResponse to string
 59  func DumpResponse(resp *http.Response) (string, error) {
 60  	// httputil.DumpResponse does not work with websockets
 61  	if resp.StatusCode == http.StatusContinue {
 62  		raw := resp.Status + "\n"
 63  		for h, v := range resp.Header {
 64  			raw += fmt.Sprintf("%s: %s\n", h, v)
 65  		}
 66  		return raw, nil
 67  	}
 68  
 69  	raw, err := httputil.DumpResponse(resp, true)
 70  	return string(raw), err
 71  }
 72  
 73  // DumpRequestRaw 导出请求包原文
 74  func DumpRequestRaw(req *retryablehttp.Request) string {
 75  	b := strings.Builder{}
 76  
 77  	b.WriteString(fmt.Sprintf("%s %s"+client.NewLine, req.Method, req.URL))
 78  
 79  	for k, v := range req.Header {
 80  		value := strings.Join(v, " ")
 81  		if value != "" {
 82  			b.WriteString(fmt.Sprintf("%s:%s"+client.NewLine, k, value))
 83  		} else {
 84  			b.WriteString(fmt.Sprintf("%s"+client.NewLine, k))
 85  		}
 86  	}
 87  
 88  	b.WriteString(client.NewLine)
 89  	body, err := req.BodyBytes()
 90  	if err != nil {
 91  		body = nil
 92  	}
 93  	if body != nil {
 94  		var buf bytes.Buffer
 95  		tee := io.TeeReader(req.Body, &buf)
 96  		body, err := ioutil.ReadAll(tee)
 97  		if err != nil {
 98  			return b.String()
 99  		}
100  		b.Write(body)
101  
102  	}
103  	return b.String()
104  }
105  
106  // ParseRequest from raw string
107  func ParseRequest(req string, unsafe bool) (method, path string, headers map[string]string, body string, err error) {
108  	headers = make(map[string]string)
109  	reader := bufio.NewReader(strings.NewReader(req))
110  	s, err := reader.ReadString('\n')
111  	if err != nil {
112  		err = fmt.Errorf("could not read request: %s", err)
113  		return
114  	}
115  	parts := strings.Split(s, " ")
116  	if len(parts) < requestParts {
117  		err = fmt.Errorf("malformed request supplied")
118  		return
119  	}
120  	method = parts[0]
121  
122  	for {
123  		line, readErr := reader.ReadString('\n')
124  		line = strings.TrimSpace(line)
125  
126  		if readErr != nil || line == "" {
127  			break
128  		}
129  
130  		// Unsafe skips all checks
131  		p := strings.SplitN(line, ":", headerParts)
132  		key := p[0]
133  		value := ""
134  		if len(p) == headerParts {
135  			value = p[1]
136  		}
137  
138  		if !unsafe {
139  			if len(p) != headerParts {
140  				continue
141  			}
142  
143  			if strings.EqualFold(key, "content-length") {
144  				continue
145  			}
146  
147  			key = strings.TrimSpace(key)
148  			value = strings.TrimSpace(value)
149  		}
150  
151  		headers[key] = value
152  	}
153  
154  	// Handle case with the full http url in path. In that case,
155  	// ignore any host header that we encounter and use the path as request URL
156  	if strings.HasPrefix(parts[1], "http") {
157  		var parsed *url.URL
158  		parsed, err = url.Parse(parts[1])
159  		if err != nil {
160  			err = fmt.Errorf("could not parse request URL: %s", err)
161  			return
162  		}
163  		path = parts[1]
164  		headers["Host"] = parsed.Host
165  	} else {
166  		path = parts[1]
167  	}
168  
169  	// Set the request body
170  	b, err := ioutil.ReadAll(reader)
171  	if err != nil {
172  		err = fmt.Errorf("could not read request body: %s", err)
173  		return
174  	}
175  	body = string(b)
176  
177  	return method, path, headers, body, nil
178  }