Files
phishingclub/backend/vendor/github.com/go-rod/rod/utils.go
T
Ronni Skansing f5e6f53d75 update vendor
Signed-off-by: Ronni Skansing <rskansing@gmail.com>
2026-05-24 09:57:07 +02:00

272 lines
5.8 KiB
Go

package rod
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"os"
"path/filepath"
"reflect"
"regexp"
"runtime/debug"
"sync"
"time"
"github.com/go-rod/rod/lib/cdp"
"github.com/go-rod/rod/lib/proto"
"github.com/go-rod/rod/lib/utils"
)
// CDPClient is usually used to make rod side-effect free. Such as proxy all IO of rod.
type CDPClient interface {
Event() <-chan *cdp.Event
Call(ctx context.Context, sessionID, method string, params interface{}) ([]byte, error)
}
// Message represents a cdp.Event.
type Message struct {
SessionID proto.TargetSessionID
Method string
lock *sync.Mutex
data json.RawMessage
event reflect.Value
}
// Load data into e, returns true if e matches the event type.
func (msg *Message) Load(e proto.Event) bool {
if msg.Method != e.ProtoEvent() {
return false
}
eVal := reflect.ValueOf(e)
if eVal.Kind() != reflect.Ptr {
return true
}
eVal = reflect.Indirect(eVal)
msg.lock.Lock()
defer msg.lock.Unlock()
if msg.data == nil {
eVal.Set(msg.event)
return true
}
utils.E(json.Unmarshal(msg.data, e))
msg.event = eVal
msg.data = nil
return true
}
// DefaultLogger for rod.
var DefaultLogger = log.New(os.Stdout, "[rod] ", log.LstdFlags)
// DefaultSleeper generates the default sleeper for retry, it uses backoff to grow the interval.
// The growth looks like:
//
// A(0) = 100ms, A(n) = A(n-1) * random[1.9, 2.1), A(n) < 1s
//
// Why the default is not RequestAnimationFrame or DOM change events is because of if a retry never
// ends it can easily flood the program. But you can always easily config it into what you want.
var DefaultSleeper = func() utils.Sleeper {
return utils.BackoffSleeper(100*time.Millisecond, time.Second, nil)
}
// NewPagePool instance.
func NewPagePool(limit int) Pool[Page] {
return NewPool[Page](limit)
}
// NewBrowserPool instance.
func NewBrowserPool(limit int) Pool[Browser] {
return NewPool[Browser](limit)
}
// Pool is used to thread-safely limit the number of elements at the same time.
// It's a common practice to use a channel to limit concurrency, it's not special for rod.
// This helper is more like an example to use Go Channel.
// Reference: https://golang.org/doc/effective_go#channels
type Pool[T any] chan *T
// NewPool instance.
func NewPool[T any](limit int) Pool[T] {
p := make(chan *T, limit)
for i := 0; i < limit; i++ {
p <- nil
}
return p
}
// Get a elem from the pool, allow error. Use the [Pool[T].Put] to make it reusable later.
func (p Pool[T]) Get(create func() (*T, error)) (elem *T, err error) {
elem = <-p
if elem == nil {
elem, err = create()
}
return
}
// Put an elem back to the pool.
func (p Pool[T]) Put(elem *T) {
p <- elem
}
// Cleanup helper.
func (p Pool[T]) Cleanup(iteratee func(*T)) {
for i := 0; i < cap(p); i++ {
select {
case elem := <-p:
if elem != nil {
iteratee(elem)
}
default:
}
}
}
var _ io.ReadCloser = &StreamReader{}
// StreamReader for browser data stream.
type StreamReader struct {
Offset *int
c proto.Client
handle proto.IOStreamHandle
buf *bytes.Buffer
}
// NewStreamReader instance.
func NewStreamReader(c proto.Client, h proto.IOStreamHandle) *StreamReader {
return &StreamReader{
c: c,
handle: h,
buf: &bytes.Buffer{},
}
}
func (sr *StreamReader) Read(p []byte) (n int, err error) {
res, err := proto.IORead{
Handle: sr.handle,
Offset: sr.Offset,
}.Call(sr.c)
if err != nil {
return 0, err
}
if !res.EOF {
var bin []byte
if res.Base64Encoded {
bin, err = base64.StdEncoding.DecodeString(res.Data)
if err != nil {
return 0, err
}
} else {
bin = []byte(res.Data)
}
_, _ = sr.buf.Write(bin)
}
return sr.buf.Read(p)
}
// Close the stream, discard any temporary backing storage.
func (sr *StreamReader) Close() error {
return proto.IOClose{Handle: sr.handle}.Call(sr.c)
}
// Try try fn with recover, return the panic as rod.ErrTry.
func Try(fn func()) (err error) {
defer func() {
if val := recover(); val != nil {
err = &TryError{val, string(debug.Stack())}
}
}()
fn()
return err
}
func genRegMatcher(includes, excludes []string) func(string) bool {
regIncludes := make([]*regexp.Regexp, len(includes))
for i, p := range includes {
regIncludes[i] = regexp.MustCompile(p)
}
regExcludes := make([]*regexp.Regexp, len(excludes))
for i, p := range excludes {
regExcludes[i] = regexp.MustCompile(p)
}
return func(s string) bool {
for _, include := range regIncludes {
if include.MatchString(s) {
for _, exclude := range regExcludes {
if exclude.MatchString(s) {
goto end
}
}
return true
}
}
end:
return false
}
}
type saveFileType int
const (
saveFileTypeScreenshot saveFileType = iota
saveFileTypePDF
)
func saveFile(fileType saveFileType, bin []byte, toFile []string) error {
if len(toFile) == 0 {
return nil
}
if toFile[0] == "" {
stamp := fmt.Sprintf("%d", time.Now().UnixNano())
switch fileType {
case saveFileTypeScreenshot:
toFile = []string{"tmp", "screenshots", stamp + ".png"}
case saveFileTypePDF:
toFile = []string{"tmp", "pdf", stamp + ".pdf"}
}
}
return utils.OutputFile(filepath.Join(toFile...), bin)
}
func httHTML(w http.ResponseWriter, body string) {
w.Header().Add("Content-Type", "text/html; charset=utf-8")
_, _ = w.Write([]byte(body))
}
func mustToJSONForDev(value interface{}) string {
buf := new(bytes.Buffer)
enc := json.NewEncoder(buf)
enc.SetEscapeHTML(false)
utils.E(enc.Encode(value))
return buf.String()
}
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/Data_URIs
var regDataURI = regexp.MustCompile(`\Adata:(.+?)?(;base64)?,`)
func parseDataURI(uri string) (string, []byte) {
matches := regDataURI.FindStringSubmatch(uri)
l := len(matches[0])
contentType := matches[1]
bin, _ := base64.StdEncoding.DecodeString(uri[l:])
return contentType, bin
}