diff --git a/ee/tabby-ui/components/prompt-form.tsx b/ee/tabby-ui/components/prompt-form.tsx index 7b29e11..097068e 100644 --- a/ee/tabby-ui/components/prompt-form.tsx +++ b/ee/tabby-ui/components/prompt-form.tsx @@ -1,6 +1,6 @@ import * as React from 'react' import { UseChatHelpers } from 'ai/react' -import { debounce, has } from 'lodash-es' +import { debounce, has, isEqual } from 'lodash-es' import useSWR from 'swr' import { useEnterSubmit } from '@/lib/hooks/use-enter-submit' @@ -26,6 +26,7 @@ import { TooltipContent, TooltipTrigger } from '@/components/ui/tooltip' +import { useSession } from '@/lib/tabby/auth' export interface PromptProps extends Pick { @@ -45,7 +46,6 @@ function PromptFormRenderer( const [queryCompletionUrl, setQueryCompletionUrl] = React.useState< string | null >(null) - const latestFetchKey = React.useRef('') const inputRef = React.useRef(null) // store the input selection for replacing inputValue const prevInputSelectionEnd = React.useRef() @@ -56,11 +56,11 @@ function PromptFormRenderer( Record >({}) - useSWR(queryCompletionUrl, fetcher, { + const { data } = useSession(); + useSWR([queryCompletionUrl, data?.accessToken], fetcher, { revalidateOnFocus: false, dedupingInterval: 0, - onSuccess: (data, key) => { - if (key !== latestFetchKey.current) return + onSuccess: (data) => { setOptions(data?.hits ?? []) } }) @@ -102,7 +102,6 @@ function PromptFormRenderer( if (queryName) { const query = encodeURIComponent(`name:${queryName} AND kind:function`) const url = `/v1beta/search?q=${query}` - latestFetchKey.current = url setQueryCompletionUrl(url) } else { setOptions([]) diff --git a/ee/tabby-ui/lib/hooks/use-health.tsx b/ee/tabby-ui/lib/hooks/use-health.tsx index 459dcb6..6502993 100644 --- a/ee/tabby-ui/lib/hooks/use-health.tsx +++ b/ee/tabby-ui/lib/hooks/use-health.tsx @@ -1,9 +1,9 @@ 'use client' -import { SWRResponse } from 'swr' -import useSWRImmutable from 'swr/immutable' +import useSWR, { SWRResponse } from 'swr' import fetcher from '@/lib/tabby/fetcher' +import { useSession } from '../tabby/auth' export interface HealthInfo { device: 'metal' | 'cpu' | 'cuda' @@ -19,5 +19,6 @@ export interface HealthInfo { } export function useHealth(): SWRResponse { - return useSWRImmutable('/v1/health', fetcher) + const { data } = useSession() + return useSWR(['/v1/health', data?.accessToken], fetcher) } diff --git a/ee/tabby-ui/lib/hooks/use-patch-fetch.ts b/ee/tabby-ui/lib/hooks/use-patch-fetch.ts index 408f4bc..5bb917b 100644 --- a/ee/tabby-ui/lib/hooks/use-patch-fetch.ts +++ b/ee/tabby-ui/lib/hooks/use-patch-fetch.ts @@ -5,30 +5,43 @@ import { StreamingTextResponse, type AIStreamCallbacksAndOptions } from 'ai' +import { useSession } from '../tabby/auth' const serverUrl = process.env.NEXT_PUBLIC_TABBY_SERVER_URL || '' export function usePatchFetch() { + const { data } = useSession() + useEffect(() => { - const fetch = window.fetch + if (!(window as any)._originFetch) { + (window as any)._originFetch = window.fetch; + } + + const fetch = (window as any)._originFetch as (typeof window.fetch); window.fetch = async function (url, options) { if (url !== '/api/chat') { return fetch(url, options) } + const headers: HeadersInit = { + 'Content-Type': 'application/json', + } + + if (data?.accessToken) { + headers["Authorization"] = `Bearer ${data?.accessToken}`; + } + const res = await fetch(`${serverUrl}/v1beta/chat/completions`, { ...options, method: 'POST', - headers: { - 'Content-Type': 'application/json' - } + headers, }) const stream = StreamAdapter(res, undefined) return new StreamingTextResponse(stream) } - }, []) + }, [data?.accessToken]) } const utf8Decoder = new TextDecoder('utf-8') diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx index 43261c6..0b6d3d1 100644 --- a/ee/tabby-ui/lib/tabby/auth.tsx +++ b/ee/tabby-ui/lib/tabby/auth.tsx @@ -210,6 +210,7 @@ function useSignOut(): () => Promise { interface User { email: string isAdmin: boolean + accessToken: string } type Session = @@ -231,7 +232,8 @@ function useSession(): Session { return { data: { email: user.email, - isAdmin: user.is_admin + isAdmin: user.is_admin, + accessToken: authState.data.accessToken }, status: authState.status } diff --git a/ee/tabby-ui/lib/tabby/fetcher.ts b/ee/tabby-ui/lib/tabby/fetcher.ts index e1c4671..4654eae 100644 --- a/ee/tabby-ui/lib/tabby/fetcher.ts +++ b/ee/tabby-ui/lib/tabby/fetcher.ts @@ -1,9 +1,12 @@ -export default function fetcher(url: string): Promise { - if (process.env.NODE_ENV === 'production') { - return fetch(url).then(x => x.json()) - } else { - return fetch(`${process.env.NEXT_PUBLIC_TABBY_SERVER_URL}${url}`).then(x => - x.json() - ) +export default function tokenFetcher([url, token]: Array): Promise { + const headers = new Headers(); + if (token) { + headers.append("authorization", `Bearer ${token}`) } -} + + if (process.env.NODE_ENV !== 'production') { + url = `${process.env.NEXT_PUBLIC_TABBY_SERVER_URL}${url}`; + } + + return fetch(url!, { headers }).then(x => x.json()) +} \ No newline at end of file diff --git a/ee/tabby-ui/lib/tabby/gql.ts b/ee/tabby-ui/lib/tabby/gql.ts index bfe74a8..bf5f4ce 100644 --- a/ee/tabby-ui/lib/tabby/gql.ts +++ b/ee/tabby-ui/lib/tabby/gql.ts @@ -2,6 +2,7 @@ import { TypedDocumentNode } from '@graphql-typed-document-node/core' import { GraphQLClient, Variables } from 'graphql-request' import { GraphQLResponse } from 'graphql-request/build/esm/types' import useSWR, { SWRConfiguration, SWRResponse } from 'swr' +import { useSession } from './auth' export const gqlClient = new GraphQLClient( `${process.env.NEXT_PUBLIC_TABBY_SERVER_URL ?? ''}/graphql` @@ -26,10 +27,18 @@ export function useGraphQLForm< onError?: (path: string, message: string) => void } ) { - const onSubmit = async (values: TVariables) => { + const { data } = useSession(); + const accessToken = data?.accessToken; + const onSubmit = async (variables: TVariables) => { let res try { - res = await gqlClient.request(document, values) + res = await gqlClient.request({ + document, + variables, + requestHeaders: accessToken ? { + "authorization": `Bearer ${accessToken}` + } : undefined + }) } catch (err) { const { errors = [] } = (err as any).response as GraphQLResponse for (const error of errors) { @@ -61,9 +70,16 @@ export function useGraphQLQuery< variables?: TVariables, swrConfiguration?: SWRConfiguration ): SWRResponse { + const { data } = useSession(); return useSWR( - [document, variables], - ([document, variables]) => gqlClient.request(document, variables), + [document, variables, data?.accessToken], + ([document, variables, accessToken]) => gqlClient.request({ + document, + variables, + requestHeaders: accessToken ? { + "authorization": `Bearer ${accessToken}` + } : undefined + }), swrConfiguration ) } diff --git a/ee/tabby-webserver/src/schema/auth.rs b/ee/tabby-webserver/src/schema/auth.rs index f8b7529..90fb951 100644 --- a/ee/tabby-webserver/src/schema/auth.rs +++ b/ee/tabby-webserver/src/schema/auth.rs @@ -42,8 +42,12 @@ fn jwt_token_secret() -> String { let jwt_secret = match std::env::var("TABBY_WEBSERVER_JWT_TOKEN_SECRET") { Ok(x) => x, Err(_) => { - warn!( - r"TABBY_WEBSERVER_JWT_TOKEN_SECRET is not set. Tabby generates a one-time (non-persisted) JWT token solely for testing purposes." + eprintln!(" + \x1b[93;1mJWT secret is not set\x1b[0m + + Tabby server will generate a one-time (non-persisted) JWT secret for the current process. + Please set the \x1b[94mTABBY_WEBSERVER_JWT_TOKEN_SECRET\x1b[0m environment variable for production usage. +" ); Uuid::new_v4().to_string() } @@ -51,6 +55,7 @@ fn jwt_token_secret() -> String { if uuid::Uuid::parse_str(&jwt_secret).is_err() { warn!("JWT token secret needs to be in standard uuid format to ensure its security, you might generate one at https://www.uuidgenerator.net"); + std::process::exit(1) } jwt_secret @@ -280,7 +285,7 @@ pub trait AuthenticationService: Send + Sync { &self, refresh_token: String, ) -> std::result::Result; - async fn verify_access_token(&self, access_token: String) -> Result; + async fn verify_access_token(&self, access_token: &str) -> Result; async fn is_admin_initialized(&self) -> Result; async fn create_invitation(&self, email: String) -> Result; @@ -293,7 +298,7 @@ mod tests { use super::*; #[test] fn test_generate_jwt() { - let claims = Claims::new(UserInfo::new("test".to_string(), false)); + let claims = Claims::new(UserInfo::new("test".to_string(), false, "cde".to_owned())); let token = generate_jwt(claims).unwrap(); assert!(!token.is_empty()) @@ -301,12 +306,13 @@ mod tests { #[test] fn test_validate_jwt() { - let claims = Claims::new(UserInfo::new("test".to_string(), false)); + let user = UserInfo::new("test".to_string(), false, "cde".to_owned()); + let claims = Claims::new(user.clone()); let token = generate_jwt(claims).unwrap(); let claims = validate_jwt(&token).unwrap(); assert_eq!( claims.user_info(), - &UserInfo::new("test".to_string(), false) + &user, ); } diff --git a/ee/tabby-webserver/src/schema/mod.rs b/ee/tabby-webserver/src/schema/mod.rs index 502d8ca..70b286a 100644 --- a/ee/tabby-webserver/src/schema/mod.rs +++ b/ee/tabby-webserver/src/schema/mod.rs @@ -71,13 +71,37 @@ pub struct Query; #[graphql_object(context = Context)] impl Query { - async fn workers(ctx: &Context) -> Vec { - ctx.locator.worker().list_workers().await + async fn workers(ctx: &Context) -> Result> { + if ctx.locator.auth().is_admin_initialized().await? { + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + let workers = ctx.locator.worker().list_workers().await; + return Ok(workers); + } + } + Err(CoreError::Unauthorized( + "Only admin is able to read workers", + )) + } else { + Ok(ctx.locator.worker().list_workers().await) + } } async fn registration_token(ctx: &Context) -> Result { - let token = ctx.locator.worker().read_registration_token().await?; - Ok(token) + if ctx.locator.auth().is_admin_initialized().await? { + if let Some(claims) = &ctx.claims { + if claims.user_info().is_admin() { + let token = ctx.locator.worker().read_registration_token().await?; + return Ok(token); + } + } + Err(CoreError::Unauthorized( + "Only admin is able to read registeration_token", + )) + } else { + let token = ctx.locator.worker().read_registration_token().await?; + Ok(token) + } } async fn is_admin_initialized(ctx: &Context) -> Result { @@ -142,7 +166,7 @@ impl Mutation { } async fn verify_token(ctx: &Context, token: String) -> Result { - Ok(ctx.locator.auth().verify_access_token(token).await?) + Ok(ctx.locator.auth().verify_access_token(&token).await?) } async fn refresh_token( diff --git a/ee/tabby-webserver/src/service/auth.rs b/ee/tabby-webserver/src/service/auth.rs index 269a598..0b5345c 100644 --- a/ee/tabby-webserver/src/service/auth.rs +++ b/ee/tabby-webserver/src/service/auth.rs @@ -220,8 +220,8 @@ impl AuthenticationService for DbConn { Ok(resp) } - async fn verify_access_token(&self, access_token: String) -> Result { - let claims = validate_jwt(&access_token)?; + async fn verify_access_token(&self, access_token: &str) -> Result { + let claims = validate_jwt(access_token)?; let resp = VerifyTokenResponse::new(claims); Ok(resp) } diff --git a/ee/tabby-webserver/src/service/db/users.rs b/ee/tabby-webserver/src/service/db/users.rs index 9b5888a..afb56a5 100644 --- a/ee/tabby-webserver/src/service/db/users.rs +++ b/ee/tabby-webserver/src/service/db/users.rs @@ -18,7 +18,7 @@ pub struct User { pub is_admin: bool, /// To authenticate IDE extensions / plugins to access code completion / chat api endpoints. - auth_token: String, + pub auth_token: String, } impl User { @@ -54,7 +54,7 @@ impl DbConn { let mut stmt = c.prepare( r#"INSERT INTO users (email, password_encrypted, is_admin, auth_token) VALUES (?, ?, ?, ?)"#, )?; - let id = stmt.insert((email, password_encrypted, is_admin, Uuid::new_v4().to_string()))?; + let id = stmt.insert((email, password_encrypted, is_admin, generate_auth_token()))?; Ok(id) }) .await?; @@ -132,6 +132,11 @@ impl DbConn { } } +fn generate_auth_token() -> String { + let uuid = Uuid::new_v4().to_string().replace("-", ""); + format!("auth_{}", uuid) +} + #[cfg(test)] mod tests { use super::*; diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs index 46f237f..1eea828 100644 --- a/ee/tabby-webserver/src/service/mod.rs +++ b/ee/tabby-webserver/src/service/mod.rs @@ -52,7 +52,7 @@ impl ServerContext { // Authorization is enabled && self.db_conn.is_admin_initialized().await.unwrap_or(false) { - let auth_token = { + let token = { let authorization = request .headers() .get("authorization") @@ -71,8 +71,10 @@ impl ServerContext { } }; - if let Some(auth_token) = auth_token { - if !self.db_conn.verify_auth_token(auth_token).await { + if let Some(token) = token { + if !self.db_conn.verify_access_token(token).await.is_ok() + && !self.db_conn.verify_auth_token(token).await + { return false; } } else {