Return invalid_token 401 error for expired oauth token

Author: futurGH

pegasus/lib/api/account_/permissions.ml

@@ -95,6 +95,7 @@ let post_handler =
                       Oauth.Queries.delete_oauth_tokens_by_client ctx.db ~did
                         ~client_id
                     in
+                    Oauth.Dpop.revoke_tokens_for_did did ;
                     Dream.redirect ctx.req "/account/permissions"
                 | None ->
                     Dream.redirect ctx.req "/account/permissions" )
@@ -113,6 +114,7 @@ let post_handler =
                   Oauth.Queries.delete_oauth_tokens_by_device ctx.db ~did
                     ~last_ip ~last_user_agent
                 in
+                Oauth.Dpop.revoke_tokens_for_did did ;
                 Dream.redirect ctx.req "/account/permissions"
             | _ ->
                 Dream.redirect ctx.req "/account/permissions" )

pegasus/lib/auth.ml

@@ -258,7 +258,7 @@ module Verifiers = struct
         ~url:(Dream.target req) ~dpop_header ()
     with
     | Error "use_dpop_nonce" ->
-        Lwt.return_error @@ Errors.use_dpop_nonce_auth ()
+        Lwt.return_error @@ Errors.dpop_auth "use_dpop_nonce"
     | Error e ->
         Log.debug (fun log -> log "dpop error: %s" e) ;
         Lwt.return_error @@ Errors.invalid_request ("dpop error: " ^ e)
@@ -279,7 +279,7 @@ module Verifiers = struct
           ~access_token:token ()
       with
       | Error "use_dpop_nonce" ->
-          Lwt.return_error @@ Errors.use_dpop_nonce_resource ()
+          Lwt.return_error @@ Errors.dpop_resource "use_dpop_nonce"
       | Error e ->
           Log.debug (fun log -> log "dpop error: %s" e) ;
           Lwt.return_error @@ Errors.invalid_request ("dpop error: " ^ e)
@@ -292,6 +292,7 @@ module Verifiers = struct
             let open Yojson.Safe.Util in
             try
               let did = claims |> member "sub" |> to_string in
+              let iat = claims |> member "iat" |> to_int in
               let exp = claims |> member "exp" |> to_int in
               let jkt_claim =
                 claims |> member "cnf" |> member "jkt" |> to_string
@@ -307,6 +308,8 @@ module Verifiers = struct
               else if exp < now then
                 Lwt.return_error
                 @@ Errors.invalid_request ~name:"ExpiredToken" "token expired"
+              else if Oauth.Dpop.is_token_revoked ~did ~iat then
+                Lwt.return_error @@ Errors.dpop_resource "invalid_token"
               else if not (Oauth.Scopes.has_atproto scopes) then
                 Lwt.return_error
                 @@ Errors.invalid_request ~name:"InvalidToken"

pegasus/lib/errors.ml

@@ -8,11 +8,11 @@ exception NotFoundError of (string * string)
 
 exception Redirect of string
 
-(* HTTP 400, { error: "use_dpop_nonce" } — https://datatracker.ietf.org/doc/html/rfc9449#section-8 *)
-exception UseDpopNonceAuthError
+(* HTTP 400, { error: "..." } — https://datatracker.ietf.org/doc/html/rfc9449#section-8 *)
+exception DpopAuthError of string
 
-(* HTTP 401, WWW-Authenticate — https://datatracker.ietf.org/doc/html/rfc9449#section-9 *)
-exception UseDpopNonceResourceError
+(* HTTP 401, WWW-Authenticate=DPoP error="..." — https://datatracker.ietf.org/doc/html/rfc9449#section-9 *)
+exception DpopResourceError of string
 
 let is_xrpc_error = function
   | InvalidRequestError _
@@ -34,9 +34,9 @@ let auth_required ?(name = "AuthRequired") msg = raise (AuthError (name, msg))
 
 let not_found ?(name = "NotFound") msg = raise (NotFoundError (name, msg))
 
-let use_dpop_nonce_auth () = raise UseDpopNonceAuthError
+let dpop_auth error = raise (DpopAuthError error)
 
-let use_dpop_nonce_resource () = raise UseDpopNonceResourceError
+let dpop_resource error = raise (DpopResourceError error)
 
 let printer = function
   | InvalidRequestError (error, message) ->
@@ -47,10 +47,10 @@ let printer = function
       Some (Printf.sprintf "Auth error (%s): %s" error message)
   | NotFoundError (error, message) ->
       Some (Printf.sprintf "Not found (%s): %s" error message)
-  | UseDpopNonceAuthError ->
-      Some "Use DPoP nonce"
-  | UseDpopNonceResourceError ->
-      Some "Use DPoP nonce"
+  | DpopAuthError error ->
+      Some (Printf.sprintf "DPoP auth error (%s)" error)
+  | DpopResourceError error ->
+      Some (Printf.sprintf "DPoP resource error (%s)" error)
   | _ ->
       None
 
@@ -72,18 +72,18 @@ let exn_to_response exn =
   | NotFoundError (error, message) ->
       Log.debug (fun log -> log "not found error: %s - %s" error message) ;
       format_response error message `Not_Found
-  | UseDpopNonceAuthError ->
-      Log.debug (fun log -> log "use_dpop_nonce auth error") ;
+  | DpopAuthError e ->
+      Log.debug (fun log -> log "dpop auth error") ;
       Dream.json ~status:`Bad_Request
         ~headers:[("Access-Control-Expose-Headers", "DPoP-Nonce")]
-        {|{ "error": "use_dpop_nonce" }|}
-  | UseDpopNonceResourceError ->
-      Log.debug (fun log -> log "use_dpop_nonce resource error") ;
+        (Format.sprintf {|{ "error": "%s" }|} e)
+  | DpopResourceError e ->
+      Log.debug (fun log -> log "dpop resource error") ;
       Dream.json ~status:`Unauthorized
         ~headers:
-          [ ("WWW-Authenticate", {|DPoP error="use_dpop_nonce"|})
+          [ ("WWW-Authenticate", Format.sprintf {|DPoP error="%s"|} e)
           ; ("Access-Control-Expose-Headers", "DPoP-Nonce, WWW-Authenticate") ]
-        {|{ "error": "use_dpop_nonce" }|}
+        (Format.sprintf {|{ "error": "%s" }|} e)
   | e ->
       Log.warn (fun log ->
           log "unexpected internal error: %s" (Printexc.to_string e) ) ;

pegasus/lib/oauth/dpop.ml

@@ -20,6 +20,28 @@ let cleanup_jti_cache () =
     (fun _ expires_at -> if expires_at > now then Some expires_at else None)
     jti_cache
 
+let revocation_cache : (string, int) Hashtbl.t = Hashtbl.create 1000
+
+let cleanup_revocation_cache () =
+  let now_s = int_of_float (Unix.gettimeofday ()) in
+  let max_token_age_s = Constants.access_token_expiry_ms / 1000 in
+  Hashtbl.filter_map_inplace
+    (fun _ revoked_at ->
+      if now_s - revoked_at < max_token_age_s then Some revoked_at else None )
+    revocation_cache
+
+let revoke_tokens_for_did did =
+  let now_s = int_of_float (Unix.gettimeofday ()) in
+  Hashtbl.replace revocation_cache did now_s ;
+  if Hashtbl.length revocation_cache mod 50 = 0 then cleanup_revocation_cache ()
+
+let is_token_revoked ~did ~iat =
+  match Hashtbl.find_opt revocation_cache did with
+  | Some revoked_at ->
+      iat < revoked_at
+  | None ->
+      false
+
 let compute_nonce secret counter =
   let data = Bytes.create 8 in
   Bytes.set_int64_be data 0 (Int64.of_int counter) ;

pegasus/lib/xrpc.ml

@@ -108,7 +108,7 @@ let add_dpop_nonce_if_needed res =
   in
   let () =
     let to_expose =
-      (* see comments on UseDpopNonce____Error in errors.ml *)
+      (* see comments on Dpop____Error in errors.ml *)
       if Dream.status res = `Unauthorized then "DPoP-Nonce, WWW-Authenticate"
       else if Dream.status res = `Bad_Request then "DPoP-Nonce"
       else ""
@@ -146,7 +146,7 @@ let handler ?(auth : Auth.Verifiers.t = Any)
           let%lwt res = exn_to_response e in
           Lwt.return
             ( match e with
-            | UseDpopNonceAuthError | UseDpopNonceResourceError ->
+            | DpopAuthError _ | DpopResourceError _ ->
                 add_dpop_nonce_if_needed res
             | _ ->
                 res )
@@ -155,7 +155,7 @@ let handler ?(auth : Auth.Verifiers.t = Any)
         Dream.redirect init.req r
     | Rate_limiter.Rate_limit_exceeded status ->
         rate_limit_response status
-    | (UseDpopNonceAuthError | UseDpopNonceResourceError) as e ->
+    | (DpopAuthError _ | DpopResourceError _) as e ->
         let%lwt res = exn_to_response e in
         Lwt.return (add_dpop_nonce_if_needed res)
     | e ->