/
chrome.go
117 lines (91 loc) · 2.33 KB
/
chrome.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package chrome
import (
"context"
"sync"
"github.com/chromedp/cdproto/page"
"github.com/chromedp/chromedp"
"github.com/seinshah/cvci/internal/pkg/types"
)
type options struct {
pageSize types.PageSize
pageMargin types.PageMargin
}
type Headless struct {
config options
}
type Option func(*options)
func WithPageSize(size types.PageSize) Option {
return func(o *options) {
o.pageSize = size
}
}
func WithPageMargin(margin types.PageMargin) Option {
return func(o *options) {
o.pageMargin = margin
}
}
func NewHeadless(opts ...Option) *Headless {
instanceOpts := options{}
for _, opt := range opts {
opt(&instanceOpts)
}
return &Headless{
config: instanceOpts,
}
}
func (h *Headless) Generate(ctx context.Context, content []byte) ([]byte, error) {
newCtx, cancel := chromedp.NewContext(ctx)
defer cancel()
var result []byte
printTask := chromedp.Tasks{
chromedp.Navigate("about:blank"),
chromedp.ActionFunc(h.getLoadContentAction(string(content))),
chromedp.ActionFunc(h.getPrintToPDFAction(&result)),
}
if err := chromedp.Run(newCtx, printTask); err != nil {
return nil, err
}
return result, nil
}
func (h *Headless) getLoadContentAction(html string) func(ctx context.Context) error {
return func(ctx context.Context) error {
loadCtx, loadCancel := context.WithCancel(ctx)
defer loadCancel()
var wg sync.WaitGroup
wg.Add(1)
chromedp.ListenTarget(loadCtx, func(ev interface{}) {
if _, ok := ev.(*page.EventLoadEventFired); ok {
loadCancel()
wg.Done()
}
})
frameTree, err := page.GetFrameTree().Do(ctx)
if err != nil {
return err
}
if err = page.SetDocumentContent(frameTree.Frame.ID, html).Do(ctx); err != nil {
return err
}
wg.Wait()
return nil
}
}
func (h *Headless) getPrintToPDFAction(output *[]byte) func(ctx context.Context) error {
return func(ctx context.Context) error {
var err error
if *output, _, err = page.PrintToPDF().
WithDisplayHeaderFooter(false).
WithPrintBackground(true).
WithScale(1).
WithPaperWidth(h.config.pageSize.GetWidthInch()).
WithPaperHeight(h.config.pageSize.GetHeightInch()).
WithMarginTop(h.config.pageMargin.Top).
WithMarginRight(h.config.pageMargin.Right).
WithMarginBottom(h.config.pageMargin.Bottom).
WithMarginLeft(h.config.pageMargin.Left).
Do(ctx); err != nil {
return err
}
return nil
}
}