feat: 添加socks5本地代理设置

Author: bincooo

chat.go

@@ -3,11 +3,14 @@ package edge
 import (
 	"bytes"
 	"context"
+	"crypto/tls"
 	"encoding/json"
 	"errors"
 	"github.com/RomiChan/websocket"
 	"github.com/google/uuid"
+	"golang.org/x/net/proxy"
 	"io"
+	"net"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -414,9 +417,25 @@ func (c *Chat) newConn() (*wsConn, error) {
 		if err != nil {
 			return nil, err
 		}
-		dialer = &websocket.Dialer{
-			Proxy:            http.ProxyURL(purl),
-			HandshakeTimeout: 45 * time.Second,
+
+		if purl.Scheme == "http" || purl.Scheme == "https" {
+			dialer = &websocket.Dialer{
+				Proxy:            http.ProxyURL(purl),
+				HandshakeTimeout: 45 * time.Second,
+			}
+		}
+
+		if purl.Scheme == "socks5" {
+			dialer = &websocket.Dialer{
+				NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+					d, e := proxy.SOCKS5("tcp", purl.Host, nil, proxy.Direct)
+					if e != nil {
+						return nil, e
+					}
+					return d.Dial(network, addr)
+				},
+				HandshakeTimeout: 45 * time.Second,
+			}
 		}
 	}
 
@@ -553,7 +572,7 @@ func (c *Chat) Delete() error {
 	}
 
 	request.Header = c.newHeader()
-	client, err := c.newClient()
+	client, err := newClient(c.proxies)
 	if err != nil {
 		return &ChatError{"delete", err}
 	}
@@ -600,10 +619,7 @@ func (c *Chat) Delete() error {
 	request.Header = c.newHeader()
 	request.Header.Set("Authorization", "Bearer "+authorization)
 	request.Header.Set("Content-Type", "application/json")
-	client, err = c.newClient()
-	if err != nil {
-		return &ChatError{"delete", err}
-	}
+
 	r, err = client.Do(request)
 	if err != nil {
 		return &ChatError{"delete", err}
@@ -630,7 +646,7 @@ func (c *Chat) LoadPlugins(names ...string) (plugins []string, err error) {
 	}
 
 	request.Header = c.newHeader()
-	client, err := c.newClient()
+	client, err := newClient(c.proxies)
 	if err != nil {
 		return nil, err
 	}
@@ -675,7 +691,7 @@ func (c *Chat) newConversation() (*Conversation, error) {
 	}
 
 	request.Header = c.newHeader()
-	client, err := c.newClient()
+	client, err := newClient(c.proxies)
 	if err != nil {
 		return nil, err
 	}
@@ -715,19 +731,44 @@ func (c *Chat) newConversation() (*Conversation, error) {
 	return &conv, nil
 }
 
-func (c *Chat) newClient() (*http.Client, error) {
+func newClient(proxies string) (*http.Client, error) {
 	client := http.DefaultClient
-	if c.proxies != "" {
-		purl, err := url.Parse(c.proxies)
+	if proxies != "" {
+		proxiesUrl, err := url.Parse(proxies)
 		if err != nil {
 			return nil, err
 		}
-		client = &http.Client{
-			Transport: &http.Transport{
-				Proxy: http.ProxyURL(purl),
-			},
+
+		if proxiesUrl.Scheme == "http" || proxiesUrl.Scheme == "https" {
+			client = &http.Client{
+				Transport: &http.Transport{
+					Proxy: http.ProxyURL(proxiesUrl),
+					TLSClientConfig: &tls.Config{
+						InsecureSkipVerify: true,
+					},
+				},
+			}
+		}
+
+		// socks5://127.0.0.1:7890
+		if proxiesUrl.Scheme == "socks5" {
+			client = &http.Client{
+				Transport: &http.Transport{
+					DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+						dialer, e := proxy.SOCKS5("tcp", proxiesUrl.Host, nil, proxy.Direct)
+						if e != nil {
+							return nil, e
+						}
+						return dialer.Dial(network, addr)
+					},
+					TLSClientConfig: &tls.Config{
+						InsecureSkipVerify: true,
+					},
+				},
+			}
 		}
 	}
+
 	return client, nil
 }
 

examples/chat_test.go

@@ -35,13 +35,13 @@ var pMessages = []edge.ChatMessage{
 }
 
 func Test_Plugins(t *testing.T) {
-	options, err := edge.NewDefaultOptions(cookie, "https://bincooo-single-proxy.hf.space/copilot")
+	options, err := edge.NewDefaultOptions(cookie, "")
 	if err != nil {
 		t.Fatal(err)
 	}
 
 	chat := edge.New(options.
-		Proxies("http://127.0.0.1:7890").
+		Proxies("socks5://127.0.0.1:7890").
 		Model(edge.ModelCreative).
 		Notebook(true))
 
@@ -109,7 +109,7 @@ func Test_classification(t *testing.T) {
 	prompt := "查看上面提供的内容，并总结"
 
 	chat := edge.New(options.
-		Proxies("http://127.0.0.1:7890").
+		Proxies("socks5://127.0.0.1:7890").
 		Model(edge.ModelCreative).
 		Notebook(true))
 	partialResponse, err := chat.Reply(context.Background(), strings.Replace(template, "{{prompt}}", prompt, -1), nil, nil)
@@ -147,7 +147,7 @@ func Test_messages(t *testing.T) {
 	options.KievAuth(KievAuth, RwBf)
 	// Sydney 模式需要自行维护历史对话
 	chat := edge.New(options.
-		Proxies("http://127.0.0.1:7890").
+		Proxies("socks5://127.0.0.1:7890").
 		Model(edge.ModelCreative).
 		Temperature(1.0).
 		TopicToE(true))
@@ -237,7 +237,7 @@ func Test_image(t *testing.T) {
 		t.Fatal(err)
 	}
 	options.KievAuth(KievAuth, RwBf)
-	options.Proxies("http://127.0.0.1:7890")
+	options.Proxies("socks5://127.0.0.1:7890")
 	options.Model(edge.ModelSydney)
 	chat := edge.New(options)
 	file := "/Users/bincooo/Desktop/Screenshot 2023-11-11 at 11.21.23.png"

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/google/uuid v1.3.0
 	github.com/joho/godotenv v1.5.1
 	github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646
+	golang.org/x/net v0.22.0
 )
 
-require golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect
+require golang.org/x/sys v0.18.0 // indirect

go.sum

@@ -6,5 +6,7 @@ github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
 github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
 github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646 h1:zYyBkD/k9seD2A7fsi6Oo2LfFZAehjjQMERAvZLEDnQ=
 github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646/go.mod h1:jpp1/29i3P1S/RLdc7JQKbRpFeM1dOBd8T9ki5s+AY8=
-golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ=
-golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/net v0.22.0 h1:9sGLhx7iRIHEiX0oAJ3MRZMUCElJgy7Br1nO+AMN3Tc=
+golang.org/x/net v0.22.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
+golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
+golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=