From 9a8e7af831f9e66ee028c6e336c35825bedacfaa Mon Sep 17 00:00:00 2001 From: Luxferre Date: Tue, 1 Sep 2026 11:02:25 +0300 Subject: [PATCH] socks5 support --- README.md | 2 ++ main.go | 32 +++++++++++++++++++-- main_test.go | 80 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 112 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index a5b8df4..c98ea10 100644 --- a/README.md +++ b/README.md @@ -153,6 +153,8 @@ If the API rejects the request with an `Invalid assistant message: content or to - `bantam_tools_dir` (optional path to a directory of extra shell tools; the Go port appends `"Extra shell tools can be found at "` to the system prompt at startup when set. The `BANTAM_TOOLS_DIR` environment variable overrides this and is checked first; if neither is set, nothing is appended. Note: this is an agent-internal hint, not forwarded to the API.) +The Go port also supports SOCKS5 proxying via the `SOCKS_PROXY` (or `socks_proxy`) environment variable (e.g. `SOCKS_PROXY=socks5://127.0.0.1:1080` or `SOCKS_PROXY=127.0.0.1:1080`), falling back to standard `HTTP_PROXY` / `HTTPS_PROXY` environment variables. + The Go port's built-in editor tracks the cursor with its own column math (terminal auto-wrap aware) and redraws from the first line of the buffer, so wrapped input stays clean at any terminal width. When enabled, the interactive console uses a subtle ANSI palette: the pending-request status `...requesting...` is darkened bold, reasoning markers are cyan, reasoning text is dim, `[tool call: ...]` traces are yellow, `[tool result: ...]` headers are green, tool result bodies are dim (red for tool errors/unknown tools), and errors/network retries are red. Tool result payloads fed back to the LLM are never colored. diff --git a/main.go b/main.go index ffb9f67..699be72 100644 --- a/main.go +++ b/main.go @@ -15,6 +15,7 @@ import ( "io" "net" "net/http" + "net/url" "os" "os/exec" "os/signal" @@ -61,9 +62,24 @@ func internalKey(k string) bool { return false } +func proxyFromEnv(req *http.Request) (*url.URL, error) { + socks := strings.TrimSpace(os.Getenv("SOCKS_PROXY")) + if socks == "" { + socks = strings.TrimSpace(os.Getenv("socks_proxy")) + } + if socks != "" { + if !strings.Contains(socks, "://") { + socks = "socks5://" + socks + } + return url.Parse(socks) + } + return http.ProxyFromEnvironment(req) +} + // llmTransport is a shared HTTP transport reused across all LLM calls so that // connections are pooled instead of recreated per request. var llmTransport = &http.Transport{ + Proxy: proxyFromEnv, DialContext: (&net.Dialer{Timeout: 300 * time.Second}).DialContext, } @@ -77,7 +93,13 @@ func atoiD(s string, d int) int { } func queryModelsContextWindow(cfg *Cfg) int { - client := &http.Client{Timeout: 3 * time.Second} + client := &http.Client{ + Transport: &http.Transport{ + Proxy: proxyFromEnv, + DialContext: (&net.Dialer{Timeout: 3 * time.Second}).DialContext, + }, + Timeout: 3 * time.Second, + } req, err := http.NewRequest("GET", strings.TrimRight(cfg.Endpoint, "/")+"/models", nil) if err != nil { return 0 } req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)") @@ -143,7 +165,13 @@ func fetchContextWindow(cfg *Cfg) int { func listModels(cfg *Cfg) (string, error) { t := cfg.Timeout if t < 10 { t = 10 } - client := &http.Client{Timeout: time.Duration(t) * time.Second} + client := &http.Client{ + Transport: &http.Transport{ + Proxy: proxyFromEnv, + DialContext: (&net.Dialer{Timeout: time.Duration(t) * time.Second}).DialContext, + }, + Timeout: time.Duration(t) * time.Second, + } req, err := http.NewRequest("GET", strings.TrimRight(cfg.Endpoint, "/")+"/models", nil) if err != nil { return "", err } req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)") diff --git a/main_test.go b/main_test.go index 7ed95bc..8218f27 100644 --- a/main_test.go +++ b/main_test.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "fmt" + "net" "net/http" "net/http/httptest" "os" @@ -2199,3 +2200,82 @@ func TestFilterTextPrintableOnly(t *testing.T) { t.Errorf("filterText = %q, want %q", got, "ok\t\nab") } } + +func TestProxyFromEnv(t *testing.T) { + os.Unsetenv("SOCKS_PROXY") + os.Unsetenv("socks_proxy") + req, _ := http.NewRequest("GET", "http://example.com", nil) + _, err := proxyFromEnv(req) + if err != nil { + t.Fatalf("proxyFromEnv err: %v", err) + } + + os.Setenv("SOCKS_PROXY", "127.0.0.1:1080") + defer os.Unsetenv("SOCKS_PROXY") + u, err := proxyFromEnv(req) + if err != nil { + t.Fatalf("proxyFromEnv with 127.0.0.1:1080: %v", err) + } + if u == nil || u.Scheme != "socks5" || u.Host != "127.0.0.1:1080" { + t.Fatalf("unexpected url: %v", u) + } + + os.Setenv("SOCKS_PROXY", "socks5://localhost:9050") + u, err = proxyFromEnv(req) + if err != nil { + t.Fatalf("proxyFromEnv with socks5://: %v", err) + } + if u == nil || u.Scheme != "socks5" || u.Host != "localhost:9050" { + t.Fatalf("unexpected url: %v", u) + } + + os.Unsetenv("SOCKS_PROXY") + os.Setenv("socks_proxy", "socks5h://user:pass@127.0.0.1:1080") + defer os.Unsetenv("socks_proxy") + u, err = proxyFromEnv(req) + if err != nil { + t.Fatalf("proxyFromEnv with socks_proxy: %v", err) + } + if u == nil || u.Scheme != "socks5h" || u.User.Username() != "user" { + t.Fatalf("unexpected url: %v", u) + } +} + +func TestSOCKS5ProxySupport(t *testing.T) { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer l.Close() + + handshakeDone := make(chan bool, 1) + go func() { + conn, err := l.Accept() + if err != nil { + return + } + defer conn.Close() + buf := make([]byte, 256) + n, err := conn.Read(buf) + if err == nil && n >= 2 && buf[0] == 0x05 { + conn.Write([]byte{0x05, 0x00}) + handshakeDone <- true + } + }() + + os.Setenv("SOCKS_PROXY", "socks5://"+l.Addr().String()) + defer os.Unsetenv("SOCKS_PROXY") + + tr := &http.Transport{ + Proxy: proxyFromEnv, + DialContext: (&net.Dialer{Timeout: 1 * time.Second}).DialContext, + } + client := &http.Client{Transport: tr, Timeout: 1 * time.Second} + client.Get("http://example.com/test") + + select { + case <-handshakeDone: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for SOCKS5 handshake through proxy") + } +}