@@ -4,103 +4,126 @@ import { clearAllIndexedDB } from "./clearAllIndexedDB";
4
4
import { restoreIndexedDB } from "./restoreIndexedDB" ;
5
5
import { restoreStorage } from "./restoreStorage" ;
6
6
7
+ import _ from "lodash" ;
7
8
import type { Browser } from "../../types" ;
8
- import { defaultOptions , SaveStateOptions } from "../saveState" ;
9
+ import { DEVTOOLS_PROTOCOL , WEBDRIVER_PROTOCOL } from "../../../constants/config" ;
10
+ import { defaultOptions , getWebdriverFrames , SaveStateData , SaveStateOptions } from "../saveState" ;
9
11
import { Protocol } from "devtools-protocol" ;
10
- import { DumpIndexDB } from "../saveState/dumpIndexedDB" ;
11
12
12
13
type RestoreStateOptions = SaveStateOptions ;
13
14
14
- const normalizeCookies = ( cookies : Array < Protocol . Network . CookieParam > ) : Array < Protocol . Network . CookieParam > =>
15
- cookies
16
- . filter ( c => c . name && c . value && c . domain )
17
- . map ( c => {
18
- const cookie : Partial < Protocol . Network . CookieParam > = {
19
- name : c . name ,
20
- value : c . value ,
21
- domain : c . domain ,
22
- path : c . path || "/" ,
23
- } ;
24
-
25
- if ( c . expires ) {
26
- cookie . expires =
27
- typeof c . expires === "string"
28
- ? Math . floor ( new Date ( c . expires ) . getTime ( ) / 1000 )
29
- : Math . floor ( c . expires as number ) ;
30
- }
31
-
32
- if ( c . secure !== undefined ) cookie . secure = c . secure ;
33
- if ( c . httpOnly !== undefined ) cookie . httpOnly = c . httpOnly ;
34
- if ( c . sameSite ) cookie . sameSite = c . sameSite ;
35
-
36
- if ( ! c . domain ?. startsWith ( "." ) ) {
37
- cookie . url = `https://${ c . domain } ` ;
38
- } else {
39
- cookie . url = `https://${ c . domain . replace ( / ^ \. / , "" ) } ` ;
40
- }
41
-
42
- return cookie as Protocol . Network . CookieParam ;
43
- } ) ;
44
-
45
- type FrameData = {
46
- localStorage : Record < string , string > ;
47
- sessionStorage : Record < string , string > ;
48
- indexDB : Record < string , DumpIndexDB > ;
49
- } ;
50
-
51
- type RestoreState = {
52
- cookies : Array < Protocol . Network . CookieParam > ;
53
- framesData : Record < string , FrameData > ;
54
- } ;
55
-
56
15
export default ( browser : Browser ) : void => {
57
16
const { publicAPI : session } = browser ;
58
17
59
18
session . addCommand ( "restoreState" , async ( _options : RestoreStateOptions ) => {
60
19
const options = { ...defaultOptions , ..._options } ;
61
20
62
- const restoreState : RestoreState = await fs . readJson ( options . path ) ;
21
+ const restoreState : SaveStateData = await fs . readJson ( options . path ) ;
63
22
64
- const puppeteer = await session . getPuppeteer ( ) ;
65
- const pages = await puppeteer . pages ( ) ;
66
- const page = pages [ 0 ] ;
67
- const frames = page . frames ( ) ;
23
+ switch ( browser . config . automationProtocol ) {
24
+ case WEBDRIVER_PROTOCOL : {
25
+ if ( restoreState . cookies && options . cookies ) {
26
+ await session . setCookies ( restoreState . cookies ) ;
27
+ }
68
28
69
- if ( restoreState . cookies && options . cookies ) {
70
- const normalized = normalizeCookies ( restoreState . cookies ) ;
29
+ if ( restoreState . framesData ) {
30
+ await session . switchToParentFrame ( ) ;
71
31
72
- await page . setCookie ( ...normalized ) ;
73
- }
32
+ const frames = await getWebdriverFrames ( session ) ;
74
33
75
- for ( const frame of frames ) {
76
- const origin = new URL ( frame . url ( ) ) . origin ;
34
+ for ( let i = - 1 ; i < frames . length ; i ++ ) {
35
+ await session . switchToParentFrame ( ) ;
77
36
78
- if ( origin === "null" || ! restoreState . framesData [ origin ] ) {
79
- continue ;
80
- }
37
+ // start with -1 for get data from main page
38
+ if ( i > - 1 ) {
39
+ await session . switchFrame ( frames [ i ] ) ;
40
+ }
81
41
82
- const frameData = restoreState . framesData [ origin ] ;
42
+ const origin = await session . execute < string , [ ] > ( ( ) => window . location . origin ) ;
83
43
84
- if ( frameData . localStorage && options . localStorage ) {
85
- await frame . evaluate (
86
- restoreStorage ,
87
- frameData . localStorage as Record < string , string > ,
88
- "localStorage" as const ,
89
- ) ;
90
- }
44
+ const frameData = restoreState . framesData [ origin ] ;
91
45
92
- if ( frameData . sessionStorage && options . sessionStorage ) {
93
- await frame . evaluate (
94
- restoreStorage ,
95
- frameData . sessionStorage as Record < string , string > ,
96
- "sessionStorage" as const ,
97
- ) ;
98
- }
46
+ if ( frameData ) {
47
+ if ( frameData . localStorage && options . localStorage ) {
48
+ await session . execute < void , [ Record < string , string > , "localStorage" ] > (
49
+ restoreStorage ,
50
+ frameData . localStorage ,
51
+ "localStorage" ,
52
+ ) ;
53
+ }
99
54
100
- if ( frameData . indexDB ) {
101
- // @todo : Doesn't work now
102
- await frame . evaluate ( clearAllIndexedDB ) ;
103
- await frame . evaluate ( restoreIndexedDB , frameData . indexDB ) ;
55
+ if ( frameData . sessionStorage && options . sessionStorage ) {
56
+ await session . execute < void , [ Record < string , string > , "sessionStorage" ] > (
57
+ restoreStorage ,
58
+ frameData . sessionStorage ,
59
+ "sessionStorage" ,
60
+ ) ;
61
+ }
62
+
63
+ if ( frameData . indexDB && options . indexDB ) {
64
+ // @todo : Doesn't work now
65
+ await session . execute ( clearAllIndexedDB ) ;
66
+ await session . execute ( restoreIndexedDB , frameData . indexDB ) ;
67
+ }
68
+ }
69
+ }
70
+
71
+ await session . switchToParentFrame ( ) ;
72
+ }
73
+
74
+ break ;
75
+ }
76
+ case DEVTOOLS_PROTOCOL : {
77
+ const puppeteer = await session . getPuppeteer ( ) ;
78
+ const pages = await puppeteer . pages ( ) ;
79
+ const page = pages [ 0 ] ;
80
+ const frames = page . frames ( ) ;
81
+
82
+ if ( restoreState . cookies && options . cookies ) {
83
+ await page . setCookie (
84
+ ...restoreState . cookies . map ( cookie => ( {
85
+ ...cookie ,
86
+ sameSite : _ . startCase ( _ . toLower ( cookie . sameSite ) ) as Protocol . Network . CookieSameSite ,
87
+ } ) ) ,
88
+ ) ;
89
+ }
90
+
91
+ for ( const frame of frames ) {
92
+ const origin = new URL ( frame . url ( ) ) . origin ;
93
+
94
+ if ( origin === "null" || ! restoreState . framesData [ origin ] ) {
95
+ continue ;
96
+ }
97
+
98
+ const frameData = restoreState . framesData [ origin ] ;
99
+
100
+ if ( ! frameData ) {
101
+ continue ;
102
+ }
103
+
104
+ if ( frameData . localStorage && options . localStorage ) {
105
+ await frame . evaluate (
106
+ restoreStorage ,
107
+ frameData . localStorage as Record < string , string > ,
108
+ "localStorage" as const ,
109
+ ) ;
110
+ }
111
+
112
+ if ( frameData . sessionStorage && options . sessionStorage ) {
113
+ await frame . evaluate (
114
+ restoreStorage ,
115
+ frameData . sessionStorage as Record < string , string > ,
116
+ "sessionStorage" as const ,
117
+ ) ;
118
+ }
119
+
120
+ if ( frameData . indexDB ) {
121
+ // @todo : Doesn't work now
122
+ await frame . evaluate ( clearAllIndexedDB ) ;
123
+ await frame . evaluate ( restoreIndexedDB , frameData . indexDB ) ;
124
+ }
125
+ }
126
+ break ;
104
127
}
105
128
}
106
129
} ) ;
0 commit comments