From efe9cbc066d8c3566b6f812b442c18229bc14b94 Mon Sep 17 00:00:00 2001 From: Alexei Shevchenko Date: Mon, 27 Nov 2023 02:35:29 +0300 Subject: [PATCH] Feature/docker (#60) * + in-docker mode --- .dockerignore | 5 ++ .github/workflows/dockerhub.yml | 39 +++++++++++ .golangci.yml | 4 +- .goreleaser.yml | 3 +- README.md | 3 +- cmd/decompose/main.go | 53 +++++++++++++- docker/Dockerfile | 22 ++++++ docker/README.md | 15 ++++ go.mod | 3 +- go.sum | 12 +++- internal/client/defaults.go | 119 +++++++++++++++++++++++++++----- internal/client/docker.go | 32 +++++---- internal/client/docker_test.go | 75 +++++++++++++++++++- internal/graph/build_test.go | 12 +++- internal/graph/config.go | 6 +- internal/graph/load_test.go | 12 +++- 16 files changed, 366 insertions(+), 49 deletions(-) create mode 100644 .dockerignore create mode 100644 .github/workflows/dockerhub.yml create mode 100644 docker/Dockerfile create mode 100644 docker/README.md diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..cad9db4 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,5 @@ +*.out +*.md +./bin +./dist +./examples diff --git a/.github/workflows/dockerhub.yml b/.github/workflows/dockerhub.yml new file mode 100644 index 0000000..3648f59 --- /dev/null +++ b/.github/workflows/dockerhub.yml @@ -0,0 +1,39 @@ +name: Publish Docker image + +on: + release: + types: [published] + +jobs: + update_registry: + name: Build and push Docker image + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ github.actor }} + password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }} + - name: Check out the repo + uses: actions/checkout@v4 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + - name: Build and push + uses: docker/build-push-action@v5 + with: + context: . + file: ./docker/Dockerfile + platforms: linux/amd64 + tags: | + ${{ env.GITHUB_REPOSITORY }}:${{ env.GITHUB_REF_NAME }} + ${{ env.GITHUB_REPOSITORY }}:latest + push: true + - name: Update Docker Hub Readme for dolt image + uses: peter-evans/dockerhub-description@v3 + with: + username: ${{ github.actor }} + password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }} + repository: ${{ env.GITHUB_REPOSITORY }} + readme-filepath: ./docker/README.md diff --git a/.golangci.yml b/.golangci.yml index 2dcf95a..9023d4d 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -19,6 +19,7 @@ linters: - nosnakecase - exhaustruct - inamedparam + - exhaustive - interfacer - varnamelen - scopelint @@ -60,9 +61,6 @@ issues: text: "G204" # G204: Subprocess launched with a potential tainted input linters: - gosec - - path: internal/client/docker\.go - linters: - - exhaustive # None mode, impossible - path: internal/builder/json\.go linters: - errchkjson diff --git a/.goreleaser.yml b/.goreleaser.yml index 0ebe80b..f8b95d3 100644 --- a/.goreleaser.yml +++ b/.goreleaser.yml @@ -39,7 +39,8 @@ archives: {{- .Os }}_ {{- if eq .Arch "amd64" }}x86_64 {{- else if eq .Arch "386" }}i386 - {{- else }}{{ .Arch }}{{ end }} + {{- else }}{{ .Arch }}_ + {{- .Tag }}{{ end }} format_overrides: - goos: windows format: zip diff --git a/README.md b/README.md index 1450db4..341918d 100644 --- a/README.md +++ b/README.md @@ -73,7 +73,7 @@ possible flags with default values: -cluster string json file with clusterization rules, or auto: for auto-clustering, similarity is float in (0.0, 1.0] range -follow string - follow only this container by name + follow only this container by name(s), comma-separated or from @file -format string output format: json, dot, yaml, stat, tree or sdsl for structurizr dsl (default "json") -full @@ -105,6 +105,7 @@ possible flags with default values: - `DOCKER_HOST` - connection uri - `DOCKER_CERT_PATH` - directory path containing key.pem, cert.pm and ca.pem - `DOCKER_TLS_VERIFY` - enable client TLS verification +- `IN_DOCKER_PROC_ROOT` - for in-docker scenario - root for host-mounted /proc ## json stream format diff --git a/cmd/decompose/main.go b/cmd/decompose/main.go index f23de97..8bf1fb4 100644 --- a/cmd/decompose/main.go +++ b/cmd/decompose/main.go @@ -3,6 +3,7 @@ package main import ( + "bufio" "errors" "flag" "fmt" @@ -14,6 +15,8 @@ import ( "strconv" "strings" + "github.com/s0rg/set" + "github.com/s0rg/decompose/internal/builder" "github.com/s0rg/decompose/internal/client" "github.com/s0rg/decompose/internal/cluster" @@ -79,11 +82,10 @@ func setupFlags() { flag.BoolVar(&fLocal, "local", false, "skip external hosts") flag.BoolVar(&fFull, "full", false, "extract full process info: (cmd, args, env) and volumes info") flag.BoolVar(&fNoLoops, "no-loops", false, "remove connection loops (node to itself) from output") - flag.StringVar(&fOut, "out", defaultOutput, "output: filename or \"-\" for stdout") flag.StringVar(&fMeta, "meta", "", "json file with metadata for enrichment") flag.StringVar(&fProto, "proto", defaultProto, "protocol to scan: tcp, udp or all") - flag.StringVar(&fFollow, "follow", "", "follow only this container by name") + flag.StringVar(&fFollow, "follow", "", "follow only this container by name(s), comma-separated or from @file") flag.StringVar( &fCluster, "cluster", @@ -194,6 +196,51 @@ func makeClusterizer( return rv, nil } +func loadFile( + s set.Unordered[string], + v string, +) (err error) { + return feed(v, func(r io.Reader) (err error) { + sc := bufio.NewScanner(r) + + for sc.Scan() { + s.Add(sc.Text()) + } + + if err = sc.Err(); err != nil { + return fmt.Errorf("read: %w", err) + } + + return nil + }) +} + +func loadSet(v string) (rv set.Unordered[string]) { + rv = make(set.Unordered[string]) + + if v == "" { + return + } + + const ( + doggy = "@" + comma = "," + ) + + switch { + case strings.HasPrefix(v, doggy): + if err := loadFile(rv, v[1:]); err != nil { + log.Println("follow:", err) + } + case strings.Contains(v, comma): + set.Load(rv, strings.Split(v, comma)...) + default: + rv.Add(v) + } + + return rv +} + func prepareConfig() ( cfg *graph.Config, nwr graph.NamedWriter, @@ -247,7 +294,7 @@ func prepareConfig() ( Builder: bildr, Meta: meta, Proto: proto, - Follow: fFollow, + Follow: loadSet(fFollow), OnlyLocal: fLocal, FullInfo: fFull, NoLoops: fNoLoops, diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..bca46eb --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,22 @@ +FROM golang:1.21 AS builder + +ADD .. /go/src/github.com/s0rg/decompose +WORKDIR /go/src/github.com/s0rg/decompose + +RUN make build + +FROM scratch + +ARG BUILD_DATE +ARG BUILD_REV + +COPY --from=builder /go/src/github.com/s0rg/decompose/bin/decompose /decompose + +ENTRYPOINT ["/decompose"] + +LABEL org.opencontainers.image.created="${BUILD_DATE}" \ + org.opencontainers.image.revision="${BUILD_REV}" \ + org.opencontainers.image.title="decompose" \ + org.opencontainers.image.authors="s0rg" \ + org.opencontainers.image.vendor="s0rg" \ + org.opencontainers.image.source="https://github.com/s0rg/decompose" diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 0000000..82c31ae --- /dev/null +++ b/docker/README.md @@ -0,0 +1,15 @@ +# decompose + +Reverse-engineering tool for docker environments. + +# how to run + +``` +docker run \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -v /:/rootfs:ro \ + -e IN_DOCKER_PROC_ROOT=/rootfs \ + s0rg/decompose:latest -format stat +``` + +[more options and documentaion](https://github.com/s0rg/decompose) diff --git a/go.mod b/go.mod index 210facc..4fe3dc9 100644 --- a/go.mod +++ b/go.mod @@ -3,9 +3,10 @@ module github.com/s0rg/decompose go 1.21.4 require ( - github.com/antonmedv/expr v1.15.4 + github.com/antonmedv/expr v1.15.5 github.com/docker/docker v24.0.7+incompatible github.com/emicklei/dot v1.6.0 + github.com/otterize/go-procnet v0.1.1 github.com/s0rg/set v1.2.0 github.com/s0rg/trie v1.3.0 gopkg.in/yaml.v3 v3.0.1 diff --git a/go.sum b/go.sum index 6a75bac..cfda051 100644 --- a/go.sum +++ b/go.sum @@ -2,9 +2,10 @@ github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOEl github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E= github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow= github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM= -github.com/antonmedv/expr v1.15.4 h1:CrNads8WDnDVJNWt/FeUINBO+vDNjurEwT7SoQN132o= -github.com/antonmedv/expr v1.15.4/go.mod h1:0E/6TxnOlRNp81GMzX9QfDPAmHo2Phg00y4JUv1ihsE= +github.com/antonmedv/expr v1.15.5 h1:y0Iz3cEwmpRz5/r3w4qQR0MfIqJGdGM1zbhD/v0G5Vg= +github.com/antonmedv/expr v1.15.5/go.mod h1:0E/6TxnOlRNp81GMzX9QfDPAmHo2Phg00y4JUv1ihsE= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/distribution/reference v0.5.0 h1:/FUIFXtfc/x2gpa5/VGfiGLuOIdYa1t65IKK2OFGvA0= @@ -41,6 +42,8 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/image-spec v1.0.2 h1:9yCKha/T5XdGtO0q9Q9a6T5NUCsTn/DrBg0D7ufOcFM= github.com/opencontainers/image-spec v1.0.2/go.mod h1:BtxoFyWECRxE4U/7sNtV5W15zMzWCbyJoFRP3s7yZA0= +github.com/otterize/go-procnet v0.1.1 h1:5vRwX35VrsWcy2uP05sA4PmwpRoAu2L4vMJou4og8Kk= +github.com/otterize/go-procnet v0.1.1/go.mod h1:WEm282HzrSVBZg6DX2fNB4dpVHBPTCjzHWvqOfauV+Q= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= @@ -51,6 +54,10 @@ github.com/s0rg/set v1.2.0 h1:53b207YMktNQJXYei/oHuTR5oOO2e9+eieZOncYsh9g= github.com/s0rg/set v1.2.0/go.mod h1:xz3nDbjF4nyMLvAHvmE7rigXpNrKKTsi6iANznIB1/4= github.com/s0rg/trie v1.3.0 h1:VeMjAJdVvAMt06QlrLJs9B7tplwyxGtCPPZHfvW4Duo= github.com/s0rg/trie v1.3.0/go.mod h1:P+hJUWvPu/imKrsdzOrVswr8Mme6GgFtZfBKojYYkfk= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= @@ -97,6 +104,7 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gotest.tools/v3 v3.5.0 h1:Ljk6PdHdOhAb5aDMWXjDLMMhph+BpztA4v1QdqEW2eY= diff --git a/internal/client/defaults.go b/internal/client/defaults.go index 3443f62..ee80f0d 100644 --- a/internal/client/defaults.go +++ b/internal/client/defaults.go @@ -5,16 +5,35 @@ package client import ( "context" "fmt" - "io" - "os/exec" + "net" + "os" + "path/filepath" "strconv" "time" "github.com/docker/docker/client" + "github.com/otterize/go-procnet/procnet" + "github.com/s0rg/decompose/internal/graph" ) -const pingTimeout = time.Second +const ( + pingTimeout = time.Second + procENV = "IN_DOCKER_PROC_ROOT" + procDefault = "/proc" + procNET = "net" + procTCP4 = "tcp" + procTCP6 = "tcp6" + procUDP4 = "udp" + procUDP6 = "udp6" +) + +var procROOT = procDefault + +type ( + connCB = func(localIP, remoteIP net.IP, localPort, remotePort uint16, kind string) + sockCB func(s procnet.SockTabEntry) +) func Default() (rv DockerClient, err error) { var dc *client.Client @@ -34,33 +53,97 @@ func Default() (rv DockerClient, err error) { return nil, fmt.Errorf("ping: %w", err) } + if root := os.Getenv(procENV); root != "" { + procROOT = filepath.Join(root, procDefault) + } + return dc, nil } +func scanTCP(pid string, onconn connCB) (err error) { + onsock := func(s procnet.SockTabEntry) { + switch s.State { + case procnet.Listen, procnet.Established: + onconn( + s.LocalAddr.IP, + s.RemoteAddr.IP, + s.LocalAddr.Port, + s.RemoteAddr.Port, + procTCP4, + ) + default: + } + } + + pathTCP4 := filepath.Join(procROOT, pid, procNET, procTCP4) + if err = scan(pathTCP4, onsock); err != nil { + return fmt.Errorf("tcp4: %w", err) + } + + pathTCP6 := filepath.Join(procROOT, pid, procNET, procTCP6) + if err = scan(pathTCP6, onsock); err != nil { + return fmt.Errorf("tcp6: %w", err) + } + + return nil +} + +func scanUDP(pid string, onconn connCB) (err error) { + onsock := func(s procnet.SockTabEntry) { + onconn( + s.LocalAddr.IP, + s.RemoteAddr.IP, + s.LocalAddr.Port, + s.RemoteAddr.Port, + procUDP4, + ) + } + + pathUDP4 := filepath.Join(procROOT, pid, procNET, procUDP4) + if err = scan(pathUDP4, onsock); err != nil { + return fmt.Errorf("udp4: %w", err) + } + + pathUDP6 := filepath.Join(procROOT, pid, procNET, procUDP6) + if err = scan(pathUDP6, onsock); err != nil { + return fmt.Errorf("udp6: %w", err) + } + + return nil +} + +func scan(path string, onsock sockCB) (err error) { + socks, err := procnet.SocksFromPath(path) + if err != nil { + return fmt.Errorf("path: %w", err) + } + + for _, s := range socks { + onsock(s) + } + + return nil +} + func Nsenter( - ctx context.Context, pid int, proto graph.NetProto, - parse func(io.Reader) error, + onconn connCB, ) ( err error, ) { - arg := append([]string{"-t", strconv.Itoa(pid), "-n"}, netstat(proto)...) - cmd := exec.CommandContext(ctx, nsenterCmd, arg...) - - pipe, err := cmd.StdoutPipe() - if err != nil { - return fmt.Errorf("pipe: %w", err) - } - - defer pipe.Close() + spid := strconv.Itoa(pid) - if err = cmd.Start(); err != nil { - return fmt.Errorf("exec: %w", err) + if proto == graph.ALL || proto == graph.TCP { + if err = scanTCP(spid, onconn); err != nil { + return fmt.Errorf("scan: %w", err) + } } - if err = parse(pipe); err != nil { - return fmt.Errorf("parse: %w", err) + if proto == graph.ALL || proto == graph.UDP { + if err = scanUDP(spid, onconn); err != nil { + return fmt.Errorf("scan: %w", err) + } } return nil diff --git a/internal/client/docker.go b/internal/client/docker.go index 213269a..6c815a5 100644 --- a/internal/client/docker.go +++ b/internal/client/docker.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "net" "slices" "strings" @@ -18,14 +19,13 @@ import ( const ( stateRunning = "running" - nsenterCmd = "nsenter" netstatCmd = "netstat" ) var ErrModeNone = errors.New("mode not set") type createClient func() (DockerClient, error) -type nsEnter func(context.Context, int, graph.NetProto, func(io.Reader) error) error +type nsEnter func(int, graph.NetProto, func(localIP, remoteIP net.IP, localPort, remotePort uint16, kind string)) error type DockerClient interface { ContainerList(context.Context, types.ContainerListOptions) ([]types.Container, error) @@ -143,17 +143,15 @@ func (d *Docker) connections( cid string, proto graph.NetProto, ) (rv []*graph.Connection, err error) { - parse := func(r io.Reader) (err error) { - if rv, err = graph.ParseNetstat(r); err != nil { - return fmt.Errorf("parse: %w", err) - } - - return nil - } - switch d.opt.Mode { case InContainer: - err = d.connectionsContainer(ctx, cid, proto, parse) + err = d.connectionsContainer(ctx, cid, proto, func(r io.Reader) (err error) { + if rv, err = graph.ParseNetstat(r); err != nil { + return fmt.Errorf("parse: %w", err) + } + + return nil + }) case LinuxNsenter: if pid == 0 { info, ierr := d.cli.ContainerInspect(ctx, cid) @@ -164,7 +162,17 @@ func (d *Docker) connections( pid = info.State.Pid } - err = d.opt.Nsenter(ctx, pid, proto, parse) + err = d.opt.Nsenter(pid, proto, func(locIP, remIP net.IP, locPort, remPort uint16, kind string) { + sk, _ := graph.ParseNetProto(kind) + + rv = append(rv, &graph.Connection{ + LocalIP: locIP, + RemoteIP: remIP, + LocalPort: locPort, + RemotePort: remPort, + Proto: sk, + }) + }) } if err != nil { diff --git a/internal/client/docker_test.go b/internal/client/docker_test.go index d7cc93a..8c86cf2 100644 --- a/internal/client/docker_test.go +++ b/internal/client/docker_test.go @@ -5,7 +5,7 @@ import ( "bytes" "context" "errors" - "io" + "net" "testing" "github.com/docker/docker/api/types" @@ -774,7 +774,11 @@ func TestDockerClientNsEnterConnectionsError(t *testing.T) { return rv } - failEnter := func(_ context.Context, _ int, _ graph.NetProto, _ func(io.Reader) error) error { + failEnter := func(_ int, _ graph.NetProto, _ func( + locIP, remIP net.IP, + locPort, remPort uint16, + kind string, + )) error { return testErr } @@ -801,3 +805,70 @@ func TestDockerClientNsEnterConnectionsError(t *testing.T) { t.Fail() } } + +func TestDockerClientNsEnterOk(t *testing.T) { + t.Parallel() + + cm := &clientMock{} + + cm.OnList = func() (rv []types.Container) { + return []types.Container{ + { + ID: "1", + Names: []string{"test"}, + Image: "test-image", + State: "running", + NetworkSettings: &types.SummaryNetworkSettings{ + Networks: map[string]*network.EndpointSettings{ + "test-net": { + EndpointID: "1", + IPAddress: "1.1.1.1", + }, + "empty-id": { + IPAddress: "1.1.1.2", + }, + }, + }, + }, + } + } + + cm.OnInspect = func() (rv types.ContainerJSON) { + rv.ContainerJSONBase = &types.ContainerJSONBase{} + rv.State = &types.ContainerState{Pid: 1} + + return rv + } + + testEnter := func(_ int, _ graph.NetProto, fn func( + locIP, remIP net.IP, + locPort, remPort uint16, + kind string, + )) error { + fn(net.IP{}, net.IP{}, 0, 0, "tcp") + + return nil + } + + cli, err := client.NewDocker( + client.WithClientCreator(func() (client.DockerClient, error) { + return cm, nil + }), + client.WithMode(client.LinuxNsenter), + client.WithNsEnter(testEnter), + ) + if err != nil { + t.Fatal("client:", err) + } + + _, err = cli.Containers( + context.Background(), + graph.ALL, + false, + nil, + voidProgress, + ) + if err != nil { + t.Fatal("containers:", err) + } +} diff --git a/internal/graph/build_test.go b/internal/graph/build_test.go index c405f1d..e632a63 100644 --- a/internal/graph/build_test.go +++ b/internal/graph/build_test.go @@ -6,6 +6,8 @@ import ( "net" "testing" + "github.com/s0rg/set" + "github.com/s0rg/decompose/internal/graph" "github.com/s0rg/decompose/internal/node" ) @@ -180,14 +182,17 @@ func TestBuildSimple(t *testing.T) { func TestBuildFollow(t *testing.T) { t.Parallel() + flw := make(set.Unordered[string]) + flw.Add("1") + cli := testClientWithEnv() bld := &testBuilder{} ext := &testEnricher{} cfg := &graph.Config{ Builder: bld, + Follow: flw, Meta: ext, Proto: graph.ALL, - Follow: "1", } if err := graph.Build(cfg, cli); err != nil { @@ -224,10 +229,13 @@ func TestBuildLocal(t *testing.T) { func TestBuildNoNodes(t *testing.T) { t.Parallel() + flw := make(set.Unordered[string]) + flw.Add("4") + cli := testClientWithEnv() cfg := &graph.Config{ + Follow: flw, Proto: graph.ALL, - Follow: "4", } if err := graph.Build(cfg, cli); err == nil { diff --git a/internal/graph/config.go b/internal/graph/config.go index bab7cd7..6d8719c 100644 --- a/internal/graph/config.go +++ b/internal/graph/config.go @@ -1,9 +1,11 @@ package graph +import "github.com/s0rg/set" + type Config struct { Builder Builder Meta Enricher - Follow string + Follow set.Unordered[string] SkipEnv []string Proto NetProto OnlyLocal bool @@ -12,7 +14,7 @@ type Config struct { } func (c *Config) MatchName(v string) (yes bool) { - return c.Follow == "" || v == c.Follow + return c.Follow.Len() == 0 || c.Follow.Has(v) } func (c *Config) MatchProto(v string) (yes bool) { diff --git a/internal/graph/load_test.go b/internal/graph/load_test.go index ccbc14c..12c7bc2 100644 --- a/internal/graph/load_test.go +++ b/internal/graph/load_test.go @@ -5,6 +5,8 @@ import ( "errors" "testing" + "github.com/s0rg/set" + "github.com/s0rg/decompose/internal/graph" ) @@ -256,10 +258,13 @@ func TestLoaderEdgesFollowNone(t *testing.T) { bldr := &testBuilder{} ext := &testEnricher{} + flw := make(set.Unordered[string]) + flw.Add("foo") + cfg := &graph.Config{ Builder: bldr, + Follow: flw, Meta: ext, - Follow: "foo", Proto: graph.ALL, } @@ -311,10 +316,13 @@ func TestLoaderEdgesFollowOne(t *testing.T) { bldr := &testBuilder{} ext := &testEnricher{} + flw := make(set.Unordered[string]) + flw.Add("test3") + cfg := &graph.Config{ Builder: bldr, + Follow: flw, Meta: ext, - Follow: "test3", Proto: graph.UDP, }