diff options
| -rw-r--r-- | xray/api.go | 170 |
1 files changed, 90 insertions, 80 deletions
diff --git a/xray/api.go b/xray/api.go index ac384ed4..2dce4dcb 100644 --- a/xray/api.go +++ b/xray/api.go @@ -31,24 +31,27 @@ type XrayAPI struct { isConnected bool } -func (x *XrayAPI) Init(apiPort int) (err error) { - if apiPort == 0 { - return common.NewError("xray api port wrong:", apiPort) +func (x *XrayAPI) Init(apiPort int) error { + if apiPort <= 0 { + return fmt.Errorf("invalid Xray API port: %d", apiPort) } - conn, err := grpc.NewClient(fmt.Sprintf("127.0.0.1:%v", apiPort), grpc.WithTransportCredentials(insecure.NewCredentials())) + + addr := fmt.Sprintf("127.0.0.1:%d", apiPort) + conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { - return err + return fmt.Errorf("failed to connect to Xray API: %w", err) } + x.grpcClient = conn x.isConnected = true - hsClient := command.NewHandlerServiceClient(x.grpcClient) - ssClient := statsService.NewStatsServiceClient(x.grpcClient) + hsClient := command.NewHandlerServiceClient(conn) + ssClient := statsService.NewStatsServiceClient(conn) x.HandlerServiceClient = &hsClient x.StatsServiceClient = &ssClient - return + return nil } func (x *XrayAPI) Close() { @@ -149,94 +152,101 @@ func (x *XrayAPI) AddUser(Protocol string, inboundTag string, user map[string]in return err } -func (x *XrayAPI) RemoveUser(inboundTag string, email string) error { - client := *x.HandlerServiceClient - _, err := client.AlterInbound(context.Background(), &command.AlterInboundRequest{ - Tag: inboundTag, - Operation: serial.ToTypedMessage(&command.RemoveUserOperation{ - Email: email, - }), - }) - return err +func (x *XrayAPI) RemoveUser(inboundTag, email string) error { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + op := &command.RemoveUserOperation{Email: email} + req := &command.AlterInboundRequest{ + Tag: inboundTag, + Operation: serial.ToTypedMessage(op), + } + + _, err := (*x.HandlerServiceClient).AlterInbound(ctx, req) + if err != nil { + return fmt.Errorf("failed to remove user: %w", err) + } + + return nil } func (x *XrayAPI) GetTraffic(reset bool) ([]*Traffic, []*ClientTraffic, error) { if x.grpcClient == nil { return nil, nil, common.NewError("xray api is not initialized") } - trafficRegex := regexp.MustCompile("(inbound|outbound)>>>([^>]+)>>>traffic>>>(downlink|uplink)") - ClientTrafficRegex := regexp.MustCompile("(user)>>>([^>]+)>>>traffic>>>(downlink|uplink)") - client := *x.StatsServiceClient + trafficRegex := regexp.MustCompile(`(inbound|outbound)>>>([^>]+)>>>traffic>>>(downlink|uplink)`) + clientTrafficRegex := regexp.MustCompile(`user>>>([^>]+)>>>traffic>>>(downlink|uplink)`) + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) defer cancel() - request := &statsService.QueryStatsRequest{ - Reset_: reset, - } - resp, err := client.QueryStats(ctx, request) + + resp, err := (*x.StatsServiceClient).QueryStats(ctx, &statsService.QueryStatsRequest{Reset_: reset}) if err != nil { return nil, nil, err } - tagTrafficMap := map[string]*Traffic{} - emailTrafficMap := map[string]*ClientTraffic{} - clientTraffics := make([]*ClientTraffic, 0) - traffics := make([]*Traffic, 0) + tagTrafficMap := make(map[string]*Traffic) + emailTrafficMap := make(map[string]*ClientTraffic) + for _, stat := range resp.GetStat() { - matchs := trafficRegex.FindStringSubmatch(stat.Name) - if len(matchs) < 3 { - - matchs := ClientTrafficRegex.FindStringSubmatch(stat.Name) - if len(matchs) < 3 { - continue - } else { - - isUser := matchs[1] == "user" - email := matchs[2] - isDown := matchs[3] == "downlink" - if !isUser { - continue - } - traffic, ok := emailTrafficMap[email] - if !ok { - traffic = &ClientTraffic{ - Email: email, - } - emailTrafficMap[email] = traffic - clientTraffics = append(clientTraffics, traffic) - } - if isDown { - traffic.Down = stat.Value - } else { - traffic.Up = stat.Value - } - - } - continue + if matches := trafficRegex.FindStringSubmatch(stat.Name); len(matches) == 4 { + processTraffic(matches, stat.Value, tagTrafficMap) + } else if matches := clientTrafficRegex.FindStringSubmatch(stat.Name); len(matches) == 3 { + processClientTraffic(matches, stat.Value, emailTrafficMap) } - isInbound := matchs[1] == "inbound" - isOutbound := matchs[1] == "outbound" - tag := matchs[2] - isDown := matchs[3] == "downlink" - if tag == "api" { - continue - } - traffic, ok := tagTrafficMap[tag] - if !ok { - traffic = &Traffic{ - IsInbound: isInbound, - IsOutbound: isOutbound, - Tag: tag, - } - tagTrafficMap[tag] = traffic - traffics = append(traffics, traffic) - } - if isDown { - traffic.Down = stat.Value - } else { - traffic.Up = stat.Value + } + + return mapToSlice(tagTrafficMap), mapToSlice(emailTrafficMap), nil +} + +func processTraffic(matches []string, value int64, trafficMap map[string]*Traffic) { + isInbound := matches[1] == "inbound" + tag := matches[2] + isDown := matches[3] == "downlink" + + if tag == "api" { + return + } + + traffic, ok := trafficMap[tag] + if !ok { + traffic = &Traffic{ + IsInbound: isInbound, + IsOutbound: !isInbound, + Tag: tag, } + trafficMap[tag] = traffic } - return traffics, clientTraffics, nil + if isDown { + traffic.Down = value + } else { + traffic.Up = value + } +} + +func processClientTraffic(matches []string, value int64, clientTrafficMap map[string]*ClientTraffic) { + email := matches[1] + isDown := matches[2] == "downlink" + + traffic, ok := clientTrafficMap[email] + if !ok { + traffic = &ClientTraffic{Email: email} + clientTrafficMap[email] = traffic + } + + if isDown { + traffic.Down = value + } else { + traffic.Up = value + } +} + +func mapToSlice[T any](m map[string]*T) []*T { + result := make([]*T, 0, len(m)) + for _, v := range m { + result = append(result, v) + } + return result } |