objective categorical abstract machine language personal data server
Use JWT for oauth refresh with separate per-token and per-session expiry
+178
-98
+57
-91
pegasus/lib/api/oauth_/token.ml
+57
-91
pegasus/lib/api/oauth_/token.ml
···
1
1
open Oauth
2
+
module Oauth_token = Oauth.Token
2
3
3
4
let post_handler =
4
5
Xrpc.handler ~auth:DPoP (fun ctx ->
···
65
66
Errors.invalid_request "DPoP key mismatch"
66
67
| _ ->
67
68
() ) ;
68
-
let token_id =
69
-
"tok-"
70
-
^ Uuidm.to_string
71
-
(Uuidm.v4_gen
72
-
(Random.State.make_self_init ())
73
-
() )
74
-
in
75
-
let refresh_token =
76
-
"ref-"
77
-
^ Uuidm.to_string
78
-
(Uuidm.v4_gen
79
-
(Random.State.make_self_init ())
80
-
() )
81
-
in
82
-
let now_sec = int_of_float (Unix.gettimeofday ()) in
83
69
let now_ms = Util.Time.now_ms () in
84
-
let expires_in =
85
-
Constants.access_token_expiry_ms / 1000
86
-
in
87
-
let exp_sec = now_sec + expires_in in
88
-
let expires_at = exp_sec * 1000 in
70
+
let now_sec = now_ms / 1000 in
89
71
(* expand scopes before creating token *)
90
72
let%lwt expanded_scopes =
91
73
let parsed = Scopes.parse_scopes orig_req.scope in
92
74
let%lwt expanded = Scopes.expand_scopes parsed in
93
75
Lwt.return (Scopes.scopes_to_string expanded)
94
76
in
95
-
let claims =
96
-
`Assoc
97
-
[ ("jti", `String token_id)
98
-
; ("sub", `String did)
99
-
; ("iat", `Int now_sec)
100
-
; ("exp", `Int exp_sec)
101
-
; ("scope", `String expanded_scopes)
102
-
; ("aud", `String Env.host_endpoint)
103
-
; ("cnf", `Assoc [("jkt", `String proof.jkt)])
104
-
]
105
-
in
106
-
let access_token =
107
-
Jwt.sign_jwt claims ~typ:"at+jwt"
108
-
~signing_key:Env.jwt_key
77
+
let access_token, refresh_token =
78
+
Oauth_token.generate_tokens ~did
79
+
~scope:expanded_scopes ~jkt:proof.jkt
80
+
~now:now_sec ()
109
81
in
110
82
let auth_ip =
111
83
Option.value code_rec.authorized_ip ~default:ip
···
119
91
in
120
92
let%lwt () =
121
93
Queries.insert_oauth_token ctx.db
122
-
{ refresh_token
94
+
{ refresh_token= refresh_token.token
123
95
; client_id= req.client_id
124
96
; did
125
97
; dpop_jkt= proof.jkt
126
98
; scope= expanded_scopes
127
99
; created_at= now_ms
128
100
; last_refreshed_at= now_ms
129
-
; expires_at
101
+
; session_expires_at=
102
+
Oauth_token.session_expires_at_ms now_ms
103
+
; expires_at= access_token.expires_at
130
104
; last_ip= auth_ip
131
105
; last_user_agent= auth_user_agent }
132
106
in
···
137
111
; ("Cache-Control", "no-store") ]
138
112
@@ Yojson.Safe.to_string
139
113
@@ `Assoc
140
-
[ ("access_token", `String access_token)
114
+
[ ("access_token", `String access_token.token)
141
115
; ("token_type", `String "DPoP")
142
-
; ("refresh_token", `String refresh_token)
143
-
; ("expires_in", `Int expires_in)
116
+
; ("refresh_token", `String refresh_token.token)
117
+
; ("expires_in", `Int access_token.expires_in)
144
118
; ("scope", `String expanded_scopes)
145
119
; ("sub", `String did) ] ) ) ) ) )
146
120
| "refresh_token" -> (
···
148
122
| None ->
149
123
Errors.invalid_request "refresh_token required"
150
124
| Some refresh_token -> (
151
-
let%lwt token_record =
152
-
Queries.get_oauth_token_by_refresh ctx.db refresh_token
153
-
in
154
-
match token_record with
155
-
| None ->
125
+
let now_ms = Util.Time.now_ms () in
126
+
let now_sec = now_ms / 1000 in
127
+
match
128
+
Oauth_token.verify_refresh_token ~now:now_sec refresh_token
129
+
with
130
+
| Error Oauth_token.Expired ->
131
+
Errors.invalid_request "expired refresh token"
132
+
| Error (Oauth_token.Invalid _) ->
156
133
Errors.invalid_request "invalid refresh token"
157
-
| Some session ->
158
-
if session.client_id <> req.client_id then
159
-
Errors.invalid_request "client_id mismatch"
160
-
else if session.dpop_jkt <> proof.jkt then
161
-
Errors.invalid_request "DPoP key mismatch"
162
-
else
163
-
let new_token_id =
164
-
"tok-"
165
-
^ Uuidm.to_string
166
-
(Uuidm.v4_gen (Random.State.make_self_init ()) ())
167
-
in
168
-
let new_refresh =
169
-
"ref-"
170
-
^ Uuidm.to_string
171
-
(Uuidm.v4_gen (Random.State.make_self_init ()) ())
172
-
in
173
-
let now_sec = int_of_float (Unix.gettimeofday ()) in
174
-
let expires_in = Constants.access_token_expiry_ms / 1000 in
175
-
let exp_sec = now_sec + expires_in in
176
-
let new_expires_at = exp_sec * 1000 in
177
-
let claims =
178
-
`Assoc
179
-
[ ("jti", `String new_token_id)
180
-
; ("sub", `String session.did)
181
-
; ("iat", `Int now_sec)
182
-
; ("exp", `Int exp_sec)
183
-
; ("scope", `String session.scope)
184
-
; ("aud", `String Env.host_endpoint)
185
-
; ("cnf", `Assoc [("jkt", `String proof.jkt)]) ]
186
-
in
187
-
let new_access_token =
188
-
Jwt.sign_jwt claims ~typ:"at+jwt" ~signing_key:Env.jwt_key
189
-
in
190
-
let%lwt () =
191
-
Queries.update_oauth_token ctx.db
192
-
~old_refresh_token:refresh_token
193
-
~new_refresh_token:new_refresh ~expires_at:new_expires_at
194
-
in
195
-
Dream.json ~headers:[("Cache-Control", "no-store")]
196
-
@@ Yojson.Safe.to_string
197
-
@@ `Assoc
198
-
[ ("access_token", `String new_access_token)
199
-
; ("token_type", `String "DPoP")
200
-
; ("refresh_token", `String new_refresh)
201
-
; ("expires_in", `Int expires_in)
202
-
; ("scope", `String session.scope)
203
-
; ("sub", `String session.did) ] ) )
134
+
| Ok verified_refresh -> (
135
+
let%lwt token_record =
136
+
Queries.get_oauth_token_by_refresh ctx.db refresh_token
137
+
in
138
+
match token_record with
139
+
| None ->
140
+
Errors.invalid_request "invalid refresh token"
141
+
| Some session ->
142
+
if session.client_id <> req.client_id then
143
+
Errors.invalid_request "client_id mismatch"
144
+
else if session.dpop_jkt <> proof.jkt then
145
+
Errors.invalid_request "DPoP key mismatch"
146
+
else if verified_refresh.sub <> session.did then
147
+
Errors.invalid_request "refresh token subject mismatch"
148
+
else if now_ms >= session.session_expires_at then
149
+
Errors.invalid_request "session expired"
150
+
else
151
+
let access_token, new_refresh =
152
+
Oauth_token.generate_tokens ~did:session.did
153
+
~scope:session.scope ~jkt:proof.jkt ~now:now_sec ()
154
+
in
155
+
let%lwt () =
156
+
Queries.update_oauth_token ctx.db
157
+
~old_refresh_token:refresh_token
158
+
~new_refresh_token:new_refresh.token
159
+
~expires_at:access_token.expires_at
160
+
in
161
+
Dream.json ~headers:[("Cache-Control", "no-store")]
162
+
@@ Yojson.Safe.to_string
163
+
@@ `Assoc
164
+
[ ("access_token", `String access_token.token)
165
+
; ("token_type", `String "DPoP")
166
+
; ("refresh_token", `String new_refresh.token)
167
+
; ("expires_in", `Int access_token.expires_in)
168
+
; ("scope", `String session.scope)
169
+
; ("sub", `String session.did) ] ) ) )
204
170
| _ ->
205
171
Errors.invalid_request ("unsupported grant_type: " ^ req.grant_type) )
+22
pegasus/lib/migrations/data_store/008_oauth_session_expiry.sql
+22
pegasus/lib/migrations/data_store/008_oauth_session_expiry.sql
···
1
+
ALTER TABLE oauth_tokens
2
+
ADD COLUMN session_expires_at INTEGER NOT NULL DEFAULT 0;
3
+
4
+
UPDATE oauth_tokens
5
+
SET session_expires_at =
6
+
CASE
7
+
WHEN typeof(created_at) = 'integer' THEN created_at + 31536000000
8
+
WHEN typeof(created_at) = 'text' THEN (unixepoch(created_at) * 1000) + 31536000000
9
+
ELSE (unixepoch() * 1000) + 31536000000
10
+
END
11
+
WHERE session_expires_at = 0;
12
+
13
+
CREATE INDEX IF NOT EXISTS oauth_tokens_session_expires_idx
14
+
ON oauth_tokens(session_expires_at);
15
+
16
+
DROP TRIGGER IF EXISTS cleanup_expired_oauth_tokens;
17
+
18
+
CREATE TRIGGER IF NOT EXISTS cleanup_expired_oauth_tokens
19
+
AFTER INSERT ON oauth_tokens
20
+
BEGIN
21
+
DELETE FROM oauth_tokens WHERE session_expires_at < unixepoch() * 1000;
22
+
END;
+4
pegasus/lib/oauth/constants.ml
+4
pegasus/lib/oauth/constants.ml
+9
-7
pegasus/lib/oauth/queries.ml
+9
-7
pegasus/lib/oauth/queries.ml
···
87
87
@@ [%rapper
88
88
execute
89
89
{sql|
90
-
INSERT INTO oauth_tokens (refresh_token, client_id, did, dpop_jkt, scope, created_at, expires_at, last_refreshed_at, last_ip, last_user_agent)
91
-
VALUES (%string{refresh_token}, %string{client_id}, %string{did}, %string{dpop_jkt}, %string{scope}, %int{created_at}, %int{expires_at}, %int{last_refreshed_at}, %string{last_ip}, %string?{last_user_agent})
90
+
INSERT INTO oauth_tokens (refresh_token, client_id, did, dpop_jkt, scope, created_at, last_refreshed_at, session_expires_at, expires_at, last_ip, last_user_agent)
91
+
VALUES (%string{refresh_token}, %string{client_id}, %string{did}, %string{dpop_jkt}, %string{scope}, %int{created_at}, %int{last_refreshed_at}, %int{session_expires_at}, %int{expires_at}, %string{last_ip}, %string?{last_user_agent})
92
92
|sql}
93
93
record_in]
94
94
token
···
99
99
get_opt
100
100
{sql|
101
101
SELECT @string{refresh_token}, @string{client_id}, @string{did},
102
-
@string{dpop_jkt}, @string{scope}, @int{created_at}, @int{expires_at},
103
-
@int{last_refreshed_at}, @string{last_ip}, @string?{last_user_agent}
102
+
@string{dpop_jkt}, @string{scope}, @int{created_at},
103
+
@int{last_refreshed_at}, @int{session_expires_at}, @int{expires_at},
104
+
@string{last_ip}, @string?{last_user_agent}
104
105
FROM oauth_tokens
105
106
WHERE refresh_token = %string{refresh_token}
106
107
|sql}
···
136
137
get_many
137
138
{sql|
138
139
SELECT @string{refresh_token}, @string{client_id}, @string{did},
139
-
@string{dpop_jkt}, @string{scope}, @int{created_at}, @int{expires_at},
140
-
@int{last_refreshed_at}, @string{last_ip}, @string?{last_user_agent}
140
+
@string{dpop_jkt}, @string{scope}, @int{created_at},
141
+
@int{last_refreshed_at}, @int{session_expires_at}, @int{expires_at},
142
+
@string{last_ip}, @string?{last_user_agent}
141
143
FROM oauth_tokens
142
144
WHERE did = %string{did}
143
-
ORDER BY expires_at ASC
145
+
ORDER BY session_expires_at ASC
144
146
|sql}
145
147
record_out]
146
148
~did
+83
pegasus/lib/oauth/token.ml
+83
pegasus/lib/oauth/token.ml
···
1
+
type access_token =
2
+
{token: string; token_id: string; expires_in: int; expires_at: int}
3
+
4
+
type refresh_token = {token: string; token_id: string; expires_at: int}
5
+
6
+
type refresh_token_error = Expired | Invalid of string
7
+
8
+
let refresh_token_scope = "com.atproto.oauth.refresh"
9
+
10
+
let next_token_id prefix =
11
+
prefix ^ Uuidm.to_string (Uuidm.v4_gen (Random.State.make_self_init ()) ())
12
+
13
+
let access_expires_in_s = Constants.access_token_expiry_ms / 1000
14
+
15
+
let refresh_expires_in_s = Constants.refresh_token_expiry_ms / 1000
16
+
17
+
let session_expires_at_ms created_at_ms =
18
+
created_at_ms + Constants.session_expiry_ms
19
+
20
+
let generate_access ?(now = Util.Time.now_s ()) ?token_id ~did ~scope ~jkt () =
21
+
let token_id = Option.value token_id ~default:(next_token_id "tok-") in
22
+
let exp = now + access_expires_in_s in
23
+
let claims =
24
+
`Assoc
25
+
[ ("jti", `String token_id)
26
+
; ("sub", `String did)
27
+
; ("iat", `Int now)
28
+
; ("exp", `Int exp)
29
+
; ("scope", `String scope)
30
+
; ("aud", `String Env.host_endpoint)
31
+
; ("cnf", `Assoc [("jkt", `String jkt)]) ]
32
+
in
33
+
{ token= Jwt.sign_jwt claims ~typ:"at+jwt" ~signing_key:Env.jwt_key
34
+
; token_id
35
+
; expires_in= access_expires_in_s
36
+
; expires_at= exp * 1000 }
37
+
38
+
let generate_refresh ?(now = Util.Time.now_s ()) ?token_id ~did () =
39
+
let token_id = Option.value token_id ~default:(next_token_id "ref-") in
40
+
let exp = now + refresh_expires_in_s in
41
+
let claims =
42
+
Jwt.symmetric_jwt_to_yojson
43
+
{ scope= refresh_token_scope
44
+
; aud= Env.host_endpoint
45
+
; sub= did
46
+
; iat= now
47
+
; exp
48
+
; jti= token_id }
49
+
in
50
+
{ token= Jwt.sign_jwt claims ~typ:"refresh+jwt" ~signing_key:Env.jwt_key
51
+
; token_id
52
+
; expires_at= exp * 1000 }
53
+
54
+
let generate_tokens ?(now = Util.Time.now_s ()) ?token_id ~did ~scope ~jkt () =
55
+
( generate_access ~now ?token_id ~did ~scope ~jkt ()
56
+
, generate_refresh ~now ?token_id ~did () )
57
+
58
+
let verify_refresh_token ?(now = Util.Time.now_s ()) token =
59
+
match Jwt.verify_jwt token ~pubkey:Env.jwt_pubkey with
60
+
| Error e ->
61
+
Error (Invalid e)
62
+
| Ok (header, payload) -> (
63
+
let open Yojson.Safe.Util in
64
+
match header |> member "typ" |> to_string_option with
65
+
| Some "refresh+jwt" -> (
66
+
match Jwt.symmetric_jwt_of_yojson payload with
67
+
| Error e ->
68
+
Error (Invalid e)
69
+
| Ok claims ->
70
+
if claims.scope <> refresh_token_scope then
71
+
Error (Invalid "invalid scope")
72
+
else if claims.aud <> Env.host_endpoint then
73
+
Error (Invalid "invalid audience")
74
+
else if claims.sub = "" then Error (Invalid "missing sub")
75
+
else if claims.jti = "" then Error (Invalid "missing jti")
76
+
else if now < claims.iat then
77
+
Error (Invalid "token issued in the future")
78
+
else if now >= claims.exp then Error Expired
79
+
else Ok claims )
80
+
| Some _ ->
81
+
Error (Invalid "invalid token type")
82
+
| None ->
83
+
Error (Invalid "missing token type") )
+1
pegasus/lib/oauth/types.ml
+1
pegasus/lib/oauth/types.ml
+2
pegasus/lib/util/time.ml
+2
pegasus/lib/util/time.ml