Skip to content

Commit 76a0769

Browse files
subash-onyxsubash-onyx
committed
Enhance credential management with multi-auth support and improved validation
1 parent c5adbe4 commit 76a0769

File tree

6 files changed

+393
-130
lines changed

6 files changed

+393
-130
lines changed

‎backend/onyx/connectors/blob/connector.py‎

Lines changed: 58 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import os
2+
import time
23
from datetime import datetime
34
from datetime import timezone
45
from io import BytesIO
@@ -7,9 +8,11 @@
78

89
import boto3 # type: ignore
910
from botocore.client import Config # type: ignore
11+
from botocore.credentials import RefreshableCredentials
1012
from botocore.exceptions import ClientError
1113
from botocore.exceptions import NoCredentialsError
1214
from botocore.exceptions import PartialCredentialsError
15+
from botocore.session import get_session
1316
from mypy_boto3_s3 import S3Client # type: ignore
1417

1518
from onyx.configs.app_configs import INDEX_BATCH_SIZE
@@ -61,7 +64,7 @@ def set_allow_images(self, allow_images: bool) -> None:
6164
def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
6265
"""Checks for boto3 credentials based on the bucket type.
6366
(1) R2: Access Key ID, Secret Access Key, Account ID
64-
(2) S3: AWS Access Key ID, AWS Secret Access Key
67+
(2) S3: AWS Access Key ID, AWS Secret Access Key or IAM role
6568
(3) GOOGLE_CLOUD_STORAGE: Access Key ID, Secret Access Key, Project ID
6669
(4) OCI_STORAGE: Namespace, Region, Access Key ID, Secret Access Key
6770
@@ -95,17 +98,61 @@ def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None
9598
)
9699

97100
elif self.bucket_type == BlobType.S3:
98-
if not all(
99-
credentials.get(key)
100-
for key in ["aws_access_key_id", "aws_secret_access_key"]
101-
):
102-
raise ConnectorMissingCredentialError("Amazon S3")
103-
104-
session = boto3.Session(
105-
aws_access_key_id=credentials["aws_access_key_id"],
106-
aws_secret_access_key=credentials["aws_secret_access_key"],
101+
# For S3, we can use either access keys or IAM roles.
102+
authentication_method = credentials.get(
103+
"authentication_method", "access_key"
104+
)
105+
logger.debug(
106+
f"Using authentication method: {authentication_method} for S3 bucket."
107107
)
108-
self.s3_client = session.client("s3")
108+
if authentication_method == "access_key":
109+
logger.debug("Using access key authentication for S3 bucket.")
110+
if not all(
111+
credentials.get(key)
112+
for key in ["aws_access_key_id", "aws_secret_access_key"]
113+
):
114+
raise ConnectorMissingCredentialError("Amazon S3")
115+
116+
session = boto3.Session(
117+
aws_access_key_id=credentials["aws_access_key_id"],
118+
aws_secret_access_key=credentials["aws_secret_access_key"],
119+
)
120+
self.s3_client = session.client("s3")
121+
elif authentication_method == "iam_role":
122+
# If using IAM roles, we assume the role and let boto3 handle the credentials.
123+
role_arn = credentials.get("aws_role_arn")
124+
# create session name using timestamp
125+
if not role_arn:
126+
raise ConnectorMissingCredentialError(
127+
"Amazon S3 IAM role ARN is required for assuming role."
128+
)
129+
130+
def _refresh_credentials() -> dict[str, str]:
131+
"""Refreshes the credentials for the assumed role."""
132+
sts_client = boto3.client("sts")
133+
assumed_role_object = sts_client.assume_role(
134+
RoleArn=role_arn,
135+
RoleSessionName=f"onyx_blob_storage_{int(time.time())}",
136+
)
137+
creds = assumed_role_object["Credentials"]
138+
return {
139+
"access_key": creds["AccessKeyId"],
140+
"secret_key": creds["SecretAccessKey"],
141+
"token": creds["SessionToken"],
142+
"expiry_time": creds["Expiration"].isoformat(),
143+
}
144+
145+
refreshable = RefreshableCredentials.create_from_metadata(
146+
metadata=_refresh_credentials(),
147+
refresh_using=_refresh_credentials,
148+
method="sts-assume-role",
149+
)
150+
botocore_session = get_session()
151+
botocore_session._credentials = refreshable # type: ignore[attr-defined]
152+
session = boto3.Session(botocore_session=botocore_session)
153+
self.s3_client = session.client("s3")
154+
else:
155+
raise ConnectorValidationError("Invalid authentication method for S3. ")
109156

110157
elif self.bucket_type == BlobType.GOOGLE_CLOUD_STORAGE:
111158
if not all(

‎web/src/components/credentials/actions/CreateCredential.tsx‎

Lines changed: 103 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import {
2929
} from "@/components/IsPublicGroupSelector";
3030
import { useUser } from "@/components/user/UserProvider";
3131
import CardSection from "@/components/admin/CardSection";
32+
import { CredentialFieldsRenderer } from "./CredentialFieldsRenderer";
3233

3334
const CreateButton = ({
3435
onClick,
@@ -92,6 +93,7 @@ export default function CreateCredential({
9293
refresh?: () => void;
9394
}) {
9495
const [showAdvancedOptions, setShowAdvancedOptions] = useState(false);
96+
const [authMethod, setAuthMethod] = useState<string>();
9597
const isPaidEnterpriseFeaturesEnabled = usePaidEnterpriseFeaturesEnabled();
9698

9799
const { isAdmin } = useUser();
@@ -175,118 +177,122 @@ export default function CreateCredential({
175177
const credentialTemplate: dictionaryType = credentialTemplates[sourceType];
176178
const validationSchema = createValidationSchema(credentialTemplate);
177179

180+
// Set initial auth method for templates with multiple auth methods
181+
const templateWithAuth = credentialTemplate as any;
182+
const initialAuthMethod =
183+
templateWithAuth?.authMethods?.[0]?.value || undefined;
184+
178185
return (
179186
<Formik
180187
initialValues={
181188
{
182189
name: "",
183190
is_public: isAdmin || !isPaidEnterpriseFeaturesEnabled,
184191
groups: [],
192+
...(initialAuthMethod && {
193+
authentication_method: initialAuthMethod,
194+
}),
185195
} as formType
186196
}
187197
validationSchema={validationSchema}
188198
onSubmit={() => {}} // This will be overridden by our custom submit handlers
189199
>
190-
{(formikProps) => (
191-
<Form className="w-full flex items-stretch">
192-
{!hideSource && (
193-
<p className="text-sm">
194-
Check our
195-
<a
196-
className="text-blue-600 hover:underline"
197-
target="_blank"
198-
href={getSourceDocLink(sourceType) || ""}
199-
>
200-
{" "}
201-
docs{" "}
202-
</a>
203-
for information on setting up this connector.
204-
</p>
205-
)}
206-
<CardSection className="w-full items-start dark:bg-neutral-900 mt-4 flex flex-col gap-y-6">
207-
<TextFormField
208-
name="name"
209-
placeholder="(Optional) credential name.."
210-
label="Name:"
211-
/>
212-
{Object.entries(credentialTemplate).map(([key, val]) => {
213-
if (typeof val === "boolean") {
214-
return (
215-
<BooleanFormField
216-
key={key}
217-
name={key}
218-
label={getDisplayNameForCredentialKey(key)}
219-
/>
220-
);
221-
}
222-
return (
223-
<TextFormField
224-
key={key}
225-
name={key}
226-
placeholder={val}
227-
label={getDisplayNameForCredentialKey(key)}
228-
type={
229-
key.toLowerCase().includes("token") ||
230-
key.toLowerCase().includes("password")
231-
? "password"
232-
: "text"
233-
}
234-
/>
235-
);
236-
})}
237-
{!swapConnector && (
238-
<div className="mt-4 flex w-full flex-col sm:flex-row justify-between items-end">
239-
<div className="w-full sm:w-3/4 mb-4 sm:mb-0">
240-
{isPaidEnterpriseFeaturesEnabled && (
241-
<div className="flex flex-col items-start">
242-
{isAdmin && (
243-
<AdvancedOptionsToggle
244-
showAdvancedOptions={showAdvancedOptions}
245-
setShowAdvancedOptions={setShowAdvancedOptions}
246-
/>
247-
)}
248-
{(showAdvancedOptions || !isAdmin) && (
249-
<IsPublicGroupSelector
250-
formikProps={formikProps}
251-
objectName="credential"
252-
publicToWhom="Curators"
253-
/>
254-
)}
255-
</div>
256-
)}
257-
</div>
258-
<div className="w-full sm:w-1/4">
259-
<CreateButton
260-
onClick={() =>
261-
handleSubmit(formikProps.values, formikProps, "create")
262-
}
263-
isSubmitting={formikProps.isSubmitting}
264-
isAdmin={isAdmin}
265-
groups={formikProps.values.groups}
266-
/>
200+
{(formikProps) => {
201+
// Update authentication_method in formik when authMethod changes
202+
if (
203+
authMethod &&
204+
formikProps.values.authentication_method !== authMethod
205+
) {
206+
formikProps.setFieldValue("authentication_method", authMethod);
207+
}
208+
209+
return (
210+
<Form className="w-full flex items-stretch">
211+
{!hideSource && (
212+
<p className="text-sm">
213+
Check our
214+
<a
215+
className="text-blue-600 hover:underline"
216+
target="_blank"
217+
href={getSourceDocLink(sourceType) || ""}
218+
>
219+
{" "}
220+
docs{" "}
221+
</a>
222+
for information on setting up this connector.
223+
</p>
224+
)}
225+
<CardSection className="w-full items-start dark:bg-neutral-900 mt-4 flex flex-col gap-y-6">
226+
<TextFormField
227+
name="name"
228+
placeholder="(Optional) credential name.."
229+
label="Name:"
230+
/>
231+
232+
<CredentialFieldsRenderer
233+
credentialTemplate={credentialTemplate}
234+
authMethod={authMethod || initialAuthMethod}
235+
setAuthMethod={setAuthMethod}
236+
/>
237+
238+
{!swapConnector && (
239+
<div className="mt-4 flex w-full flex-col sm:flex-row justify-between items-end">
240+
<div className="w-full sm:w-3/4 mb-4 sm:mb-0">
241+
{isPaidEnterpriseFeaturesEnabled && (
242+
<div className="flex flex-col items-start">
243+
{isAdmin && (
244+
<AdvancedOptionsToggle
245+
showAdvancedOptions={showAdvancedOptions}
246+
setShowAdvancedOptions={setShowAdvancedOptions}
247+
/>
248+
)}
249+
{(showAdvancedOptions || !isAdmin) && (
250+
<IsPublicGroupSelector
251+
formikProps={formikProps}
252+
objectName="credential"
253+
publicToWhom="Curators"
254+
/>
255+
)}
256+
</div>
257+
)}
258+
</div>
259+
<div className="w-full sm:w-1/4">
260+
<CreateButton
261+
onClick={() =>
262+
handleSubmit(formikProps.values, formikProps, "create")
263+
}
264+
isSubmitting={formikProps.isSubmitting}
265+
isAdmin={isAdmin}
266+
groups={formikProps.values.groups}
267+
/>
268+
</div>
267269
</div>
270+
)}
271+
</CardSection>
272+
{swapConnector && (
273+
<div className="flex gap-x-4 w-full mt-8 justify-end">
274+
<Button
275+
className="bg-rose-500 hover:bg-rose-400 border-rose-800"
276+
onClick={() =>
277+
handleSubmit(
278+
formikProps.values,
279+
formikProps,
280+
"createAndSwap"
281+
)
282+
}
283+
type="button"
284+
disabled={formikProps.isSubmitting}
285+
>
286+
<div className="flex gap-x-2 items-center w-full border-none">
287+
<FaAccusoft />
288+
<p>Create</p>
289+
</div>
290+
</Button>
268291
</div>
269292
)}
270-
</CardSection>
271-
{swapConnector && (
272-
<div className="flex gap-x-4 w-full mt-8 justify-end">
273-
<Button
274-
className="bg-rose-500 hover:bg-rose-400 border-rose-800"
275-
onClick={() =>
276-
handleSubmit(formikProps.values, formikProps, "createAndSwap")
277-
}
278-
type="button"
279-
disabled={formikProps.isSubmitting}
280-
>
281-
<div className="flex gap-x-2 items-center w-full border-none">
282-
<FaAccusoft />
283-
<p>Create</p>
284-
</div>
285-
</Button>
286-
</div>
287-
)}
288-
</Form>
289-
)}
293+
</Form>
294+
);
295+
}}
290296
</Formik>
291297
);
292298
}

0 commit comments

Comments
 (0)