Add allowed domains settings and disallow prediction based on this list

packages/server/src/index.ts

@@ -1219,7 +1219,23 @@ export class App {
             upload.array('files'),
             (req: Request, res: Response, next: NextFunction) => getRateLimiter(req, res, next),
             async (req: Request, res: Response) => {
+                const chatflow = await this.AppDataSource.getRepository(ChatFlow).findOneBy({
+                    id: req.params.id
+                })
+                if (!chatflow) return res.status(404).send(`Chatflow ${req.params.id} not found`)
+                let isDomainAllowed = true
+                if (chatflow.chatbotConfig) {
+                    const parsedConfig = JSON.parse(chatflow.chatbotConfig)
+                    if (parsedConfig.allowedDomains && parsedConfig.allowedDomains.length > 0) {
+                        isDomainAllowed = parsedConfig.allowedDomains.includes(req.headers.host)
+                    }
+                }
+
+                if (isDomainAllowed) {
                     await this.buildChatflow(req, res, socketIO)
+                } else {
+                    return res.status(401).send(`This domain is not allowed to access chatflow ${req.params.id}`)
+                }
             }
         )
 

packages/ui/src/menu-items/settings.js

@@ -1,8 +1,17 @@
 // assets
-import { IconTrash, IconFileUpload, IconFileExport, IconCopy, IconSearch, IconMessage, IconPictureInPictureOff } from '@tabler/icons'
+import {
+    IconTrash,
+    IconFileUpload,
+    IconFileExport,
+    IconCopy,
+    IconSearch,
+    IconMessage,
+    IconPictureInPictureOff,
+    IconLink
+} from '@tabler/icons'
 
 // constant
-const icons = { IconTrash, IconFileUpload, IconFileExport, IconCopy, IconSearch, IconMessage, IconPictureInPictureOff }
+const icons = { IconTrash, IconFileUpload, IconFileExport, IconCopy, IconSearch, IconMessage, IconPictureInPictureOff, IconLink }
 
 // ==============================|| SETTINGS MENU ITEMS ||============================== //
 
@@ -25,6 +34,13 @@ const settings = {
             url: '',
             icon: icons.IconMessage
         },
+        {
+            id: 'allowedDomains',
+            title: 'Allowed Domains',
+            type: 'item',
+            url: '',
+            icon: icons.IconLink
+        },
         {
             id: 'duplicateChatflow',
             title: 'Duplicate Chatflow',

packages/ui/src/ui-component/dialog/AllowedDomainsDialog.js

@@ -0,0 +1,238 @@
+import { createPortal } from 'react-dom'
+import { useDispatch } from 'react-redux'
+import { useState, useEffect } from 'react'
+import PropTypes from 'prop-types'
+import { enqueueSnackbar as enqueueSnackbarAction, closeSnackbar as closeSnackbarAction, SET_CHATFLOW } from 'store/actions'
+
+// material-ui
+import {
+    Button,
+    IconButton,
+    Dialog,
+    DialogContent,
+    OutlinedInput,
+    DialogTitle,
+    DialogActions,
+    Box,
+    List,
+    InputAdornment
+} from '@mui/material'
+import { IconX, IconTrash, IconPlus } from '@tabler/icons'
+
+// Project import
+import { StyledButton } from 'ui-component/button/StyledButton'
+
+// store
+import { HIDE_CANVAS_DIALOG, SHOW_CANVAS_DIALOG } from 'store/actions'
+import useNotifier from 'utils/useNotifier'
+
+// API
+import chatflowsApi from 'api/chatflows'
+
+const AllowedDomainsDialog = ({ show, dialogProps, onCancel, onConfirm }) => {
+    const portalElement = document.getElementById('portal')
+    const dispatch = useDispatch()
+
+    useNotifier()
+
+    const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args))
+    const closeSnackbar = (...args) => dispatch(closeSnackbarAction(...args))
+
+    const [inputFields, setInputFields] = useState([
+        {
+            origin: ''
+        }
+    ])
+
+    const [chatbotConfig, setChatbotConfig] = useState({})
+
+    const addInputField = () => {
+        setInputFields([
+            ...inputFields,
+            {
+                origin: ''
+            }
+        ])
+    }
+    const removeInputFields = (index) => {
+        const rows = [...inputFields]
+        rows.splice(index, 1)
+        setInputFields(rows)
+    }
+
+    const handleChange = (index, evnt) => {
+        const { name, value } = evnt.target
+        const list = [...inputFields]
+        list[index][name] = value
+        setInputFields(list)
+    }
+
+    const onSave = async () => {
+        try {
+            let value = {
+                allowedOrigins: {
+                    ...inputFields
+                }
+            }
+            chatbotConfig.allowedOrigins = value.allowedOrigins
+            const saveResp = await chatflowsApi.updateChatflow(dialogProps.chatflow.id, {
+                chatbotConfig: JSON.stringify(chatbotConfig)
+            })
+            if (saveResp.data) {
+                enqueueSnackbar({
+                    message: 'Allowed Origins Saved',
+                    options: {
+                        key: new Date().getTime() + Math.random(),
+                        variant: 'success',
+                        action: (key) => (
+                            <Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
+                                <IconX />
+                            </Button>
+                        )
+                    }
+                })
+                dispatch({ type: SET_CHATFLOW, chatflow: saveResp.data })
+            }
+            onConfirm()
+        } catch (error) {
+            const errorData = error.response.data || `${error.response.status}: ${error.response.statusText}`
+            enqueueSnackbar({
+                message: `Failed to save Allowed Origins: ${errorData}`,
+                options: {
+                    key: new Date().getTime() + Math.random(),
+                    variant: 'error',
+                    persist: true,
+                    action: (key) => (
+                        <Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
+                            <IconX />
+                        </Button>
+                    )
+                }
+            })
+        }
+    }
+
+    useEffect(() => {
+        if (dialogProps.chatflow && dialogProps.chatflow.chatbotConfig) {
+            try {
+                let chatbotConfig = JSON.parse(dialogProps.chatflow.chatbotConfig)
+                setChatbotConfig(chatbotConfig || {})
+                if (chatbotConfig.allowedOrigins) {
+                    let inputFields = []
+                    Object.getOwnPropertyNames(chatbotConfig.allowedOrigins).forEach((key) => {
+                        if (chatbotConfig.allowedOrigins[key]) {
+                            inputFields.push(chatbotConfig.allowedOrigins[key])
+                        }
+                    })
+                    setInputFields(inputFields)
+                } else {
+                    setInputFields([
+                        {
+                            origin: ''
+                        }
+                    ])
+                }
+            } catch (e) {
+                setInputFields([
+                    {
+                        origin: ''
+                    }
+                ])
+            }
+        }
+
+        return () => {}
+    }, [dialogProps])
+
+    useEffect(() => {
+        if (show) dispatch({ type: SHOW_CANVAS_DIALOG })
+        else dispatch({ type: HIDE_CANVAS_DIALOG })
+        return () => dispatch({ type: HIDE_CANVAS_DIALOG })
+    }, [show, dispatch])
+
+    const component = show ? (
+        <Dialog
+            onClose={onCancel}
+            open={show}
+            fullWidth
+            maxWidth='sm'
+            aria-labelledby='alert-dialog-title'
+            aria-describedby='alert-dialog-description'
+        >
+            <DialogTitle sx={{ fontSize: '1rem' }} id='alert-dialog-title'>
+                {dialogProps.title || 'Allowed Origins'}
+            </DialogTitle>
+            <DialogContent>
+                <div
+                    style={{
+                        display: 'flex',
+                        flexDirection: 'column'
+                    }}
+                >
+                    <span>Your chatbot will only work when used from the following domains</span>
+                </div>
+                <Box sx={{ '& > :not(style)': { m: 1 }, pt: 2 }}>
+                    <List>
+                        {inputFields.map((data, index) => {
+                            return (
+                                <div key={index} style={{ display: 'flex', width: '100%' }}>
+                                    <Box sx={{ width: '100%', mb: 1 }}>
+                                        <OutlinedInput
+                                            sx={{ width: '100%' }}
+                                            key={index}
+                                            type='text'
+                                            onChange={(e) => handleChange(index, e)}
+                                            size='small'
+                                            value={data.prompt}
+                                            name='prompt'
+                                            endAdornment={
+                                                <InputAdornment position='end' sx={{ padding: '2px' }}>
+                                                    {inputFields.length > 1 && (
+                                                        <IconButton
+                                                            sx={{ height: 30, width: 30 }}
+                                                            size='small'
+                                                            color='error'
+                                                            disabled={inputFields.length === 1}
+                                                            onClick={() => removeInputFields(index)}
+                                                            edge='end'
+                                                        >
+                                                            <IconTrash />
+                                                        </IconButton>
+                                                    )}
+                                                </InputAdornment>
+                                            }
+                                        />
+                                    </Box>
+                                    <Box sx={{ width: '5%', mb: 1 }}>
+                                        {index === inputFields.length - 1 && (
+                                            <IconButton color='primary' onClick={addInputField}>
+                                                <IconPlus />
+                                            </IconButton>
+                                        )}
+                                    </Box>
+                                </div>
+                            )
+                        })}
+                    </List>
+                </Box>
+            </DialogContent>
+            <DialogActions>
+                <Button onClick={onCancel}>Cancel</Button>
+                <StyledButton variant='contained' onClick={onSave}>
+                    Save
+                </StyledButton>
+            </DialogActions>
+        </Dialog>
+    ) : null
+
+    return createPortal(component, portalElement)
+}
+
+AllowedDomainsDialog.propTypes = {
+    show: PropTypes.bool,
+    dialogProps: PropTypes.object,
+    onCancel: PropTypes.func,
+    onConfirm: PropTypes.func
+}
+
+export default AllowedDomainsDialog

packages/ui/src/views/canvas/CanvasHeader.js

@@ -17,6 +17,7 @@ import APICodeDialog from 'views/chatflows/APICodeDialog'
 import AnalyseFlowDialog from 'ui-component/dialog/AnalyseFlowDialog'
 import ViewMessagesDialog from 'ui-component/dialog/ViewMessagesDialog'
 import StarterPromptsDialog from 'ui-component/dialog/StarterPromptsDialog'
+import AllowedDomainsDialog from 'ui-component/dialog/AllowedDomainsDialog'
 
 // API
 import chatflowsApi from 'api/chatflows'
@@ -50,6 +51,8 @@ const CanvasHeader = ({ chatflow, handleSaveFlow, handleDeleteFlow, handleLoadFl
     const [conversationStartersDialogProps, setConversationStartersDialogProps] = useState({})
     const [viewMessagesDialogOpen, setViewMessagesDialogOpen] = useState(false)
     const [viewMessagesDialogProps, setViewMessagesDialogProps] = useState({})
+    const [allowedDomainsDialogOpen, setAllowedDomainsDialogOpen] = useState(false)
+    const [allowedDomainsDialogProps, setAllowedDomainsDialogProps] = useState({})
 
     const updateChatflowApi = useApi(chatflowsApi.updateChatflow)
     const canvas = useSelector((state) => state.canvas)
@@ -65,6 +68,12 @@ const CanvasHeader = ({ chatflow, handleSaveFlow, handleDeleteFlow, handleLoadFl
                 chatflow: chatflow
             })
             setConversationStartersDialogOpen(true)
+        } else if (setting === 'allowedDomains') {
+            setAllowedDomainsDialogProps({
+                title: 'Starter Prompts - ' + chatflow.name,
+                chatflow: chatflow
+            })
+            setAllowedDomainsDialogOpen(true)
         } else if (setting === 'analyseChatflow') {
             setAnalyseDialogProps({
                 title: 'Analyse Chatflow',
@@ -391,6 +400,12 @@ const CanvasHeader = ({ chatflow, handleSaveFlow, handleDeleteFlow, handleLoadFl
                 onConfirm={() => setConversationStartersDialogOpen(false)}
                 onCancel={() => setConversationStartersDialogOpen(false)}
             />
+            <AllowedDomainsDialog
+                show={allowedDomainsDialogOpen}
+                dialogProps={allowedDomainsDialogProps}
+                onConfirm={() => setAllowedDomainsDialogOpen(false)}
+                onCancel={() => setAllowedDomainsDialogOpen(false)}
+            />
             <ViewMessagesDialog
                 show={viewMessagesDialogOpen}
                 dialogProps={viewMessagesDialogProps}