instrument: fix duplicated instrumentation statements on already instrumented frameworks

internal/instrument/chiv5.go

@@ -9,6 +9,17 @@ import (
 	"github.com/dave/dst"
 )
 
+func instrumentChiV5(stmt *dst.AssignStmt) []dst.Stmt {
+	if !isChiV5(stmt) {
+		return nil
+	}
+	stmt.Decorations().Start.Prepend(dd_instrumented)
+	return []dst.Stmt{
+		stmt,
+		chiV5Middleware(stmt),
+	}
+}
+
 func isChiV5(stmt *dst.AssignStmt) bool {
 	rhs := stmt.Rhs[0]
 	f, ok := funcIdent(rhs)
@@ -21,7 +32,7 @@ func chiV5Middleware(got *dst.AssignStmt) dst.Stmt {
 		return nil
 	}
 	stmt := useMiddleware(iden.Name, "ChiV5Middleware")
-	wrap(stmt)
+	markAsInstrumented(stmt)
 	return stmt
 }
 

internal/instrument/chiv5_test.go

@@ -0,0 +1,96 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2023-present Datadog, Inc.
+
+package instrument
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/datadog/orchestrion/internal/config"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestChiV5(t *testing.T) {
+	var codeTpl = `package main
+
+import %s
+
+func register() {
+	%s
+}
+`
+	var wantTpl = `package main
+
+import (
+	"github.com/datadog/orchestrion/instrument"
+	%s
+)
+
+func register() {
+	//dd:instrumented
+	%s
+	//dd:startinstrument
+	%s
+	//dd:endinstrument
+}
+`
+
+	tests := []struct {
+		pkg  string
+		stmt string
+		want string
+		tmpl string
+	}{
+		{pkg: `"github.com/go-chi/chi/v5"`, stmt: `r := chi.NewRouter()`, want: `r.Use(instrument.ChiV5Middleware())`, tmpl: wantTpl},
+		{pkg: `chi "github.com/go-chi/chi/v5"`, stmt: `r := chi.NewRouter()`, want: `r.Use(instrument.ChiV5Middleware())`, tmpl: wantTpl},
+		{pkg: `chiv5 "github.com/go-chi/chi/v5"`, stmt: `r := chiv5.NewRouter()`, want: `r.Use(instrument.ChiV5Middleware())`, tmpl: wantTpl},
+	}
+
+	for i, tc := range tests {
+		t.Run(fmt.Sprintf("tc-%d", i), func(t *testing.T) {
+			code := fmt.Sprintf(codeTpl, tc.pkg, tc.stmt)
+			reader, err := InstrumentFile("test", strings.NewReader(code), config.Config{})
+			require.Nil(t, err)
+			got, err := io.ReadAll(reader)
+			require.Nil(t, err)
+			want := fmt.Sprintf(tc.tmpl, tc.pkg, tc.stmt, tc.want)
+			require.Equal(t, want, string(got))
+
+			reader, err = UninstrumentFile("test", strings.NewReader(want), config.Config{})
+			require.Nil(t, err)
+			orig, err := io.ReadAll(reader)
+			require.Nil(t, err)
+			require.Equal(t, code, string(orig))
+		})
+	}
+}
+
+func TestChiV5Duplicates(t *testing.T) {
+	var tpl = `package main
+
+import (
+	"github.com/datadog/orchestrion/instrument"
+	"github.com/go-chi/chi/v5"
+)
+
+func chiV5Server() {
+	//dd:instrumented
+	r := chi.NewRouter()
+	//dd:startinstrument
+	r.Use(instrument.ChiV5Middleware())
+	//dd:endinstrument
+}
+`
+
+	reader, err := InstrumentFile("test", strings.NewReader(tpl), config.Config{})
+	require.Nil(t, err)
+	got, err := io.ReadAll(reader)
+	require.Nil(t, err)
+	require.Equal(t, tpl, string(got))
+}

internal/instrument/echov4.go

@@ -7,6 +7,17 @@ package instrument
 
 import "github.com/dave/dst"
 
+func instrumentEchoV4(stmt *dst.AssignStmt) []dst.Stmt {
+	if !isEchoV4(stmt) {
+		return nil
+	}
+	stmt.Decorations().Start.Prepend(dd_instrumented)
+	return []dst.Stmt{
+		stmt,
+		echoV4Middleware(stmt),
+	}
+}
+
 func isEchoV4(stmt *dst.AssignStmt) bool {
 	rhs := stmt.Rhs[0]
 	f, ok := funcIdent(rhs)
@@ -19,7 +30,7 @@ func echoV4Middleware(got *dst.AssignStmt) dst.Stmt {
 		return nil
 	}
 	stmt := useMiddleware(iden.Name, "EchoV4Middleware")
-	wrap(stmt)
+	markAsInstrumented(stmt)
 	return stmt
 }
 

internal/instrument/echov4_test.go

@@ -0,0 +1,96 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2023-present Datadog, Inc.
+
+package instrument
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/datadog/orchestrion/internal/config"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestEchoV4(t *testing.T) {
+	var codeTpl = `package main
+
+import %s
+
+func register() {
+	%s
+}
+`
+	var wantTpl = `package main
+
+import (
+	"github.com/datadog/orchestrion/instrument"
+	%s
+)
+
+func register() {
+	//dd:instrumented
+	%s
+	//dd:startinstrument
+	%s
+	//dd:endinstrument
+}
+`
+
+	tests := []struct {
+		pkg  string
+		stmt string
+		want string
+		tmpl string
+	}{
+		{pkg: `"github.com/labstack/echo/v4"`, stmt: `r := echo.New()`, want: `r.Use(instrument.EchoV4Middleware())`, tmpl: wantTpl},
+		{pkg: `echo "github.com/labstack/echo/v4"`, stmt: `r := echo.New()`, want: `r.Use(instrument.EchoV4Middleware())`, tmpl: wantTpl},
+		{pkg: `echov4 "github.com/labstack/echo/v4"`, stmt: `r := echov4.New()`, want: `r.Use(instrument.EchoV4Middleware())`, tmpl: wantTpl},
+	}
+
+	for i, tc := range tests {
+		t.Run(fmt.Sprintf("tc-%d", i), func(t *testing.T) {
+			code := fmt.Sprintf(codeTpl, tc.pkg, tc.stmt)
+			reader, err := InstrumentFile("test", strings.NewReader(code), config.Config{})
+			require.Nil(t, err)
+			got, err := io.ReadAll(reader)
+			require.Nil(t, err)
+			want := fmt.Sprintf(tc.tmpl, tc.pkg, tc.stmt, tc.want)
+			require.Equal(t, want, string(got))
+
+			reader, err = UninstrumentFile("test", strings.NewReader(want), config.Config{})
+			require.Nil(t, err)
+			orig, err := io.ReadAll(reader)
+			require.Nil(t, err)
+			require.Equal(t, code, string(orig))
+		})
+	}
+}
+
+func TestEchoV4Duplicates(t *testing.T) {
+	var tpl = `package main
+
+import (
+	"github.com/datadog/orchestrion/instrument"
+	"github.com/labstack/echo/v4"
+)
+
+func echoV4Server() {
+	//dd:instrumented
+	r := echo.New()
+	//dd:startinstrument
+	r.Use(instrument.EchoV4Middleware())
+	//dd:endinstrument
+}
+`
+
+	reader, err := InstrumentFile("test", strings.NewReader(tpl), config.Config{})
+	require.Nil(t, err)
+	got, err := io.ReadAll(reader)
+	require.Nil(t, err)
+	require.Equal(t, tpl, string(got))
+}

internal/instrument/gin.go

@@ -7,6 +7,17 @@ package instrument
 
 import "github.com/dave/dst"
 
+func instrumentGin(stmt *dst.AssignStmt) []dst.Stmt {
+	if !isGin(stmt) {
+		return nil
+	}
+	stmt.Decorations().Start.Prepend(dd_instrumented)
+	return []dst.Stmt{
+		stmt,
+		ginMiddleware(stmt),
+	}
+}
+
 func isGin(stmt *dst.AssignStmt) bool {
 	rhs := stmt.Rhs[0]
 	f, ok := funcIdent(rhs)
@@ -19,7 +30,7 @@ func ginMiddleware(got *dst.AssignStmt) dst.Stmt {
 		return nil
 	}
 	stmt := useMiddleware(iden.Name, "GinMiddleware")
-	wrap(stmt)
+	markAsInstrumented(stmt)
 	return stmt
 }
 

internal/instrument/gin_test.go

@@ -0,0 +1,102 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2023-present Datadog, Inc.
+
+package instrument
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/datadog/orchestrion/internal/config"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGin(t *testing.T) {
+	var codeTpl = `package main
+
+import "github.com/gin-gonic/gin"
+
+func register() {
+	%s
+}
+`
+	var wantTpl = `package main
+
+import (
+	"github.com/datadog/orchestrion/instrument"
+	"github.com/gin-gonic/gin"
+)
+
+func register() {
+	//dd:instrumented
+	%s
+	//dd:startinstrument
+	%s
+	//dd:endinstrument
+}
+`
+
+	tests := []struct {
+		in   string
+		want string
+		tmpl string
+	}{
+		{in: `g := gin.New()`, want: `g.Use(instrument.GinMiddleware())`, tmpl: wantTpl},
+		{in: `g := gin.Default()`, want: `g.Use(instrument.GinMiddleware())`, tmpl: wantTpl},
+	}
+
+	for i, tc := range tests {
+		t.Run(fmt.Sprintf("tc-%d", i), func(t *testing.T) {
+			code := fmt.Sprintf(codeTpl, tc.in)
+			reader, err := InstrumentFile("test", strings.NewReader(code), config.Config{})
+			require.Nil(t, err)
+			got, err := io.ReadAll(reader)
+			require.Nil(t, err)
+			want := fmt.Sprintf(tc.tmpl, tc.in, tc.want)
+			require.Equal(t, want, string(got))
+
+			reader, err = UninstrumentFile("test", strings.NewReader(want), config.Config{})
+			require.Nil(t, err)
+			orig, err := io.ReadAll(reader)
+			require.Nil(t, err)
+			require.Equal(t, code, string(orig))
+		})
+	}
+}
+
+func TestGinDuplicates(t *testing.T) {
+	var tpl = `package main
+
+import (
+	"net/http"
+
+	"github.com/datadog/orchestrion/instrument"
+	"github.com/gin-gonic/gin"
+)
+
+func ginServer() {
+	//dd:instrumented
+	r := gin.Default()
+	//dd:startinstrument
+	r.Use(instrument.GinMiddleware())
+	//dd:endinstrument
+	r.GET("/ping", func(c *gin.Context) {
+		c.JSON(http.StatusOK, gin.H{
+			"message": "pong",
+		})
+	})
+	r.Run()
+}
+`
+
+	reader, err := InstrumentFile("test", strings.NewReader(tpl), config.Config{})
+	require.Nil(t, err)
+	got, err := io.ReadAll(reader)
+	require.Nil(t, err)
+	require.Equal(t, tpl, string(got))
+}

internal/instrument/gorilla.go

@@ -16,7 +16,7 @@ func wrapGorillaMux(stmt *dst.AssignStmt) {
 	if !(f.Path == "github.com/gorilla/mux" && f.Name == "NewRouter") {
 		return
 	}
-	wrap(stmt)
+	markAsWrap(stmt)
 	call := rhs.(*dst.CallExpr)
 	call.Args = []dst.Expr{
 		&dst.CallExpr{

internal/instrument/grpc.go

@@ -28,7 +28,7 @@ func wrapGRPC(stmt *dst.AssignStmt) {
 				if !(iden.Name == targetName && iden.Path == "google.golang.org/grpc") {
 					return
 				}
-				wrap(stmt)
+				markAsWrap(stmt)
 				for _, opt := range opts {
 					fun.Args = append(fun.Args,
 						&dst.CallExpr{Fun: &dst.Ident{Name: opt, Path: "github.com/datadog/orchestrion/instrument"}},