From 9dda1ccf4db29b12fceabc2d5224a2c7026d6ca4 Mon Sep 17 00:00:00 2001
From: Christian Simon <simon@swine.de>
Date: Wed, 17 Jan 2024 15:07:08 +0000
Subject: [PATCH] Add ProcessInfo2 call

---
 README.md     |  1 +
 client.go     | 59 +++++++++++++++++++++++++++++++++++++++++++++++++--
 dotnetdiag.go | 39 +++++++++++++++++++++++++++++++---
 3 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index d8492e4..b0c177a 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,7 @@ Supported platforms:
 Implemented commands:
  - [x] StopTracing
  - [x] CollectTracing
+ - [x] ProcessInfo2
  - [ ] CollectTracing2
  - [ ] CreateCoreDump
  - [ ] AttachProfiler
diff --git a/client.go b/client.go
index f838b34..eb7dd95 100644
--- a/client.go
+++ b/client.go
@@ -1,8 +1,13 @@
 package dotnetdiag
 
 import (
+	"bytes"
+	"encoding/binary"
 	"fmt"
+	"io"
 	"net"
+
+	"github.com/pyroscope-io/dotnetdiag/nettrace"
 )
 
 // Client implement Diagnostic IPC Protocol client.
@@ -58,8 +63,8 @@ type CollectTracingConfig struct {
 // NewClient creates a new Diagnostic IPC Protocol client for the transport
 // specified - on Unix/Linux based platforms, a Unix Domain Socket will be used, and
 // on Windows, a Named Pipe will be used:
-//  - /tmp/dotnet-diagnostic-{%d:PID}-{%llu:disambiguation key}-socket (Linux/MacOS)
-//  - \\.\pipe\dotnet-diagnostic-{%d:PID} (Windows)
+//   - /tmp/dotnet-diagnostic-{%d:PID}-{%llu:disambiguation key}-socket (Linux/MacOS)
+//   - \\.\pipe\dotnet-diagnostic-{%d:PID} (Windows)
 //
 // Refer to documentation for details:
 // https://github.com/dotnet/diagnostics/blob/main/documentation/design-docs/ipc-protocol.md#transport
@@ -74,6 +79,56 @@ func NewClient(addr string, options ...Option) *Client {
 	return c
 }
 
+func (c *Client) ProcessInfo2() (*ProcessInfo2Response, error) {
+	conn, err := c.dial(c.addr)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Close()
+
+	if err = writeMessage(conn, CommandSetProcess, ProcessProcessInfo2, nil); err != nil {
+		return nil, err
+	}
+
+	var resp ProcessInfo2Response
+	header, err := readResponseHeader(conn)
+	if err != nil {
+		return nil, err
+	}
+
+	if header.CommandSet != CommandSetServer || header.CommandID != 00 {
+		return nil, fmt.Errorf("unexpected response header: commandSet=%v (expected 0xff) commandID=%v (expected 0x00)", header.CommandSet, header.CommandID)
+	}
+
+	buf := bytes.NewBuffer(nil)
+	if _, err := io.CopyN(buf, conn, int64(header.Size-headerSize)); err != nil {
+		return nil, err
+	}
+
+	if err := binary.Read(buf, binary.LittleEndian, &resp.ProcessID); err != nil {
+		return nil, fmt.Errorf("unable to read process ID: %w", err)
+	}
+
+	if err := binary.Read(buf, binary.LittleEndian, &resp.GUID); err != nil {
+		return nil, fmt.Errorf("unable to read process ID: %w", err)
+	}
+
+	// now parse the strings out
+	p := &nettrace.Parser{Buffer: buf}
+	p.UTF16NTS()
+	resp.CommandLine = p.UTF16NTS()
+	p.UTF16NTS()
+	resp.OS = p.UTF16NTS()
+	p.UTF16NTS()
+	resp.Arch = p.UTF16NTS()
+	p.UTF16NTS()
+	resp.AssemblyName = p.UTF16NTS()
+	p.UTF16NTS()
+	resp.RuntimeVersion = p.UTF16NTS()
+
+	return &resp, nil
+}
+
 // CollectTracing creates a new EventPipe session stream of NetTrace data.
 func (c *Client) CollectTracing(config CollectTracingConfig) (s *Session, err error) {
 	// Every session has its own IPC connection which cannot be reused for any
diff --git a/dotnetdiag.go b/dotnetdiag.go
index 34b2f08..8148947 100644
--- a/dotnetdiag.go
+++ b/dotnetdiag.go
@@ -39,6 +39,18 @@ const (
 	CommandSetServer = 0xFF
 )
 
+const (
+	ProcessProcessInfo uint8 = iota
+	ProcessResumeRuntime
+	ProcessProcessEnvironment
+	_ // 0x03 not used
+	ProcessProcessInfo2
+	ProcessEnablePerfMap
+	ProcessDisablePerfMap
+	ProcessApplyStartupHook
+	ProcessProcessInfo3
+)
+
 const (
 	_ = iota
 	EventPipeStopTracing
@@ -82,6 +94,16 @@ type StopTracingResponse struct {
 	SessionID uint64
 }
 
+type ProcessInfo2Response struct {
+	ProcessID      uint64
+	GUID           [16]byte
+	CommandLine    string
+	OS             string
+	Arch           string
+	AssemblyName   string
+	RuntimeVersion string
+}
+
 func writeMessage(w io.Writer, commandSet, commandID uint8, payload []byte) error {
 	bw := bufio.NewWriter(w)
 	err := binary.Write(bw, binary.LittleEndian, Header{
@@ -100,14 +122,25 @@ func writeMessage(w io.Writer, commandSet, commandID uint8, payload []byte) erro
 	return bw.Flush()
 }
 
-func readResponse(r io.Reader, v interface{}) error {
+func readResponseHeader(r io.Reader) (*Header, error) {
 	var h Header
 	if err := binary.Read(r, binary.LittleEndian, &h); err != nil {
-		return err
+		return nil, err
 	}
+
 	if h.Magic != magic {
-		return ErrHeaderMalformed
+		return nil, ErrHeaderMalformed
 	}
+
+	return &h, nil
+}
+
+func readResponse(r io.Reader, v interface{}) error {
+	h, err := readResponseHeader(r)
+	if err != nil {
+		return err
+	}
+
 	if !(h.CommandSet == CommandSetServer && h.CommandID == 0xFF) {
 		return binary.Read(r, binary.LittleEndian, v)
 	}