github.com/gospider007/requests@v0.0.0-20240506025355-c73d46169a23/response.go (about)

     1  package requests
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"context"
     7  	"errors"
     8  	"io"
     9  	"net/url"
    10  	"strconv"
    11  	"strings"
    12  
    13  	"net/http"
    14  
    15  	"github.com/gospider007/bar"
    16  	"github.com/gospider007/bs4"
    17  	"github.com/gospider007/gson"
    18  	"github.com/gospider007/re"
    19  	"github.com/gospider007/tools"
    20  	"github.com/gospider007/websocket"
    21  )
    22  
    23  type Response struct {
    24  	rawConn   *readWriteCloser
    25  	response  *http.Response
    26  	webSocket *websocket.Conn
    27  	sse       *Sse
    28  	ctx       context.Context
    29  	cnl       context.CancelFunc
    30  	content   []byte
    31  	encoding  string
    32  	stream    bool
    33  	disDecode bool
    34  	disUnzip  bool
    35  	filePath  string
    36  	bar       bool
    37  	isNewConn bool
    38  	readBody  bool
    39  }
    40  
    41  type Sse struct {
    42  	reader *bufio.Reader
    43  	raw    io.ReadCloser
    44  }
    45  type Event struct {
    46  	Data    string //data
    47  	Event   string //event
    48  	Id      string //id
    49  	Retry   int    //retry num
    50  	Comment string //comment info
    51  }
    52  
    53  func newSse(rd io.ReadCloser) *Sse {
    54  	return &Sse{raw: rd, reader: bufio.NewReader(rd)}
    55  }
    56  
    57  // recv sse envent data
    58  func (obj *Sse) Recv() (Event, error) {
    59  	var event Event
    60  	for {
    61  		readStr, err := obj.reader.ReadString('\n')
    62  		if err != nil || readStr == "\n" {
    63  			return event, err
    64  		}
    65  		reResult := re.Search(`data:\s?(.*)`, readStr)
    66  		if reResult != nil {
    67  			event.Data += reResult.Group(1)
    68  			continue
    69  		}
    70  		reResult = re.Search(`event:\s?(.*)`, readStr)
    71  
    72  		if reResult != nil {
    73  			event.Event = reResult.Group(1)
    74  			continue
    75  		}
    76  		reResult = re.Search(`id:\s?(.*)`, readStr)
    77  		if reResult != nil {
    78  			event.Id = reResult.Group(1)
    79  			continue
    80  		}
    81  		reResult = re.Search(`retry:\s?(.*)`, readStr)
    82  		if reResult != nil {
    83  			if event.Retry, err = strconv.Atoi(reResult.Group(1)); err != nil {
    84  				return event, err
    85  			}
    86  			continue
    87  		}
    88  		reResult = re.Search(`:\s?(.*)`, readStr)
    89  		if reResult != nil {
    90  			event.Comment = reResult.Group(1)
    91  			continue
    92  		}
    93  		return event, errors.New("content parse error:" + readStr)
    94  	}
    95  }
    96  
    97  // close sse
    98  func (obj *Sse) Close() error {
    99  	return obj.raw.Close()
   100  }
   101  
   102  // return websocket client
   103  func (obj *Response) WebSocket() *websocket.Conn {
   104  	return obj.webSocket
   105  }
   106  
   107  // return sse client
   108  func (obj *Response) Sse() *Sse {
   109  	return obj.sse
   110  }
   111  
   112  // return URL redirected address
   113  func (obj *Response) Location() (*url.URL, error) {
   114  	u, err := obj.response.Location()
   115  	if err == http.ErrNoLocation {
   116  		err = nil
   117  	}
   118  	return u, err
   119  }
   120  
   121  // return response Proto
   122  func (obj *Response) Proto() string {
   123  	return obj.response.Proto
   124  }
   125  
   126  // return response cookies
   127  func (obj *Response) Cookies() Cookies {
   128  	if obj.filePath != "" {
   129  		return nil
   130  	}
   131  	return obj.response.Cookies()
   132  }
   133  
   134  // return response status code
   135  func (obj *Response) StatusCode() int {
   136  	if obj.filePath != "" {
   137  		return 200
   138  	}
   139  	return obj.response.StatusCode
   140  }
   141  
   142  // return response status
   143  func (obj *Response) Status() string {
   144  	if obj.filePath != "" {
   145  		return "200 OK"
   146  	}
   147  	return obj.response.Status
   148  }
   149  
   150  // return response url
   151  func (obj *Response) Url() *url.URL {
   152  	if obj.filePath != "" {
   153  		return nil
   154  	}
   155  	return obj.response.Request.URL
   156  }
   157  
   158  // return response headers
   159  func (obj *Response) Headers() http.Header {
   160  	if obj.filePath != "" {
   161  		return http.Header{
   162  			"Content-Type": []string{obj.ContentType()},
   163  		}
   164  	}
   165  	return obj.response.Header
   166  }
   167  
   168  // change decoding with content
   169  func (obj *Response) Decode(encoding string) {
   170  	if obj.encoding != encoding {
   171  		obj.encoding = encoding
   172  		obj.SetContent(tools.Decode(obj.Content(), encoding))
   173  	}
   174  }
   175  
   176  // return content with map[string]any
   177  func (obj *Response) Map() (data map[string]any, err error) {
   178  	_, err = gson.Decode(obj.Content(), &data)
   179  	return
   180  }
   181  
   182  // return content with json and you can parse struct
   183  func (obj *Response) Json(vals ...any) (*gson.Client, error) {
   184  	return gson.Decode(obj.Content(), vals...)
   185  }
   186  
   187  // return content with string
   188  func (obj *Response) Text() string {
   189  	return tools.BytesToString(obj.Content())
   190  }
   191  
   192  // set response content with []byte
   193  func (obj *Response) SetContent(val []byte) {
   194  	obj.content = val
   195  }
   196  
   197  // return content with []byte
   198  func (obj *Response) Content() []byte {
   199  	return obj.content
   200  }
   201  
   202  // return content with parse html
   203  func (obj *Response) Html() *bs4.Client {
   204  	return bs4.NewClient(obj.Text(), obj.Url().String())
   205  }
   206  
   207  // return content type
   208  func (obj *Response) ContentType() string {
   209  	if obj.filePath != "" {
   210  		return http.DetectContentType(obj.content)
   211  	}
   212  	contentType := obj.response.Header.Get("Content-Type")
   213  	if contentType == "" {
   214  		contentType = http.DetectContentType(obj.content)
   215  	}
   216  	return contentType
   217  }
   218  
   219  // return content encoding
   220  func (obj *Response) ContentEncoding() string {
   221  	if obj.filePath != "" {
   222  		return ""
   223  	}
   224  	return obj.response.Header.Get("Content-Encoding")
   225  }
   226  
   227  // return content length
   228  func (obj *Response) ContentLength() int64 {
   229  	if obj.filePath != "" {
   230  		return int64(len(obj.content))
   231  	}
   232  	if obj.response.ContentLength >= 0 {
   233  		return obj.response.ContentLength
   234  	}
   235  	return int64(len(obj.content))
   236  }
   237  
   238  type barBody struct {
   239  	body *bytes.Buffer
   240  	bar  *bar.Client
   241  }
   242  
   243  func (obj *barBody) Write(con []byte) (int, error) {
   244  	l, err := obj.body.Write(con)
   245  	obj.bar.Add(int64(l))
   246  	return l, err
   247  }
   248  func (obj *Response) defaultDecode() bool {
   249  	return strings.Contains(obj.ContentType(), "html")
   250  }
   251  
   252  // must stream=true
   253  func (obj *Response) Conn() *connecotr {
   254  	if obj.IsStream() {
   255  		return obj.rawConn.Conn()
   256  	}
   257  	return nil
   258  }
   259  
   260  // return true if response is stream
   261  func (obj *Response) IsStream() bool {
   262  	return obj.webSocket != nil || obj.sse != nil || obj.stream
   263  }
   264  
   265  // read body
   266  func (obj *Response) ReadBody() (err error) {
   267  	if obj.IsStream() {
   268  		return errors.New("can not read stream")
   269  	}
   270  	if obj.readBody {
   271  		return errors.New("already read body")
   272  	}
   273  	obj.readBody = true
   274  	bBody := bytes.NewBuffer(nil)
   275  	if obj.bar && obj.ContentLength() > 0 {
   276  		_, err = io.Copy(&barBody{
   277  			bar:  bar.NewClient(obj.response.ContentLength),
   278  			body: bBody,
   279  		}, obj.response.Body)
   280  	} else {
   281  		_, err = io.Copy(bBody, obj.response.Body)
   282  	}
   283  	if err != nil {
   284  		obj.ForceCloseConn()
   285  		return errors.New("response read content error: " + err.Error())
   286  	}
   287  	if !obj.disDecode && obj.defaultDecode() {
   288  		obj.content, obj.encoding, _ = tools.Charset(bBody.Bytes(), obj.ContentType())
   289  	} else {
   290  		obj.content = bBody.Bytes()
   291  	}
   292  	return
   293  }
   294  
   295  // conn is new conn
   296  func (obj *Response) IsNewConn() bool {
   297  	return obj.isNewConn
   298  }
   299  
   300  // conn proxy
   301  func (obj *Response) Proxy() string {
   302  	if obj.rawConn != nil {
   303  		return obj.rawConn.Proxy()
   304  	}
   305  	return ""
   306  }
   307  
   308  // conn is in pool ?
   309  func (obj *Response) InPool() bool {
   310  	if obj.rawConn != nil {
   311  		return obj.rawConn.InPool()
   312  	}
   313  	return false
   314  }
   315  
   316  // close body
   317  func (obj *Response) CloseBody() {
   318  	obj.close(false)
   319  }
   320  
   321  // safe close conn
   322  func (obj *Response) CloseConn() {
   323  	obj.close(true)
   324  }
   325  
   326  // close
   327  func (obj *Response) close(closeConn bool) {
   328  	if obj.webSocket != nil {
   329  		obj.webSocket.Close()
   330  	}
   331  	if obj.sse != nil {
   332  		obj.sse.Close()
   333  	}
   334  	if obj.IsStream() || !obj.readBody {
   335  		obj.ForceCloseConn()
   336  	} else if obj.rawConn != nil {
   337  		if closeConn {
   338  			obj.rawConn.CloseConn()
   339  		} else {
   340  			obj.rawConn.Close()
   341  		}
   342  	}
   343  	obj.cnl() //must later
   344  }
   345  
   346  // force close conn
   347  func (obj *Response) ForceCloseConn() {
   348  	if obj.rawConn != nil {
   349  		obj.rawConn.ForceCloseConn()
   350  	}
   351  }