Merge remote-tracking branch 'upstream/main' into refactor-main

2023-07-30 18:28:10 +08:00 · 2023-07-30 18:28:10 +08:00 · 167cb35637
commit 167cb35637
parent 8db2c0173b 3fca6ff534
7 changed files with 146 additions and 8 deletions
--- a/README.md
+++ b/README.md
@ -94,7 +94,7 @@ _✨ 通过标准的 OpenAI API 格式访问所有的大模型，开箱即用
 19. 支持通过系统访问令牌访问管理 API。
 20. 支持 Cloudflare Turnstile 用户校验。
 21. 支持用户管理，支持**多种用户登录注册方式**：
-    + 邮箱登录注册以及通过邮箱进行密码重置。
+    + 邮箱登录注册（支持注册邮箱白名单）以及通过邮箱进行密码重置。
    + [GitHub 开放授权](https://github.com/settings/applications/new)。
    + 微信公众号授权（需要额外部署 [WeChat Server](https://github.com/songquanpeng/wechat-server)）。

--- a/common/constants.go
+++ b/common/constants.go
@ -45,6 +45,19 @@ var GoogleOAuthEnabled = false
 var TurnstileCheckEnabled = false
 var RegisterEnabled = true

+var EmailDomainRestrictionEnabled = false
+var EmailDomainWhitelist = []string{
+	"gmail.com",
+	"163.com",
+	"126.com",
+	"qq.com",
+	"outlook.com",
+	"hotmail.com",
+	"icloud.com",
+	"yahoo.com",
+	"foxmail.com",
+}
+
 var LogConsumeEnabled = true

 var SMTPServer = ""
--- a/controller/misc.go
+++ b/controller/misc.go
@ -6,6 +6,7 @@ import (
 	"net/http"
 	"one-api/common"
 	"one-api/model"
+	"strings"

 	"github.com/gin-gonic/gin"
 )
@ -83,6 +84,22 @@ func SendEmailVerification(c *gin.Context) {
 		})
 		return
 	}
+	if common.EmailDomainRestrictionEnabled {
+		allowed := false
+		for _, domain := range common.EmailDomainWhitelist {
+			if strings.HasSuffix(email, "@"+domain) {
+				allowed = true
+				break
+			}
+		}
+		if !allowed {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "管理员启用了邮箱域名白名单，您的邮箱地址的域名不在白名单中",
+			})
+			return
+		}
+	}
 	if model.IsEmailAlreadyTaken(email) {
 		c.JSON(http.StatusOK, gin.H{
 			"success": false,
--- a/controller/option.go
+++ b/controller/option.go
@ -58,6 +58,14 @@ func UpdateOption(c *gin.Context) {
 			})
 			return
 		}
+	case "EmailDomainRestrictionEnabled":
+		if option.Value == "true" && len(common.EmailDomainWhitelist) == 0 {
+			c.JSON(http.StatusOK, gin.H{
+				"success": false,
+				"message": "无法启用邮箱域名限制，请先填入限制的邮箱域名！",
+			})
+			return
+		}
 	case "WeChatAuthEnabled":
 		if option.Value == "true" && common.WeChatServerAddress == "" {
 			c.JSON(http.StatusOK, gin.H{
--- a/controller/relay-openai.go
+++ b/controller/relay-openai.go
@ -52,8 +52,10 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
 				}
 			}
 			if !strings.HasPrefix(data, "data:") {
+				if data[:6] != "data: " && data[:6] != "[DONE]" {
 					continue
 				}
+			}
 			dataChan <- data
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
--- a/model/option.go
+++ b/model/option.go
@ -41,6 +41,8 @@ func InitOptionMap() {
 	common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled)
 	common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled)
 	common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64)
+	common.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(common.EmailDomainRestrictionEnabled)
+	common.OptionMap["EmailDomainWhitelist"] = strings.Join(common.EmailDomainWhitelist, ",")
 	common.OptionMap["SMTPServer"] = ""
 	common.OptionMap["SMTPFrom"] = ""
 	common.OptionMap["SMTPPort"] = strconv.Itoa(common.SMTPPort)
@ -151,6 +153,8 @@ func updateOptionMap(key string, value string) (err error) {
 			common.TurnstileCheckEnabled = boolValue
 		case "RegisterEnabled":
 			common.RegisterEnabled = boolValue
+		case "EmailDomainRestrictionEnabled":
+			common.EmailDomainRestrictionEnabled = boolValue
 		case "AutomaticDisableChannelEnabled":
 			common.AutomaticDisableChannelEnabled = boolValue
 		case "ApproximateTokenEnabled":
@ -164,6 +168,8 @@ func updateOptionMap(key string, value string) (err error) {
 		}
 	}
 	switch key {
+	case "EmailDomainWhitelist":
+		common.EmailDomainWhitelist = strings.Split(value, ",")
 	case "SMTPServer":
 		common.SMTPServer = value
 	case "SMTPPort":
--- a/web/src/components/SystemSetting.js
+++ b/web/src/components/SystemSetting.js
@ -1,6 +1,6 @@
 import React, { useEffect, useState } from 'react';
-import { Divider, Form, Grid, Header, Message } from 'semantic-ui-react';
-import { API, removeTrailingSlash, showError, verifyJSON } from '../helpers';
+import { Button, Divider, Form, Grid, Header, Input, Message } from 'semantic-ui-react';
+import { API, removeTrailingSlash, showError } from '../helpers';

 const SystemSetting = () => {
  let [inputs, setInputs] = useState({
@ -32,9 +32,13 @@ const SystemSetting = () => {
    TurnstileSiteKey: '',
    TurnstileSecretKey: '',
    RegisterEnabled: '',
+    EmailDomainRestrictionEnabled: '',
+    EmailDomainWhitelist: ''
  });
  const [originInputs, setOriginInputs] = useState({});
  let [loading, setLoading] = useState(false);
+  const [EmailDomainWhitelist, setEmailDomainWhitelist] = useState([]);
+  const [restrictedDomainInput, setRestrictedDomainInput] = useState('');

  const getOptions = async () => {
    const res = await API.get('/api/option/');
@ -44,8 +48,15 @@ const SystemSetting = () => {
      data.forEach((item) => {
        newInputs[item.key] = item.value;
      });
-      setInputs(newInputs);
+      setInputs({
+        ...newInputs,
+        EmailDomainWhitelist: newInputs.EmailDomainWhitelist.split(',')
+      });
      setOriginInputs(newInputs);
+
+      setEmailDomainWhitelist(newInputs.EmailDomainWhitelist.split(',').map((item) => {
+        return { key: item, text: item, value: item };
+      }));
    } else {
      showError(message);
    }
@ -66,6 +77,7 @@ const SystemSetting = () => {
      case 'WeChatAuthEnabled':
      case 'GoogleOAuthEnabled':
      case 'TurnstileCheckEnabled':
+      case 'EmailDomainRestrictionEnabled':
      case 'RegisterEnabled':
        value = inputs[key] === 'true' ? 'false' : 'true';
        break;
@ -78,7 +90,12 @@ const SystemSetting = () => {
    });
    const { success, message } = res.data;
    if (success) {
-      setInputs((inputs) => ({ ...inputs, [key]: value }));
+      if (key === 'EmailDomainWhitelist') {
+        value = value.split(',');
+      }
+      setInputs((inputs) => ({
+        ...inputs, [key]: value
+      }));
    } else {
      showError(message);
    }
@ -100,7 +117,8 @@ const SystemSetting = () => {
      name === 'GoogleClientId' ||
      name === 'GoogleClientSecret' ||
      name === 'TurnstileSiteKey' ||
-      name === 'TurnstileSecretKey'
+      name === 'TurnstileSecretKey' ||
+      name === 'EmailDomainWhitelist'
    ) {
      setInputs((inputs) => ({ ...inputs, [name]: value }));
    } else {
@ -137,6 +155,16 @@ const SystemSetting = () => {
    }
  };

+
+  const submitEmailDomainWhitelist = async () => {
+    if (
+      originInputs['EmailDomainWhitelist'] !== inputs.EmailDomainWhitelist.join(',') &&
+      inputs.SMTPToken !== ''
+    ) {
+      await updateOption('EmailDomainWhitelist', inputs.EmailDomainWhitelist.join(','));
+    }
+  };
+
  const submitWeChat = async () => {
    if (originInputs['WeChatServerAddress'] !== inputs.WeChatServerAddress) {
      await updateOption(
@ -209,6 +237,22 @@ const SystemSetting = () => {
    }
  };

+  const submitNewRestrictedDomain = () => {
+    const localDomainList = inputs.EmailDomainWhitelist;
+    if (restrictedDomainInput !== '' && !localDomainList.includes(restrictedDomainInput)) {
+      setRestrictedDomainInput('');
+      setInputs({
+        ...inputs,
+        EmailDomainWhitelist: [...localDomainList, restrictedDomainInput],
+      });
+      setEmailDomainWhitelist([...EmailDomainWhitelist, {
+        key: restrictedDomainInput,
+        text: restrictedDomainInput,
+        value: restrictedDomainInput,
+      }]);
+    }
+  }
+
  return (
    <Grid columns={1}>
      <Grid.Column>
@ -287,6 +331,54 @@ const SystemSetting = () => {
            />
          </Form.Group>
          <Divider />
+          <Header as='h3'>
+            配置邮箱域名白名单
+            <Header.Subheader>用以防止恶意用户利用临时邮箱批量注册</Header.Subheader>
+          </Header>
+          <Form.Group widths={3}>
+            <Form.Checkbox
+              label='启用邮箱域名白名单'
+              name='EmailDomainRestrictionEnabled'
+              onChange={handleInputChange}
+              checked={inputs.EmailDomainRestrictionEnabled === 'true'}
+            />
+          </Form.Group>
+          <Form.Group widths={2}>
+            <Form.Dropdown
+              label='允许的邮箱域名'
+              placeholder='允许的邮箱域名'
+              name='EmailDomainWhitelist'
+              required
+              fluid
+              multiple
+              selection
+              onChange={handleInputChange}
+              value={inputs.EmailDomainWhitelist}
+              autoComplete='new-password'
+              options={EmailDomainWhitelist}
+            />
+            <Form.Input
+              label='添加新的允许的邮箱域名'
+              action={
+                <Button type='button' onClick={() => {
+                  submitNewRestrictedDomain();
+                }}>填入</Button>
+              }
+              onKeyDown={(e) => {
+                if (e.key === 'Enter') {
+                  submitNewRestrictedDomain();
+                }
+              }}
+              autoComplete='new-password'
+              placeholder='输入新的允许的邮箱域名'
+              value={restrictedDomainInput}
+              onChange={(e, { value }) => {
+                setRestrictedDomainInput(value);
+              }}
+            />
+          </Form.Group>
+          <Form.Button onClick={submitEmailDomainWhitelist}>保存邮箱域名白名单设置</Form.Button>
+          <Divider />
          <Header as='h3'>
            配置 SMTP
            <Header.Subheader>用以支持系统的邮件发送</Header.Subheader>
@ -332,7 +424,7 @@ const SystemSetting = () => {
              onChange={handleInputChange}
              type='password'
              autoComplete='new-password'
-              value={inputs.SMTPToken}
+              checked={inputs.RegisterEnabled === 'true'}
              placeholder='敏感信息不会发送到前端显示'
            />
          </Form.Group>