feat: Add explicit refresh token callable. #1230 (#1555)

app.go

@@ -86,6 +86,7 @@ func (app *App) send(clientID string, session *Session, data []byte) error {
 	if session.subject != anon {
 		req.Header.Set("Wave-Access-Token", session.token.AccessToken)
 		req.Header.Set("Wave-Refresh-Token", session.token.RefreshToken)
+		req.Header.Set("Wave-Session-ID", session.id)
 	}
 
 	resp, err := app.client.Do(req)

auth.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/coreos/go-oidc"
 	"github.com/google/uuid"
+	"github.com/h2oai/wave/pkg/keychain"
 	"golang.org/x/oauth2"
 )
 
@@ -459,3 +460,43 @@ func (h *LogoutHandler) redirect(w http.ResponseWriter, r *http.Request, idToken
 
 	http.Redirect(w, r, redirectURL.String(), http.StatusFound)
 }
+
+// Handles token refreshes.
+type RefreshHandler struct {
+	auth     *Auth
+	keychain *keychain.Keychain
+}
+
+func newRefreshHandler(auth *Auth, keychain *keychain.Keychain) http.Handler {
+	return &RefreshHandler{auth, keychain}
+}
+
+func (h *RefreshHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	if !h.keychain.Guard(w, r) { // API
+		return
+	}
+
+	sessionID := r.Header.Get("Wave-Session-ID")
+	session, ok := h.auth.get(sessionID)
+
+	if !ok {
+		echo(Log{"t": "refresh_session", "error": "session unavailable"})
+		http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
+		return
+	}
+
+	token, err := h.auth.ensureValidOAuth2Token(r.Context(), session.token)
+	if err != nil {
+		// Purge session and reload clients if refresh not successful?
+		echo(Log{"t": "refresh_session", "error": err.Error()})
+		http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
+		return
+	}
+
+	session.token = token
+	h.auth.set(session)
+
+	w.WriteHeader(http.StatusOK)
+	w.Header().Set("Wave-Access-Token", token.AccessToken)
+	w.Header().Set("Wave-Refresh-Token", token.RefreshToken)
+}

py/h2o_wave/server.py

@@ -64,7 +64,7 @@ class Auth:
     Represents authentication information for a given query context. Carries valid information only if single sign on is enabled.
     """
 
-    def __init__(self, username: str, subject: str, access_token: str, refresh_token: str):
+    def __init__(self, username: str, subject: str, access_token: str, refresh_token: str, session_id: str):
         self.username = username
         """The username of the user."""
         self.subject = subject
@@ -73,6 +73,22 @@ def __init__(self, username: str, subject: str, access_token: str, refresh_token
         """The access token of the user."""
         self.refresh_token = refresh_token
         """The refresh token of the user."""
+        self._session_id = session_id
+        """Session identifier. Do not access, internal use only."""
+
+    async def ensure_fresh_token(self) -> Optional[str]:
+        """
+        Explicitly refresh OIDC tokens when needed, e.g. during long-running background jobs.
+        """
+        async with httpx.AsyncClient(auth=(_config.hub_access_key_id, _config.hub_access_key_secret), verify=False) as http:
+            res = await http.get(_config.hub_address + '_auth/refresh', headers={'Wave-Session-ID': self._session_id}) 
+
+            access_token = res.headers.get('Wave-Access-Token', None)
+            refresh_token = res.headers.get('Wave-Refresh-Token', None)
+            if access_token and refresh_token:
+                self.access_token = access_token
+                self.refresh_token = refresh_token
+            return access_token
 
 
 class Query:
@@ -290,7 +306,8 @@ async def _receive(self, req: Request):
         username = req.headers.get('Wave-Username')
         access_token = req.headers.get('Wave-Access-Token')
         refresh_token = req.headers.get('Wave-Refresh-Token')
-        auth = Auth(username, subject, access_token, refresh_token)
+        session_id = req.headers.get('Wave-Session-ID')
+        auth = Auth(username, subject, access_token, refresh_token, session_id)
         args = await req.json()
 
         return PlainTextResponse('', background=BackgroundTask(self._process, client_id, auth, args))

server.go

@@ -103,6 +103,7 @@ func Run(conf ServerConf) {
 		handle("_auth/init", newLoginHandler(auth))
 		handle("_auth/callback", newAuthHandler(auth))
 		handle("_auth/logout", newLogoutHandler(auth, broker))
+		handle("_auth/refresh", newRefreshHandler(auth, conf.Keychain))
 	}
 
 	handle("_s/", newSocketServer(broker, auth, conf.Editable, conf.BaseURL)) // XXX terminate sockets when logged out

website/docs/security.md

@@ -117,7 +117,6 @@ To enable OpenID Connect, pass the following flags when starting the Wave server
 
 Once authenticated, you can access user's authentication and authorization information from your app using `q.auth` (see the [Auth](api/server#auth) class for details):
 
-
 ```py
 from h2o_wave import Q, main, app
 
@@ -127,12 +126,16 @@ async def serve(q: Q):
     print(q.auth.access_token)
 ```
 
-:::caution
-Note that access token is not refreshed automatically and it's not suited for long running jobs. The lifespan of a token
-depends on a provider settings but usually it's short. Access token is refreshed each time user performs an action i.e.
-the query handler `serve()` is called.
-:::
+Note that access token is not refreshed automatically and it's not suited for long running jobs. The lifespan of a token depends on a provider settings but usually it's short. Access token is refreshed each time user performs an action i.e. the query handler `serve()` is called. However, if your UI is blocked (no user interacitons that could automatically refresh the token) and you are performing a long-running job, and still need fresh access token, you can call `ensure_fresh_token` function that refreshes and sets the token explicitly. Additionally, it also returns the access token if needed for async token providers.
+
+```py
+from h2o_wave import Q, main, app
 
+@app('/example')
+async def serve(q: Q):
+    # Refreshes the token and makes it available in q.auth.access_token.
+    new_access_token = q.auth.ensure_fresh_token()
+```
 
 ## App Server API Access Keys