feat(datacap): implement datacap tracking for connections

cmd/cmd_run.go

@@ -15,8 +15,13 @@ import (
 	"github.com/sagernet/sing-box/option"
 	E "github.com/sagernet/sing/common/exceptions"
 	"github.com/sagernet/sing/common/json"
+	"github.com/sagernet/sing/service"
 	"github.com/spf13/cobra"
 	"gopkg.in/ini.v1"
+
+	"github.com/getlantern/lantern-box/adapter"
+	"github.com/getlantern/lantern-box/tracker/clientcontext"
+	"github.com/getlantern/lantern-box/tracker/datacap"
 )
 
 func init() {
@@ -25,6 +30,7 @@ func init() {
 	runCmd.Flags().String("geo-city-url", "https://lanterngeo.lantern.io/GeoLite2-City.mmdb.tar.gz", "URL for downloading GeoLite2-City database")
 	runCmd.Flags().String("city-database-name", "GeoLite2-City.mmdb", "Filename for storing GeoLite2-City database")
 	runCmd.Flags().String("telemetry-endpoint", "telemetry.iantem.io:443", "Telemetry endpoint for OpenTelemetry exporter")
+	runCmd.Flags().String("datacap-url", "", "Datacap server URL")
 }
 
 var runCmd = &cobra.Command{
@@ -35,7 +41,11 @@ var runCmd = &cobra.Command{
 		if err != nil {
 			return fmt.Errorf("get config flag: %w", err)
 		}
-		return run(path)
+		datacapURL, err := cmd.Flags().GetString("datacap-url")
+		if err != nil {
+			return fmt.Errorf("get datacap-url flag: %w", err)
+		}
+		return run(path, datacapURL)
 	},
 }
 
@@ -64,7 +74,7 @@ func readConfig(path string) (option.Options, error) {
 	return options, nil
 }
 
-func create(configPath string) (*box.Box, context.CancelFunc, error) {
+func create(configPath string, datacapURL string) (*box.Box, context.CancelFunc, error) {
 	options, err := readConfig(configPath)
 	if err != nil {
 		return nil, nil, fmt.Errorf("read config: %w", err)
@@ -79,6 +89,29 @@ func create(configPath string) (*box.Box, context.CancelFunc, error) {
 		return nil, nil, fmt.Errorf("create service: %w", err)
 	}
 
+	if datacapURL != "" {
+		// Add datacap tracker
+		clientCtxMgr := clientcontext.NewManager(clientcontext.MatchBounds{
+			Inbound:  []string{""},
+			Outbound: []string{""},
+		}, log.NewNOPFactory().NewLogger("tracker"))
+		instance.Router().AppendTracker(clientCtxMgr)
+		service.MustRegister[adapter.ClientContextManager](ctx, clientCtxMgr)
+
+		datacapTracker, err := datacap.NewDatacapTracker(
+			datacap.Options{
+				URL: datacapURL,
+			},
+			log.NewNOPFactory().NewLogger("datacap-tracker"),
+		)
+		if err != nil {
+			return nil, nil, fmt.Errorf("create datacap tracker: %w", err)
+		}
+		clientCtxMgr.AppendTracker(datacapTracker)
+	} else {
+		log.Warn("Datacap URL not provided, datacap tracking disabled")
+	}
+
 	osSignals := make(chan os.Signal, 1)
 	signal.Notify(osSignals, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
 	defer func() {
@@ -112,13 +145,13 @@ func closeMonitor(ctx context.Context) {
 	log.Fatal("sing-box did not close!")
 }
 
-func run(configPath string) error {
+func run(configPath string, datacapURL string) error {
 	log.Info("build info: version ", version, ", commit ", commit)
 	osSignals := make(chan os.Signal, 1)
 	signal.Notify(osSignals, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP)
 	defer signal.Stop(osSignals)
 	for {
-		instance, cancel, err := create(configPath)
+		instance, cancel, err := create(configPath, datacapURL)
 		if err != nil {
 			return err
 		}

tracker/datacap/client.go

@@ -0,0 +1,150 @@
+package datacap
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+)
+
+// Client handles communication with the datacap sidecar service.
+type Client struct {
+	httpClient *http.Client
+	baseURL    string
+}
+
+// ClientConfig holds configuration for the datacap client.
+type ClientConfig struct {
+	BaseURL            string
+	Timeout            time.Duration
+	InsecureSkipVerify bool
+}
+
+// NewClient creates a new datacap client.
+// The baseURL can be overridden by the DATACAP_URL environment variable.
+// Supports both HTTP and HTTPS. For HTTPS, uses system's trusted certificates by default.
+func NewClient(baseURL string, timeout time.Duration) *Client {
+	return NewClientWithConfig(ClientConfig{
+		BaseURL:            baseURL,
+		Timeout:            timeout,
+		InsecureSkipVerify: false,
+	})
+}
+
+// NewClientWithConfig creates a new datacap client with advanced configuration.
+func NewClientWithConfig(config ClientConfig) *Client {
+	// Check for environment variable override
+	if envURL := os.Getenv("DATACAP_URL"); envURL != "" {
+		config.BaseURL = envURL
+	}
+
+	// Ensure HTTPS if not explicitly HTTP
+	if config.BaseURL != "" && !strings.HasPrefix(config.BaseURL, "http://") && !strings.HasPrefix(config.BaseURL, "https://") {
+		config.BaseURL = "https://" + config.BaseURL
+	}
+
+	// Create HTTP client with TLS configuration
+	transport := &http.Transport{
+		TLSClientConfig: &tls.Config{
+			InsecureSkipVerify: config.InsecureSkipVerify,
+		},
+		MaxIdleConns:        100,
+		MaxIdleConnsPerHost: 10,
+		IdleConnTimeout:     90 * time.Second,
+	}
+
+	return &Client{
+		httpClient: &http.Client{
+			Timeout:   config.Timeout,
+			Transport: transport,
+		},
+		baseURL: config.BaseURL,
+	}
+}
+
+// DataCapStatus represents the response from the GET /data-cap/{deviceId} endpoint.
+type DataCapStatus struct {
+	Throttle       bool  `json:"throttle"`
+	RemainingBytes int64 `json:"remainingBytes"`
+	CapLimit       int64 `json:"capLimit"`
+	ExpiryTime     int64 `json:"expiryTime"`
+}
+
+// DataCapReport represents the request body for POST /data-cap/ endpoint.
+type DataCapReport struct {
+	DeviceID    string `json:"deviceId"`
+	CountryCode string `json:"countryCode"`
+	Platform    string `json:"platform"`
+	BytesUsed   int64  `json:"bytesUsed"`
+}
+
+func (c *Client) GetDataCapStatus(ctx context.Context, deviceID string) (*DataCapStatus, error) {
+	url := fmt.Sprintf("%s/data-cap/%s", c.baseURL, deviceID)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query datacap status: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("datacap status request failed with status %d: %s", resp.StatusCode, string(body))
+	}
+
+	var status DataCapStatus
+	if err := json.NewDecoder(resp.Body).Decode(&status); err != nil {
+		return nil, fmt.Errorf("failed to decode datacap status: %w", err)
+	}
+
+	return &status, nil
+}
+
+// ReportDataCapConsumption sends data consumption report to the sidecar.
+// Endpoint: POST /data-cap/
+// This tracks usage and returns updated cap status.
+func (c *Client) ReportDataCapConsumption(ctx context.Context, report *DataCapReport) (*DataCapStatus, error) {
+	url := fmt.Sprintf("%s/data-cap/", c.baseURL)
+
+	jsonData, err := json.Marshal(report)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal report: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Accept", "application/json")
+
+	resp, err := c.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to report datacap consumption: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
+		body, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("datacap report request failed with status %d: %s", resp.StatusCode, string(body))
+	}
+
+	var status DataCapStatus
+	if err := json.NewDecoder(resp.Body).Decode(&status); err != nil {
+		return nil, fmt.Errorf("failed to decode datacap status: %w", err)
+	}
+
+	return &status, nil
+}