Przeglądaj źródła

feat: support proxy

Yeuoly 1 rok temu
rodzic
commit
f7711fe3d6

+ 2 - 0
conf/config.yaml

@@ -7,3 +7,5 @@ max_requests: 50
 worker_timeout: 5
 python_path: /usr/local/bin/python3
 enable_network: True # please make sure there is no network risk in your environment
+proxy:
+  socks5: '192.168.64.1:7890'

+ 1 - 1
docker/amd64/dockerfile

@@ -10,7 +10,7 @@ COPY main /main
 COPY conf/config.yaml /conf/config.yaml
 RUN rm -rf /var/lib/apt/lists/* \
     && chmod +x /main \
-    && pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple jinja2 requests httpx \
+    && pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple jinja2 requests httpx PySocks httpx[socks] \
     && wget -O /opt/node-v20.11.1-linux-x64.tar.xz https://npmmirror.com/mirrors/node/v20.11.1/node-v20.11.1-linux-x64.tar.xz \
     && tar -xvf /opt/node-v20.11.1-linux-x64.tar.xz -C /opt \
     && ln -s /opt/node-v20.11.1-linux-x64/bin/node /usr/local/bin/node \

+ 1 - 1
docker/arm64/dockerfile

@@ -10,7 +10,7 @@ COPY main /main
 COPY conf/config.yaml /conf/config.yaml
 RUN rm -rf /var/lib/apt/lists/* \
     && chmod +x /main \
-    && pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple jinja2 requests httpx numpy \
+    && pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple jinja2 requests httpx numpy PySocks httpx[socks] \
     && wget -O /opt/node-v20.11.1-linux-arm64.tar.xz https://npmmirror.com/mirrors/node/v20.11.1/node-v20.11.1-linux-arm64.tar.xz \
     && tar -xvf /opt/node-v20.11.1-linux-arm64.tar.xz -C /opt \
     && ln -s /opt/node-v20.11.1-linux-arm64/bin/node /usr/local/bin/node \

+ 2 - 2
internal/core/runner/python/dependencies/network.go

@@ -1,6 +1,6 @@
 package dependencies
 
 func init() {
-	SetupDependency("httpx", "", "import httpx\nimport encodings.idna")
-	SetupDependency("requests", "", "import requests\nfrom netrc import netrc, NetrcParseError\nimport urllib3\nimport socket")
+	SetupDependency("httpx", "", "import httpx\nimport encodings.idna\nimport socks")
+	SetupDependency("requests", "", "import requests\nfrom netrc import netrc, NetrcParseError\nimport urllib3\nimport socket\nimport socks")
 }

+ 8 - 1
internal/core/runner/python/python.go

@@ -44,6 +44,8 @@ func (p *PythonRunner) Run(
 	preload string,
 	options *types.RunnerOptions,
 ) (chan []byte, chan []byte, chan bool, error) {
+	configuration := static.GetDifySandboxGlobalConfigurations()
+
 	// initialize the environment
 	untrusted_code_path, err := p.InitializeEnvironment(code, preload, options)
 	if err != nil {
@@ -63,11 +65,16 @@ func (p *PythonRunner) Run(
 
 		// create a new process
 		cmd := exec.Command(
-			static.GetDifySandboxGlobalConfigurations().PythonPath,
+			configuration.PythonPath,
 			untrusted_code_path,
 		)
 		cmd.Env = []string{}
 
+		if configuration.Proxy.Socks5 != "" {
+			cmd.Env = append(cmd.Env, fmt.Sprintf("HTTPS_PROXY=socks5://%s", configuration.Proxy.Socks5))
+			cmd.Env = append(cmd.Env, fmt.Sprintf("HTTP_PROXY=socks5://%s", configuration.Proxy.Socks5))
+		}
+
 		err = output_handler.CaptureOutput(cmd)
 		if err != nil {
 			return err

+ 6 - 0
internal/static/config.go

@@ -86,6 +86,12 @@ func InitConfig(path string) error {
 		log.Info("network has been enabled")
 	}
 
+	socks5_proxy := os.Getenv("SOCKS5_PROXY")
+	if socks5_proxy != "" {
+		difySandboxGlobalConfigurations.Proxy.Socks5 = socks5_proxy
+
+	}
+
 	return nil
 }
 

+ 3 - 0
internal/types/config.go

@@ -12,4 +12,7 @@ type DifySandboxGlobalConfigurations struct {
 	PythonPath    string `yaml:"python_path"`
 	NodejsPath    string `yaml:"nodejs_path"`
 	EnableNetwork bool   `yaml:"enable_network"`
+	Proxy         struct {
+		Socks5 string `yaml:"socks5"`
+	} `yaml:"proxy"`
 }