From 4043fccedbc52211a831942bc4297c53f76a550c Mon Sep 17 00:00:00 2001
From: ckt1031 <65409152+ckt1031@users.noreply.github.com>
Date: Fri, 14 Jul 2023 21:30:13 +0800
Subject: [PATCH] feat: support ip randomize in http header

---
 common/ip-gen.go                     | 16 ++++++++++++++++
 controller/channel-test.go           | 14 ++++++++++++++
 controller/relay-text.go             | 14 ++++++++++++++
 middleware/distributor.go            |  1 +
 model/channel.go                     |  6 +++++-
 web/src/pages/Channel/EditChannel.js | 15 ++++++++++++++-
 6 files changed, 64 insertions(+), 2 deletions(-)
 create mode 100644 common/ip-gen.go

diff --git a/common/ip-gen.go b/common/ip-gen.go
new file mode 100644
index 00000000..b09da224
--- /dev/null
+++ b/common/ip-gen.go
@@ -0,0 +1,16 @@
+package common
+
+import (
+	"fmt"
+	"math/rand"
+)
+
+func GenerateIP() string {
+	// Generate a random number between 20 and 240
+	segment2 := rand.Intn(221) + 20
+	segment3 := rand.Intn(256)
+	segment4 := rand.Intn(256)
+
+	ipAddress := fmt.Sprintf("104.%d.%d.%d", segment2, segment3, segment4)
+	return ipAddress
+}
diff --git a/controller/channel-test.go b/controller/channel-test.go
index 3f3ff2b5..54bb0c71 100644
--- a/controller/channel-test.go
+++ b/controller/channel-test.go
@@ -48,6 +48,20 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
 		req.Header.Set("Authorization", "Bearer "+channel.Key)
 	}
 	req.Header.Set("Content-Type", "application/json")
+
+	if channel.EnableIpRandomization {
+		// Generate random IP
+		ip := common.GenerateIP()
+		req.Header.Set("X-Forwarded-For", ip)
+		req.Header.Set("X-Real-IP", ip)
+		req.Header.Set("X-Client-IP", ip)
+		req.Header.Set("X-Forwarded-Host", ip)
+		req.Header.Set("X-Originating-IP", ip)
+		req.RemoteAddr = ip
+		req.Header.Set("X-Remote-IP", ip)
+		req.Header.Set("X-Remote-Addr", ip)
+	}
+
 	client := &http.Client{}
 	resp, err := client.Do(req)
 	if err != nil {
diff --git a/controller/relay-text.go b/controller/relay-text.go
index 363c7f4d..cf20fc72 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -175,6 +175,20 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
 	req.Header.Set("Accept", c.Request.Header.Get("Accept"))
 	//req.Header.Set("Connection", c.Request.Header.Get("Connection"))
+
+	if c.GetBool("enable_ip_randomization") == true {
+		// Generate random IP
+		ip := common.GenerateIP()
+		req.Header.Set("X-Forwarded-For", ip)
+		req.Header.Set("X-Real-IP", ip)
+		req.Header.Set("X-Client-IP", ip)
+		req.Header.Set("X-Forwarded-Host", ip)
+		req.Header.Set("X-Originating-IP", ip)
+		req.RemoteAddr = ip
+		req.Header.Set("X-Remote-IP", ip)
+		req.Header.Set("X-Remote-Addr", ip)
+	}
+
 	client := &http.Client{}
 	resp, err := client.Do(req)
 	if err != nil {
diff --git a/middleware/distributor.go b/middleware/distributor.go
index f985d1a7..8eb17e1c 100644
--- a/middleware/distributor.go
+++ b/middleware/distributor.go
@@ -100,6 +100,7 @@ func Distribute() func(c *gin.Context) {
 		c.Set("channel_id", channel.Id)
 		c.Set("channel_name", channel.Name)
 		c.Set("model_mapping", channel.ModelMapping)
+		c.Set("enable_ip_randomization", channel.EnableIpRandomization)
 		c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
 		c.Set("base_url", channel.BaseURL)
 		if channel.Type == common.ChannelTypeAzure {
diff --git a/model/channel.go b/model/channel.go
index 7cc9fa9b..fee4e964 100644
--- a/model/channel.go
+++ b/model/channel.go
@@ -1,8 +1,9 @@
 package model
 
 import (
-	"gorm.io/gorm"
 	"one-api/common"
+
+	"gorm.io/gorm"
 )
 
 type Channel struct {
@@ -23,6 +24,9 @@ type Channel struct {
 	Group              string  `json:"group" gorm:"type:varchar(32);default:'default'"`
 	UsedQuota          int64   `json:"used_quota" gorm:"bigint;default:0"`
 	ModelMapping       string  `json:"model_mapping" gorm:"type:varchar(1024);default:''"`
+
+	// Additional fields, default value is false
+	EnableIpRandomization bool `json:"enable_ip_randomization"`
 }
 
 func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {
diff --git a/web/src/pages/Channel/EditChannel.js b/web/src/pages/Channel/EditChannel.js
index b0ba40e3..7fd6d7c6 100644
--- a/web/src/pages/Channel/EditChannel.js
+++ b/web/src/pages/Channel/EditChannel.js
@@ -23,7 +23,8 @@ const EditChannel = () => {
     other: '',
     model_mapping: '',
     models: [],
-    groups: ['default']
+    groups: ['default'],
+    enable_ip_randomization: false,
   };
   const [batch, setBatch] = useState(false);
   const [inputs, setInputs] = useState(originInputs);
@@ -32,6 +33,7 @@ const EditChannel = () => {
   const [basicModels, setBasicModels] = useState([]);
   const [fullModels, setFullModels] = useState([]);
   const handleInputChange = (e, { name, value }) => {
+    console.log(name, value)
     setInputs((inputs) => ({ ...inputs, [name]: value }));
   };
 
@@ -264,6 +266,17 @@ const EditChannel = () => {
               handleInputChange(null, { name: 'models', value: [] });
             }}>清除所有模型</Button>
           </div>
+          <Form.Field>
+            <Form.Checkbox
+              name='enable_ip_randomization'
+              label='将IP随机地址传递给HTTP头'
+              onChange={(e, { name, checked }) => {
+                handleInputChange(e, { name, value: checked });
+              }}
+              checked={inputs.enable_ip_randomization}
+              autoComplete='new-password'
+            />
+          </Form.Field>
           <Form.Field>
             <Form.TextArea
               label='模型映射'