Skip to content

Commit

Permalink
feat: add webhook authentication in the API
Browse files Browse the repository at this point in the history
Authentication is off by default. Running with --token-webhook-url
option will enable it. If enabled, unauthenticated requests are not
allowed.

Signed-off-by: Donnie Adams <[email protected]>
  • Loading branch information
thedadams committed Sep 25, 2024
1 parent 6dd088e commit 1262ff9
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 11 deletions.
28 changes: 21 additions & 7 deletions pkg/api/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,21 +9,24 @@ import (
"github.com/gptscript-ai/otto/pkg/jwt"
"github.com/gptscript-ai/otto/pkg/storage"
apierrors "k8s.io/apimachinery/pkg/api/errors"
"k8s.io/apiserver/pkg/authentication/authenticator"
user2 "k8s.io/apiserver/pkg/authentication/user"
"k8s.io/apiserver/pkg/endpoints/request"
)

type Server struct {
client storage.Client
gptClient *gptscript.GPTScript
tokenService *jwt.TokenService
client storage.Client
gptClient *gptscript.GPTScript
tokenService *jwt.TokenService
authenticator authenticator.Request
}

func NewServer(client storage.Client, gptClient *gptscript.GPTScript, tokenService *jwt.TokenService) *Server {
func NewServer(client storage.Client, gptClient *gptscript.GPTScript, tokenService *jwt.TokenService, authn authenticator.Request) *Server {
return &Server{
client: client,
gptClient: gptClient,
tokenService: tokenService,
client: client,
gptClient: gptClient,
tokenService: tokenService,
authenticator: authn,
}
}

Expand All @@ -44,6 +47,17 @@ func (s *Server) Wrap(f HandlerFunc) http.Handler {
"otto:agentID": {tokenContext.AgentID},
},
}
} else if s.authenticator != nil {
resp, ok, err := s.authenticator.AuthenticateRequest(req)
if err != nil {
http.Error(rw, err.Error(), http.StatusUnauthorized)
return
} else if !ok {
http.Error(rw, "Unauthorized", http.StatusUnauthorized)
return
}

user = resp.User
} else {
user = &user2.DefaultInfo{}
}
Expand Down
39 changes: 39 additions & 0 deletions pkg/cookie/cookie.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
package cookie

import (
"net/http"

"k8s.io/apiserver/pkg/authentication/authenticator"
)

type Auth struct {
next authenticator.Request
}

func New(next authenticator.Request) authenticator.Request {
return &Auth{
next: next,
}
}

func (c *Auth) AuthenticateRequest(req *http.Request) (*authenticator.Response, bool, error) {
if c.next == nil {
return nil, false, nil
}

token, ok := GetCookieToken(req)
if ok && token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}

return c.next.AuthenticateRequest(req)
}

func GetCookieToken(req *http.Request) (string, bool) {
c, err := req.Cookie("A_SESS")
if err != nil {
return "", false
}

return c.Value, true
}
21 changes: 17 additions & 4 deletions pkg/services/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,16 +13,19 @@ import (
"github.com/gptscript-ai/gptscript/pkg/sdkserver"
"github.com/gptscript-ai/otto/pkg/aihelper"
"github.com/gptscript-ai/otto/pkg/api"
"github.com/gptscript-ai/otto/pkg/cookie"
"github.com/gptscript-ai/otto/pkg/events"
"github.com/gptscript-ai/otto/pkg/invoke"
"github.com/gptscript-ai/otto/pkg/jwt"
"github.com/gptscript-ai/otto/pkg/storage"
"github.com/gptscript-ai/otto/pkg/storage/scheme"
"github.com/gptscript-ai/otto/pkg/storage/services"
"github.com/gptscript-ai/otto/pkg/system"
"github.com/gptscript-ai/otto/pkg/webhook"
wclient "github.com/thedadams/workspace-provider/pkg/client"
coordinationv1 "k8s.io/api/coordination/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/apiserver/pkg/authentication/authenticator"
)

const (
Expand All @@ -31,9 +34,10 @@ const (
)

type Config struct {
HTTPListenPort int `usage:"HTTP port to listen on" default:"8080" name:"http-listen-port"`
DevMode bool `usage:"Enable development mode" default:"false" name:"dev-mode" env:"OTTO_DEV_MODE"`
AllowedOrigin string `usage:"Allowed origin for CORS"`
HTTPListenPort int `usage:"HTTP port to listen on" default:"8080" name:"http-listen-port"`
DevMode bool `usage:"Enable development mode" default:"false" name:"dev-mode" env:"OTTO_DEV_MODE"`
AllowedOrigin string `usage:"Allowed origin for CORS"`
TokenWebhookURL string `usage:"The url for the token webhook"`
services.Config
}

Expand Down Expand Up @@ -103,6 +107,15 @@ func New(ctx context.Context, config Config) (*Services, error) {
return nil, err
}

var wh authenticator.Request
if config.TokenWebhookURL != "" {
wh, err = webhook.New(scheme.Scheme, config.TokenWebhookURL)
if err != nil {
return nil, err
}
wh = cookie.New(wh)
}

var (
tokenServer = &jwt.TokenService{}
workspaceClient = wclient.New(wclient.Options{
Expand All @@ -116,7 +129,7 @@ func New(ctx context.Context, config Config) (*Services, error) {
StorageClient: storageClient,
Router: r,
GPTClient: c,
APIServer: api.NewServer(storageClient, c, tokenServer),
APIServer: api.NewServer(storageClient, c, tokenServer, wh),
TokenServer: tokenServer,
WorkspaceClient: workspaceClient,
Invoker: invoke.NewInvoker(storageClient, c, tokenServer, workspaceClient, events, config.KnowledgeTool),
Expand Down
54 changes: 54 additions & 0 deletions pkg/webhook/webhook.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package webhook

import (
"net/url"
"time"

"github.com/acorn-io/baaah/pkg/ratelimit"
"github.com/acorn-io/baaah/pkg/restconfig"
authenticationv1 "k8s.io/api/authentication/v1"
"k8s.io/apimachinery/pkg/runtime"
"k8s.io/apiserver/pkg/authentication/authenticator"
"k8s.io/apiserver/pkg/authentication/request/bearertoken"
"k8s.io/apiserver/pkg/authentication/token/cache"
"k8s.io/apiserver/pkg/server/options"
"k8s.io/apiserver/plugin/pkg/authenticator/token/webhook"
"k8s.io/client-go/rest"
)

func New(scheme *runtime.Scheme, webhookURL string) (authenticator.Request, error) {
restConfig, err := restCfg(webhookURL, scheme)
if err != nil {
return nil, err
}

wh, err := webhook.New(restConfig, authenticationv1.SchemeGroupVersion.Version, nil, *options.DefaultAuthWebhookRetryBackoff())
if err != nil {
return nil, err
}

tokenCache := cache.New(wh, false, 10*time.Second, 10*time.Second)
return bearertoken.New(tokenCache), nil
}

func restCfg(serverURL string, scheme *runtime.Scheme) (*rest.Config, error) {
u, err := url.Parse(serverURL)
if err != nil {
return nil, err
}
insecure := false
if u.Scheme == "https" && u.Host == "localhost" {
insecure = true
}

cfg := &rest.Config{
Host: serverURL,
TLSClientConfig: rest.TLSClientConfig{
Insecure: insecure,
},
RateLimiter: ratelimit.None,
}

restconfig.SetScheme(cfg, scheme)
return cfg, nil
}

0 comments on commit 1262ff9

Please sign in to comment.