isolate session network checks

common/session.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/http"
 	"strings"
 	"sync"
 	"time"
@@ -119,6 +120,43 @@ func (s *ScanSession) DialTCP(ctx context.Context, network, address string, time
 	return conn, nil
 }
 
+// HTTPDo executes an HTTP request with the session's packet limits and counters.
+func (s *ScanSession) HTTPDo(client *http.Client, req *http.Request) (*http.Response, error) {
+	if ok, err := CanSendPacketWith(s.Config, s.State); !ok {
+		s.LogError(fmt.Sprintf("HTTP请求 %s 受限: %s", req.URL.String(), err.Error()))
+		return nil, fmt.Errorf("%s", i18n.Tr("network_rate_limited", err.Error()))
+	}
+
+	resp, err := client.Do(req)
+	if err != nil {
+		s.State.IncrementTCPFailedPacketCount()
+		return nil, err
+	}
+	s.State.IncrementTCPSuccessPacketCount()
+	return resp, nil
+}
+
+// ProxyEnabled reports whether this scan session uses a network proxy.
+func (s *ScanSession) ProxyEnabled() bool {
+	if s == nil || s.Config == nil {
+		return false
+	}
+	return s.Config.Network.Socks5Proxy != "" || s.Config.Network.HTTPProxy != ""
+}
+
+// IsSOCKS5Proxy reports whether this scan session uses SOCKS5.
+func (s *ScanSession) IsSOCKS5Proxy() bool {
+	return s != nil && s.Config != nil && s.Config.Network.Socks5Proxy != ""
+}
+
+// ProxyReliable reports whether the session proxy should be treated as reliable.
+func (s *ScanSession) ProxyReliable() bool {
+	if !s.ProxyEnabled() || !s.IsSOCKS5Proxy() {
+		return true
+	}
+	return proxy.IsProxyReliable()
+}
+
 func (s *ScanSession) getDialer(timeout time.Duration) (proxy.Dialer, error) {
 	if timeout <= 0 {
 		timeout = s.Config.Timeout

common/session_test.go

@@ -1,6 +1,9 @@
 package common
 
 import (
+	"io"
+	"net/http"
+	"strings"
 	"testing"
 	"time"
 )
@@ -65,3 +68,81 @@ func TestScanSessionDialerCacheIsTimeoutAware(t *testing.T) {
 		t.Fatalf("proxy timeout = %v, want %v", got, shortTimeout)
 	}
 }
+
+func TestScanSessionHTTPDoUsesSessionState(t *testing.T) {
+	previousState := GetGlobalState()
+	globalState := NewState()
+	SetGlobalState(globalState)
+	t.Cleanup(func() { SetGlobalState(previousState) })
+
+	sessionState := NewState()
+	session := NewScanSession(NewConfig(), sessionState, &FlagVars{})
+	client := &http.Client{
+		Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: http.StatusNoContent,
+				Body:       io.NopCloser(strings.NewReader("")),
+				Header:     make(http.Header),
+			}, nil
+		}),
+	}
+	req, err := http.NewRequest(http.MethodHead, "http://example.com", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	resp, err := session.HTTPDo(client, req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_ = resp.Body.Close()
+
+	if got := sessionState.GetTCPSuccessPacketCount(); got != 1 {
+		t.Fatalf("session TCP success count = %d, want 1", got)
+	}
+	if got := globalState.GetTCPSuccessPacketCount(); got != 0 {
+		t.Fatalf("global TCP success count = %d, want 0", got)
+	}
+}
+
+func TestScanSessionProxyStateComesFromConfig(t *testing.T) {
+	direct := NewScanSession(NewConfig(), NewState(), &FlagVars{})
+	if direct.ProxyEnabled() {
+		t.Fatal("direct session should not report proxy enabled")
+	}
+	if direct.IsSOCKS5Proxy() {
+		t.Fatal("direct session should not report SOCKS5")
+	}
+	if !direct.ProxyReliable() {
+		t.Fatal("direct session should be reliable")
+	}
+
+	httpCfg := NewConfig()
+	httpCfg.Network.HTTPProxy = "http://127.0.0.1:8080"
+	httpSession := NewScanSession(httpCfg, NewState(), &FlagVars{})
+	if !httpSession.ProxyEnabled() {
+		t.Fatal("HTTP proxy session should report proxy enabled")
+	}
+	if httpSession.IsSOCKS5Proxy() {
+		t.Fatal("HTTP proxy session should not report SOCKS5")
+	}
+	if !httpSession.ProxyReliable() {
+		t.Fatal("HTTP proxy session should be reliable")
+	}
+
+	socksCfg := NewConfig()
+	socksCfg.Network.Socks5Proxy = "127.0.0.1:1080"
+	socksSession := NewScanSession(socksCfg, NewState(), &FlagVars{})
+	if !socksSession.ProxyEnabled() {
+		t.Fatal("SOCKS5 proxy session should report proxy enabled")
+	}
+	if !socksSession.IsSOCKS5Proxy() {
+		t.Fatal("SOCKS5 proxy session should report SOCKS5")
+	}
+}
+
+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return f(req)
+}