Skip to content

Commit

Permalink
feat: New service name rule
Browse files Browse the repository at this point in the history
  • Loading branch information
layou233 committed Jul 21, 2024
1 parent 76024ce commit 549a9aa
Show file tree
Hide file tree
Showing 6 changed files with 77 additions and 9 deletions.
3 changes: 2 additions & 1 deletion adapter/metadata.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ type Protocol = uint8

type Metadata struct {
ConnectionID string
ServiceName string
SniffedProtocol Protocol
SourceIP netip.Addr
DestinationHostname string
Expand All @@ -30,7 +31,7 @@ type Metadata struct {

func (m *Metadata) GenerateID() {
id := int64(fastrand.Int31())
idColor := fastrand.Intn(len(color.List))
idColor := fastrand.Int31n(int32(len(color.List)))
m.ConnectionID = color.Apply(color.List[idColor], "["+strconv.FormatInt(id, 10)+"]")
}

Expand Down
2 changes: 2 additions & 0 deletions route/rule.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@ func NewRule(logger *log.Logger, config *config.Rule, listMap map[string]set.Str
return NewLogicalAndRule(logger, config, listMap, ruleRegistry)
case "or":
return NewLogicalOrRule(logger, config, listMap, ruleRegistry)
case "ServiceName":
return NewServiceNameRule(config, listMap)
case "SourceIPVersion":
return NewSourceIPVersionRule(config)
case "SourceIP":
Expand Down
4 changes: 2 additions & 2 deletions route/rule_logic.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ func NewLogicalAndRule(logger *log.Logger, newConfig *config.Rule, listMap map[s
return &RuleLogicalAnd{logicRule}, nil
}

func (r RuleLogicalAnd) Match(metadata *adapter.Metadata) (match bool) {
func (r *RuleLogicalAnd) Match(metadata *adapter.Metadata) (match bool) {
match = true
for _, rule := range r.rules {
if !rule.Match(metadata) {
Expand Down Expand Up @@ -83,7 +83,7 @@ func NewLogicalOrRule(logger *log.Logger, newConfig *config.Rule, listMap map[st
return &RuleLogicalOr{logicRule}, nil
}

func (r RuleLogicalOr) Match(metadata *adapter.Metadata) (match bool) {
func (r *RuleLogicalOr) Match(metadata *adapter.Metadata) (match bool) {
for _, rule := range r.rules {
if rule.Match(metadata) {
match = true
Expand Down
4 changes: 2 additions & 2 deletions route/rule_minecraft.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,11 +45,11 @@ func NewMinecraftPlayerNameRule(newConfig *config.Rule, listMap map[string]set.S
}, nil
}

func (r RuleMinecraftPlayerName) Config() *config.Rule {
func (r *RuleMinecraftPlayerName) Config() *config.Rule {
return r.config
}

func (r RuleMinecraftPlayerName) Match(metadata *adapter.Metadata) (match bool) {
func (r *RuleMinecraftPlayerName) Match(metadata *adapter.Metadata) (match bool) {
if metadata.Minecraft != nil {
for _, nameSet := range r.sets {
match = nameSet.Has(metadata.Minecraft.PlayerName)
Expand Down
63 changes: 63 additions & 0 deletions route/rule_service.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
package route

import (
"encoding/json"
"fmt"
"strings"

"github.com/layou233/zbproxy/v3/adapter"
"github.com/layou233/zbproxy/v3/common/jsonx"
"github.com/layou233/zbproxy/v3/common/set"
"github.com/layou233/zbproxy/v3/config"
)

type RuleServiceName struct {
sets []set.StringSet
config *config.Rule
}

var _ Rule = (*RuleServiceName)(nil)

func NewServiceNameRule(newConfig *config.Rule, listMap map[string]set.StringSet) (Rule, error) {
var serviceList jsonx.Listable[string]
err := json.Unmarshal(newConfig.Parameter, &serviceList)
if err != nil {
return nil, fmt.Errorf("bad service name list %v: %w", newConfig.Parameter, err)
}
sets := []set.StringSet{
{}, // new set for individual names
}
for _, i := range serviceList {
if strings.HasPrefix(i, parameterListPrefix) {
i = strings.TrimPrefix(i, parameterListPrefix)
nameSet, found := listMap[i]
if !found {
return nil, fmt.Errorf("list [%v] is not found", i)
}
sets = append(sets, nameSet)
} else {
sets[0].Add(i)
}
}
return &RuleServiceName{
sets: sets,
config: newConfig,
}, nil
}

func (r RuleServiceName) Config() *config.Rule {
return r.config
}

func (r RuleServiceName) Match(metadata *adapter.Metadata) (match bool) {
for _, nameSet := range r.sets {
match = nameSet.Has(metadata.ServiceName)
if match {
break
}
}
if r.config.Invert {
match = !match
}
return
}
10 changes: 6 additions & 4 deletions service/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,13 @@ func (s *Service) listenLoop() {
s.logger.Warn().Str("service", s.config.Name).Str("ip", ipString).Msg("Rejected by access control")
continue
}
metadata := &adapter.Metadata{}
metadata := &adapter.Metadata{
ServiceName: s.config.Name,
DestinationHostname: s.config.TargetAddress,
DestinationPort: s.config.TargetPort,
SourceIP: common.MustOK(netip.AddrFromSlice(ip)).Unmap(),
}
metadata.GenerateID()
metadata.DestinationHostname = s.config.TargetAddress
metadata.DestinationPort = s.config.TargetPort
metadata.SourceIP = common.MustOK(netip.AddrFromSlice(ip)).Unmap()
s.logger.Info().Str("id", metadata.ConnectionID).Str("service", s.config.Name).
Str("ip", ipString).Msg("New inbound connection")
if s.legacyOutbound != nil {
Expand Down

0 comments on commit 549a9aa

Please sign in to comment.