Add uprobes to execDC in order to instrument SQL DML

CHANGELOG.md

@@ -18,6 +18,7 @@ OpenTelemetry Go Automatic Instrumentation adheres to [Semantic Versioning](http
 - Add HTTP status code attribute to `net/http` server instrumentation. ([#428](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/428))
 - The instrumentation scope now includes the version of the auto-instrumentation project. ([#442](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/442))
 - Add a new `WithSampler` method allowing end-users to provide their own implementation of OpenTelemetry sampler directly through the package API. ([#468](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/468)).
+- Add uprobes to `execDC` in order to instrument SQL DML ([#475](https://github.com/open-telemetry/opentelemetry-go-instrumentation/pull/475)).
 
 ### Changed
 

examples/httpPlusdb/main.go

@@ -29,15 +29,15 @@ const (
 	dbName   = "test.db"
 
 	tableDefinition = `CREATE TABLE contacts (
-							contact_id INTEGER PRIMARY KEY,
-							first_name TEXT NOT NULL,
-							last_name TEXT NOT NULL,
-							email TEXT NOT NULL UNIQUE,
-							phone TEXT NOT NULL UNIQUE);`
+	contact_id INTEGER PRIMARY KEY,
+	first_name TEXT NOT NULL,
+	last_name TEXT NOT NULL,
+	email TEXT NOT NULL,
+	phone TEXT NOT NULL);`
 
 	tableInsertion = `INSERT INTO 'contacts'
-						('first_name', 'last_name', 'email', 'phone') VALUES
-						('Moshe', 'Levi', '[email protected]', '052-1234567');`
+	('first_name', 'last_name', 'email', 'phone') VALUES
+	('Moshe', 'Levi', '[email protected]', '052-1234567');`
 )
 
 // Server is Http server that exposes multiple endpoints.
@@ -73,12 +73,6 @@ func NewServer() *Server {
 		panic(err)
 	}
 
-	_, err = database.Exec(tableInsertion)
-
-	if err != nil {
-		panic(err)
-	}
-
 	return &Server{
 		db: database,
 	}
@@ -92,6 +86,11 @@ func (s *Server) queryDb(w http.ResponseWriter, req *http.Request) {
 		panic(err)
 	}
 
+	_, err = s.db.ExecContext(ctx, tableInsertion)
+	if err != nil {
+		panic(err)
+	}
+
 	rows, err := conn.QueryContext(req.Context(), sqlQuery)
 	if err != nil {
 		panic(err)

internal/pkg/instrumentation/bpf/database/sql/bpf/probe.bpf.c

@@ -20,7 +20,7 @@
 
 char __license[] SEC("license") = "Dual MIT/GPL";
 
-#define MAX_QUERY_SIZE 100
+#define MAX_QUERY_SIZE 256
 #define MAX_CONCURRENT 50
 
 struct sql_request_t {
@@ -84,4 +84,48 @@ int uprobe_queryDC(struct pt_regs *ctx) {
 
 // This instrumentation attaches uprobe to the following function:
 // func (db *DB) queryDC(ctx, txctx context.Context, dc *driverConn, releaseConn func(error), query string, args []any)
-UPROBE_RETURN(queryDC, struct sql_request_t, sql_events, events, 3, 0, true)
+UPROBE_RETURN(queryDC, struct sql_request_t, sql_events, events, 3, 0, true)
+
+// This instrumentation attaches uprobe to the following function:
+// func (db *DB) execDC(ctx context.Context, dc *driverConn, release func(error), query string, args []any)
+SEC("uprobe/execDC")
+int uprobe_execDC(struct pt_regs *ctx) {
+    // argument positions
+    u64 context_ptr_pos = 3;
+    u64 query_str_ptr_pos = 6;
+    u64 query_str_len_pos = 7;
+
+    struct sql_request_t sql_request = {0};
+    sql_request.start_time = bpf_ktime_get_ns();
+
+    if (should_include_db_statement) {
+        // Read Query string
+        void *query_str_ptr = get_argument(ctx, query_str_ptr_pos);
+        u64 query_str_len = (u64)get_argument(ctx, query_str_len_pos);
+        u64 query_size = MAX_QUERY_SIZE < query_str_len ? MAX_QUERY_SIZE : query_str_len;
+        bpf_probe_read(sql_request.query, query_size, query_str_ptr);
+    }
+
+    // Get parent if exists
+    void *context_ptr_val = get_Go_context(ctx, 3, 0, true);
+    struct span_context *span_ctx = get_parent_span_context(context_ptr_val);
+    if (span_ctx != NULL) {
+        // Set the parent context
+        bpf_probe_read(&sql_request.psc, sizeof(sql_request.psc), span_ctx);
+        copy_byte_arrays(sql_request.psc.TraceID, sql_request.sc.TraceID, TRACE_ID_SIZE);
+        generate_random_bytes(sql_request.sc.SpanID, SPAN_ID_SIZE);
+    } else {
+        sql_request.sc = generate_span_context();
+    }
+
+    // Get key
+    void *key = get_consistent_key(ctx, context_ptr_val);
+
+    bpf_map_update_elem(&sql_events, &key, &sql_request, 0);
+    start_tracking_span(context_ptr_val, &sql_request.sc);
+    return 0;
+}
+
+// This instrumentation attaches uprobe to the following function:
+// func (db *DB) execDC(ctx context.Context, dc *driverConn, release func(error), query string, args []any)
+UPROBE_RETURN(execDC, struct sql_request_t, sql_events, events, 3, 0, true)

internal/pkg/instrumentation/bpf/database/sql/probe.go

@@ -48,7 +48,7 @@ const instrumentedPkg = "database/sql"
 // request-response.
 type Event struct {
 	context.BaseSpanProperties
-	Query [100]byte
+	Query [256]byte
 }
 
 // Probe is the database/sql instrumentation probe.
@@ -75,7 +75,10 @@ func (h *Probe) LibraryName() string {
 
 // FuncNames returns the function names from "database/sql" that are instrumented.
 func (h *Probe) FuncNames() []string {
-	return []string{"database/sql.(*DB).queryDC"}
+	return []string{
+		"database/sql.(*DB).queryDC",
+		"database/sql.(*DB).execDC",
+	}
 }
 
 // Load loads all instrumentation offsets.
@@ -139,6 +142,35 @@ func (h *Probe) Load(exec *link.Executable, target *process.TargetDetails) error
 		h.returnProbs = append(h.returnProbs, retProbe)
 	}
 
+	offset, err = target.GetFunctionOffset(h.FuncNames()[1])
+	if err != nil {
+		return err
+	}
+
+	up, err = exec.Uprobe("", h.bpfObjects.UprobeExecDC, &link.UprobeOptions{
+		Address: offset,
+	})
+	if err != nil {
+		return err
+	}
+
+	h.uprobes = append(h.uprobes, up)
+
+	retOffsets, err = target.GetFunctionReturns(h.FuncNames()[1])
+	if err != nil {
+		return err
+	}
+
+	for _, ret := range retOffsets {
+		retProbe, err := exec.Uprobe("", h.bpfObjects.UprobeExecDC_Returns, &link.UprobeOptions{
+			Address: ret,
+		})
+		if err != nil {
+			return err
+		}
+		h.returnProbs = append(h.returnProbs, retProbe)
+	}
+
 	rd, err := perf.NewReader(h.bpfObjects.Events, os.Getpagesize())
 	if err != nil {
 		return err

internal/pkg/instrumentation/bpf/database/sql/probe_test.go

@@ -44,7 +44,7 @@ func TestProbeConvertEvent(t *testing.T) {
 			SpanContext: context.EBPFSpanContext{TraceID: traceID, SpanID: spanID},
 		},
 		// "SELECT * FROM foo"
-		Query: [100]byte{0x53, 0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x2a, 0x20, 0x46, 0x52, 0x4f, 0x4d, 0x20, 0x66, 0x6f, 0x6f},
+		Query: [256]byte{0x53, 0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x2a, 0x20, 0x46, 0x52, 0x4f, 0x4d, 0x20, 0x66, 0x6f, 0x6f},
 	})
 
 	sc := trace.NewSpanContext(trace.SpanContextConfig{