fix: refresh oauth token for API GET requests made from the UI

Signed-off-by: Donnie Adams <[email protected]>

pkg/api/router/router.go

@@ -146,7 +146,7 @@ func Router(services *services.Services) (http.Handler, error) {
 	services.GatewayServer.AddRoutes(services.APIServer)
 
 	// UI
-	services.APIServer.HTTPHandle("/", services.ProxyServer.Wrap(ui.Handler(services.DevUIPort)))
+	services.APIServer.HTTPHandle("/", ui.Handler(services.DevUIPort))
 
 	return services.APIServer, nil
 }

pkg/api/server/server.go

@@ -3,12 +3,14 @@ package server
 import (
 	"errors"
 	"net/http"
+	"strings"
 
 	"github.com/gptscript-ai/go-gptscript"
 	"github.com/otto8-ai/otto8/apiclient/types"
 	"github.com/otto8-ai/otto8/pkg/api"
 	"github.com/otto8-ai/otto8/pkg/api/authn"
 	"github.com/otto8-ai/otto8/pkg/api/authz"
+	"github.com/otto8-ai/otto8/pkg/proxy"
 	"github.com/otto8-ai/otto8/pkg/storage"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 )
@@ -18,16 +20,18 @@ type Server struct {
 	gptClient     *gptscript.GPTScript
 	authenticator *authn.Authenticator
 	authorizer    *authz.Authorizer
+	proxyServer   *proxy.Proxy
 
 	mux *http.ServeMux
 }
 
-func NewServer(storageClient storage.Client, gptClient *gptscript.GPTScript, authn *authn.Authenticator, authz *authz.Authorizer) *Server {
+func NewServer(storageClient storage.Client, gptClient *gptscript.GPTScript, authn *authn.Authenticator, authz *authz.Authorizer, proxyServer *proxy.Proxy) *Server {
 	return &Server{
 		storageClient: storageClient,
 		gptClient:     gptClient,
 		authenticator: authn,
 		authorizer:    authz,
+		proxyServer:   proxyServer,
 
 		mux: http.NewServeMux(),
 	}
@@ -50,14 +54,29 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 func (s *Server) wrap(f api.HandlerFunc) http.HandlerFunc {
 	return func(rw http.ResponseWriter, req *http.Request) {
+		// If this header is set, then the session was deemed to be invalid and the request has come back around through the proxy.
+		// The cookie on the request is still invalid because the new one has not been sent back to the browser.
+		// Therefore, respond with a redirect so that the browser will redirect back to the original request with the new cookie.
+		if req.Header.Get("X-Otto-Auth-Required") == "true" {
+			http.Redirect(rw, req, req.RequestURI, http.StatusFound)
+			return
+		}
+
 		user, err := s.authenticator.Authenticate(req)
 		if err != nil {
 			http.Error(rw, err.Error(), http.StatusUnauthorized)
 			return
 		}
 
-		if !s.authorizer.Authorize(req, user) {
-			http.Error(rw, "forbidden", http.StatusForbidden)
+		if !s.authorizer.Authorize(req, user) || strings.HasPrefix(req.URL.Path, "/oauth2/") {
+			// If this is not a request coming from browser or the proxy is not enabled, then return 403.
+			if s.proxyServer == nil || req.Method != http.MethodGet || !strings.Contains(strings.ToLower(req.UserAgent()), "mozilla") {
+				http.Error(rw, "forbidden", http.StatusForbidden)
+				return
+			}
+
+			req.Header.Set("X-Otto-Auth-Required", "true")
+			s.proxyServer.ServeHTTP(rw, req)
 			return
 		}
 

pkg/proxy/proxy.go

@@ -2,7 +2,6 @@ package proxy
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"net/http"
 	"strings"
@@ -78,33 +77,8 @@ func New(serverURL string, authProviderID uint, cfg Config) (*Proxy, error) {
 	}, nil
 }
 
-func (p *Proxy) Wrap(h http.Handler) http.Handler {
-	if p == nil {
-		return h
-	}
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if strings.HasPrefix(r.URL.Path, "/static/") || strings.HasPrefix(r.URL.Path, "/ui/login/complete") {
-			// No authentication required
-			h.ServeHTTP(w, r)
-			return
-		}
-
-		// If this header is set, then the session was deemed to be invalid and the request has come back around through the proxy.
-		// The cookie on the request is still invalid because the new one has not been sent back to the browser.
-		// Therefore, respond with a redirect so that the browser will redirect back to the original request with the new cookie.
-		if r.Header.Get("X-Otto-Auth-Required") != "" {
-			http.Redirect(w, r, r.URL.RawPath, http.StatusFound)
-			return
-		}
-
-		state, err := p.proxy.LoadCookiedSession(r)
-		if strings.HasPrefix(r.URL.Path, "/oauth2") || err != nil && !errors.Is(err, http.ErrNoCookie) || state != nil && state.IsExpired() {
-			r.Header.Add("X-Otto-Auth-Required", "true")
-			p.proxy.ServeHTTP(w, r)
-		} else {
-			h.ServeHTTP(w, r)
-		}
-	})
+func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	p.proxy.ServeHTTP(w, r)
 }
 
 func (p *Proxy) AuthenticateRequest(req *http.Request) (*authenticator.Response, bool, error) {

pkg/services/config.go

@@ -232,7 +232,7 @@ func New(ctx context.Context, config Config) (*Services, error) {
 		StorageClient:         storageClient,
 		Router:                r,
 		GPTClient:             c,
-		APIServer:             server.NewServer(storageClient, c, authn.NewAuthenticator(authenticators), authz.NewAuthorizer()),
+		APIServer:             server.NewServer(storageClient, c, authn.NewAuthenticator(authenticators), authz.NewAuthorizer(), proxyServer),
 		TokenServer:           tokenServer,
 		Invoker:               invoke.NewInvoker(storageClient, c, config.Hostname, config.WorkspaceProviderType, tokenServer, events),
 		AIHelper:              aihelper.New(c, config.HelperModel),