hijack.go

package engine

// TODO: redundant from katana - the whole headless package should be replace with katana

import (
	"encoding/base64"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/proto"
)

// NewHijack create hijack from page.
func NewHijack(page *rod.Page) *Hijack {
	return &Hijack{
		page:    page,
		disable: &proto.FetchDisable{},
	}
}

// HijackHandler type
type HijackHandler = func(e *proto.FetchRequestPaused) error

// Hijack is a hijack handler
type Hijack struct {
	page    *rod.Page
	enable  *proto.FetchEnable
	disable *proto.FetchDisable
	cancel  func()
}

// SetPattern set pattern directly
func (h *Hijack) SetPattern(pattern *proto.FetchRequestPattern) {
	h.enable = &proto.FetchEnable{
		Patterns: []*proto.FetchRequestPattern{pattern},
	}
}

// Start hijack.
func (h *Hijack) Start(handler HijackHandler) func() error {
	if h.enable == nil {
		panic("hijack pattern not set")
	}

	p, cancel := h.page.WithCancel()
	h.cancel = cancel

	err := h.enable.Call(p)
	if err != nil {
		return func() error { return err }
	}

	wait := p.EachEvent(func(e *proto.FetchRequestPaused) {
		if handler != nil {
			err = handler(e)
		}
	})

	return func() error {
		wait()
		return err
	}
}

// Stop
func (h *Hijack) Stop() error {
	if h.cancel != nil {
		h.cancel()
	}
	return h.disable.Call(h.page)
}

// FetchGetResponseBody get request body.
func FetchGetResponseBody(page *rod.Page, e *proto.FetchRequestPaused) ([]byte, error) {
	m := proto.FetchGetResponseBody{
		RequestID: e.RequestID,
	}
	r, err := m.Call(page)
	if err != nil {
		return nil, err
	}

	if !r.Base64Encoded {
		return []byte(r.Body), nil
	}

	bs, err := base64.StdEncoding.DecodeString(r.Body)
	if err != nil {
		return nil, err
	}
	return bs, nil
}

// FetchContinueRequest continue request
func FetchContinueRequest(page *rod.Page, e *proto.FetchRequestPaused) error {
	m := proto.FetchContinueRequest{
		RequestID: e.RequestID,
	}
	return m.Call(page)
}