-
-
Notifications
You must be signed in to change notification settings - Fork 332
/
states.go
132 lines (112 loc) · 3.94 KB
/
states.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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
package rod
import (
"encoding/json"
"github.com/go-rod/rod/lib/proto"
"github.com/go-rod/rod/lib/utils"
)
type stateKey struct {
browserContextID proto.BrowserBrowserContextID
sessionID proto.TargetSessionID
methodName string
}
func (b *Browser) key(sessionID proto.TargetSessionID, methodName string) stateKey {
return stateKey{
browserContextID: b.BrowserContextID,
sessionID: sessionID,
methodName: methodName,
}
}
func (b *Browser) set(sessionID proto.TargetSessionID, methodName string, params json.RawMessage) {
b.states.Store(b.key(sessionID, methodName), params)
key := ""
switch methodName {
case (proto.TargetSetDiscoverTargets{}).MethodName(): // only Target domain is special
method := &proto.TargetSetDiscoverTargets{}
utils.E(json.Unmarshal(params, method))
if !method.Discover {
key = (proto.TargetSetDiscoverTargets{}).MethodName()
}
case (proto.EmulationClearDeviceMetricsOverride{}).MethodName():
key = (proto.EmulationSetDeviceMetricsOverride{}).MethodName()
case (proto.EmulationClearGeolocationOverride{}).MethodName():
key = (proto.EmulationSetGeolocationOverride{}).MethodName()
default:
domain, name := proto.ParseMethodName(methodName)
if name == "disable" {
key = domain + ".enable"
}
}
if key != "" {
b.states.Delete(b.key(sessionID, key))
}
}
// LoadState into the method, seesionID can be empty.
func (b *Browser) LoadState(sessionID proto.TargetSessionID, method proto.Payload) (has bool) {
data, has := b.states.Load(b.key(sessionID, method.MethodName()))
if has {
utils.E(json.Unmarshal(data.(json.RawMessage), method))
}
return
}
// EnableDomain and returns a recover function to restore previous state
func (b *Browser) EnableDomain(sessionID proto.TargetSessionID, method proto.Payload) (recover func()) {
_, enabled := b.states.Load(b.key(sessionID, method.MethodName()))
if !enabled {
payload, err := proto.Normalize(method)
utils.E(err)
_, _ = b.Call(b.ctx, string(sessionID), method.MethodName(), payload)
}
return func() {
if !enabled {
if method.MethodName() == (proto.TargetSetDiscoverTargets{}).MethodName() { // only Target domain is special
_ = proto.TargetSetDiscoverTargets{Discover: false}.Call(b)
return
}
domain, _ := proto.ParseMethodName(method.MethodName())
_, _ = b.Call(b.ctx, string(sessionID), domain+".disable", nil)
}
}
}
// DisableDomain and returns a recover function to restore previous state
func (b *Browser) DisableDomain(sessionID proto.TargetSessionID, method proto.Payload) (recover func()) {
_, enabled := b.states.Load(b.key(sessionID, method.MethodName()))
domain, _ := proto.ParseMethodName(method.MethodName())
if enabled {
if method.MethodName() == (proto.TargetSetDiscoverTargets{}).MethodName() { // only Target domain is special
_ = proto.TargetSetDiscoverTargets{Discover: false}.Call(b)
} else {
_, _ = b.Call(b.ctx, string(sessionID), domain+".disable", nil)
}
}
return func() {
if enabled {
payload, err := proto.Normalize(method)
utils.E(err)
_, _ = b.Call(b.ctx, string(sessionID), method.MethodName(), payload)
}
}
}
func (b *Browser) storePage(page *Page) {
b.states.Store(page.TargetID, page)
}
func (b *Browser) loadPage(id proto.TargetTargetID) *Page {
if cache, ok := b.states.Load(id); ok {
return cache.(*Page)
}
return nil
}
// LoadState into the method.
func (p *Page) LoadState(method proto.Payload) (has bool) {
return p.browser.LoadState(p.SessionID, method)
}
// EnableDomain and returns a recover function to restore previous state
func (p *Page) EnableDomain(method proto.Payload) (recover func()) {
return p.browser.Context(p.ctx).EnableDomain(p.SessionID, method)
}
// DisableDomain and returns a recover function to restore previous state
func (p *Page) DisableDomain(method proto.Payload) (recover func()) {
return p.browser.Context(p.ctx).DisableDomain(p.SessionID, method)
}
func (p *Page) cleanupStates() {
p.browser.states.Delete(p.TargetID)
}